import matplotlib.pyplot as plt
import numpy as np
from itertools import cycle
import argparse
import pickle


def is_ipython():
    ''' Check if script is run in IPython.

    Returns:
        bool: True if IPython, else False '''
    try:
        get_ipython()
        ipy = True
    except NameError:
        ipy = False

    return ipy


def load_data(file):
    ''' Load numpy data from file.

    Returns
        dict: data dictionary
    '''
    dat = np.load(file)

    return dat


def plot_parameters(dat, deparameterize=False, ref=None):
    ''' Plot the parameters in separate subplots with uncertainties.

    Args:
        dat (dict):  data dictionary
        deparameterize (bool):  flag indicating if parameters should be
          deparameterized via 2**theta
        ref:    reference value to be plotted with parameters
    '''
    if is_ipython():
        plt.ion()

    dim = dat['theta'].shape[-1]
    fig1, axes = plt.subplots(1, dim)

    if dim == 1:
        axes = [axes]

    axes[0].set_ylabel(r'$\theta$')

    t = dat['times']
    theta = dat['theta']
    P = dat['P_theta']

    col = cycle(['C0', 'C1', 'C0', 'C1'])
    ls = cycle(['-', '-', '--', '--', ':', ':', '-.', '-.'])

    col_ = next(col)
    ls_ = next(ls)
    for i, ax in enumerate(axes):
        if dim == 1:
            theta = theta.reshape((-1, 1))
            P = P.reshape((-1, 1, 1))

        if deparameterize:
            ax.plot(t, 2**theta[:, i], '-', c=col_, ls=ls_)
        else:
            ax.plot(t, theta[:, i], '-', c=col_, ls=ls_)
            ax.fill_between(t, theta[:, i] - np.sqrt(P[:, i, i]),
                            theta[:, i] + np.sqrt(P[:, i, i]), alpha=0.3,
                            color=col_)
        ax.set_xlabel(r'time')

    if ref:
        if isinstance(ref, (int, float)):
            ref = np.array(ref)

        for ax, ri in zip(axes, ref):
            if ri:
                # if deparameterize:
                ax.plot((t[0], t[-1]), (ri, )*2, '-.k', lw=2,
                        label=r'ground truth')
                # else:
                #     ax.plot((t[0], t[-1]), (np.log2(ri), )*2, '-.k', lw=2,
                #             label=r'ground truth')

    # print('theta_peak: \t {}'.format(theta[round(len(theta)/2), :]))
    print('Final value theta: \t {}'.format(theta[-1, :]))
    print('Deparameterized: 2^theta_end: \t {}'.format(2**theta[-1, :]))

    if not is_ipython():
        plt.show()


def get_parser():
    parser = argparse.ArgumentParser(
        description='''
Plot the time evolution of the ROUKF estimated parameters.

To execute in IPython::

    %run plot_roukf_parameters.py [-d] [-r N [N \
...]] file
''',
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('file', type=str, help='path to ROUKF stats file')
    parser.add_argument('-d', '--deparameterize', action='store_true',
                        help='deparameterize the parameters by 2**theta')
    parser.add_argument('-r', '--ref', metavar='N', nargs='+', default=None,
                        type=float, help='Reference values for parameters')
    return parser


if __name__ == '__main__':
    args = get_parser().parse_args()

    dat = load_data(args.file)

    plot_parameters(dat, deparameterize=args.deparameterize, ref=args.ref)