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):
    ''' 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()


    fig1, axes = plt.subplots(1, 1,figsize=(8,6))
    axes.set_ylabel(r'$\theta$',fontsize=18)

    col = cycle(['C0', 'C1', 'C2', 'C3'])
    

    for k in dat.keys():
        
        t = dat[k]['times']
        theta = dat[k]['theta']
        P = dat[k]['P_theta']

                
        theta = theta.reshape((-1, 1))
        P = P.reshape((-1, 1, 1))
        
        #theta = 2**theta*float(k)
        
        sP = np.sqrt(P[:,0,0])
        sP_up = 2**sP
        sP_down = 2**(-sP)

        col_ = next(col)
        
        axes.plot(t, theta[:, 0], '-', c=col_)
        #axes.fill_between(t, theta[:, 0]*sP_down[:], theta[:, 0]*sP_up[:], alpha=0.3,color=col_)
        axes.fill_between(t, theta[:, 0]  - np.sqrt(P[:, 0, 0]),theta[:, 0]  + np.sqrt(P[:, 0, 0]), alpha=0.3,color=col_)
        axes.set_xlabel(r'time',fontsize=18)
        axes.plot(t, t*0 + np.log(30/float(k))/np.log(2), '-', c='black',ls='--')

        
    axes.set_xlim([-0.01,0.2])
    

    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()
    files = ['10','30','60']
    dat_array = {}

    for ff in files:
        path = args.file + ff + '/theta_stats.npz'
        dat_array[ff] = load_data(path)

    plot_parameters(dat_array)