import numpy as np
#import wakeoscillator as wo

import os
import re
import pickle as pkl
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

import SSDpy as ssd


def print_version():
    from . import __version__
    print(f"VIVyd version: {__version__}")


class VIVModel:
    
    def __init__(self, nu=np.linspace(0.1, 2.0, 20), xi_s=0.003, param=None, overwrite_param=None, type=None):
        """Constructor for the VIVModel class. This class is used to define a VIV model and its parameters.
        It can be specified by the user or by using a predefined model (e.g. Hartlen-Currie or Tamura-Feng).
        Parametric studies can be performed by varying the parameters nu and xi_s, which should be provided
        as lists or numpy arrays.
        After defining the model, the user can run simulations using the VIVSimulation class.
        
        Important attributes:
        - type: the type of the model (e.g. 'User-defined', 'Hartlen-Currie', 'Tamura-Feng')
        - nu: a list of reduced velocities
        - xi_s: a list of damping ratios
        - param: a dictionary with the parameters of the generic model
        - own_param: a dictionary with the parameters of the original model, if specified (e.g. 'Hartlen-Currie', 'Tamura-Feng')
        """

        nu, xi_s = self.check_input(nu, xi_s)
        
        # create a consistent set of parameters for the model
        if type is None:
            type = 'User-defined'
            
            # default parameters
            param = {'epsilon': 0.1, 'a': 0.1, 'b': 0,
                        'A0': 1, 'A1': 0., 'A2': 0.,
                        'nu1': 1., 'sig1': 0., 'nu2': 1., 'sig2': 0.}
                   
            # check if some parameters are overwritten
            if overwrite_param is not None:
                for key in overwrite_param:
                    if key in param:
                        print(f"Overwriting parameter {key} with value {overwrite_param[key]}")
                        param[key] = overwrite_param[key]
                    else:
                        print(f"Parameter {key} not found in model's parameters")
            
        else:
            if type == 'Hartlen-Currie':
                # from their 1970 paper
                rho = 1.24  # converted from slug/ft^3 to kg/m^3
                D = 3 * 0.0254
                St = 0.2
                M = 1.583 * 0.4536
                L = 27 * 0.0254
                a = rho * D**2 * L / ( 8 * np.pi ** 2 * St ** 2 * M )
                CL0 = 0.2
                alpha = 0.02
                gamma = 4/3 * alpha / CL0**2
                own_param = {'St': St, 'CL0': CL0, 'zeta': 0.0015, 'alpha': alpha, 'a': a, 'b': 0.4, 'gamma': gamma}
                xi_s = [own_param['zeta']]
            
            elif type == 'Tamura-Feng':
                # from his 1981 paper
                own_param = {'St': 0.2, 'CL0': 0.4, 'f': 1.16, 'm*': 0.625, 'CD': 1.2, 'zeta': 0.0378}
                own_param['n'] =  0.00257
                own_param['eta'] = 0.00181
                xi_s = [own_param['eta']]
            elif type == 'Tamura 1983':
                # from his 1983 paper
                own_param = {'St': 0.2, 'CL0': 0.4, 'f': 1.16, 'm*': 0.625, 'CD': 1.2, 'zeta': 0.0378}
                own_param['n'] =  0.0034
                own_param['eta'] = 0.00285
                xi_s = [own_param['eta']]
            elif type == 'Rigo-Andrianne-Denoel':
                own_param = {'alpha': -0.09, 'gamma': 0.009, 'delta': 0.0063, 'L_eta': 1.3, 'sig_eta': 0.5}
                own_param['xi_s'] = 0.2 # arbitrary large value to simulate a fixed cylinder
            
            self.own_param = own_param
            
            # check if some parameters are overwritten
            if overwrite_param is not None:
                for key in overwrite_param:
                    if key in own_param:
                        print(f"Overwriting parameter {key} with value {overwrite_param[key]}")
                        own_param[key] = overwrite_param[key]

                        if type == 'Hartlen-Currie' and key == 'zeta':
                                xi_s = [overwrite_param[key]]
                        if (type == 'Tamura-Feng' or type == 'Tamura 1983') and key == 'eta':
                                xi_s = [overwrite_param[key]]
                    else:
                        print(f"Parameter {key} not found in model's parameters")
            
        self.type = type
        self.nu = nu
        self.xi_s = xi_s
        self.param = param


    def check_input(self, nu, xi_s):

        # if nu or xi_s are numpy array, convert them to list
        if isinstance(nu, np.ndarray):
            nu = nu.tolist()
        if isinstance(xi_s, np.ndarray):
            xi_s = xi_s.tolist()
        
        # make sure all items in nu and xi_s are iterable (lists)
        if not isinstance(nu, list):
            nu = [nu]
        if not isinstance(xi_s, list):
            xi_s = [xi_s]
        return nu, xi_s
    
    def get_parameter(self, nu, xi_s):
        # use this function to transform the parameters of an original model (e.g Tamura, Facchinetti, ...),
        # to the parameters of the generic model

        if self.type == 'User-defined':
            return self.param
        
        else:
            ownp = self.own_param
            if self.type == 'Hartlen-Currie':
                xi_s = ownp['zeta']
                a = 0
                b = ownp['gamma'] / nu**2
                eps = ownp['alpha']
                A0 = 0
                A1 = ownp['b']
                A2 = 0
                M0 = ownp['a'] * nu**2
                nu1, sig1, nu2, sig2 = 1.0, 0.0, 1.0, 0.0

            elif self.type == 'Tamura-Feng' or self.type == 'Tamura 1983':
                # notice that xi_s is modified in Tamura's model with the drag effect
                xi_s = xi_s + nu * ownp['n'] * (ownp['f'] + ownp['CD']) / ( 4* np.pi * ownp['St'])
                a = 8. * ownp['zeta'] * ownp['f']**2 / ownp['CL0']**2
                b = 0.
                eps = 2 * ownp['zeta']
                A0 = 0
                A1 = -2 * np.pi * ownp['St'] * nu
                A2 = - ownp['m*']
                M0 = - ownp['f'] * ownp['n'] * nu**2 / (2 * np.pi * ownp['St'])**2
                nu1, sig1, nu2, sig2 = 1.0, 0.0, 1.0, 0.0
            
            elif self.type == 'Rigo-Andrianne-Denoel':
                xi_s = ownp['xi_s']
                a = 0
                b = 0
                eps = 0
                A0, A1, A2, M0 = 0., 0., 0., 1.e-6
                nu1, sig1 = 0., 0.
                nu2, sig2 = ownp['L_eta'], ownp['sig_eta']
    
        return {'nu': nu, 'xi_s': xi_s, 'a': a, 'b': b, 'epsilon': eps,
                'A0': A0, 'A1': A1, 'A2': A2, 'M0': M0,
                'nu1': nu1, 'sig1': sig1, 'nu2': nu2, 'sig2': sig2}

    
    def print(self):
        print(self.type)
        print(self.param)


class VIVSimulation:
    def __init__(self, model, verbosity=1,filemask='VIV_results'):
        self.model = model
        self.verbosity = verbosity
        self.sims = []

        # create a list of simulations to be performed
        for nu in model.nu:
            for xi_s in model.xi_s:
                self.sims.append({'param': model.get_parameter(nu, xi_s)})
                        
        # find next available file number in the results directory
        self.dir = 'results' + os.sep + model.type
        self.filemask = filemask
        self.filename = self.get_data_filename()

        # run the simulations
        i = 1

        for sim in self.sims:
            print (f"Running simulation {i} of {len(self.sims)}")
            
            t, y, skip = self.simulation(sim['param'])

            # compute the envelopes of the response
            _, env_y = ssd.signal.envelopes(y[:,0])
            _, env_q = ssd.signal.envelopes(y[:,2])
            
            # compute basic statistics of cylinder displacement and lift force
            stats_y = self.compute_stats(y[:, 0],  skip)
            stats_q = self.compute_stats(y[:, 2],  skip)
            stats_env_y = self.compute_stats(env_y, skip)
            stats_env_q = self.compute_stats(env_q, skip)
            
            # analyse phase and frequency of the response
            f1, f2, phase_shift, spectra = self.analyze_phase_and_freq(t[skip:], y[skip:, 0], y[skip:, 2])

            # store this result in the sims dictionary
            self.sims[i-1]['time'] = t
            self.sims[i-1]['time series'] = {'y': y[:, 0], 'q': y[:, 2]}
            self.sims[i-1]['envelopes'] = {'y': env_y, 'q': env_q}
            self.sims[i-1]['statistics'] = {'y': stats_y, 'q': stats_q, 'env_y': stats_env_y, 'env_q': stats_env_q}
            self.sims[i-1]['phase'] = phase_shift
            self.sims[i-1]['spectra'] = spectra
            self.sims[i-1]['leading_freq'] = {'y': f1, 'q': f2}
            
            # plot details of the simulation if there are not too many to display
            if len(self.sims) <= 10:
                self.plot_results(self.sims[i-1])
            
            if self.verbosity > 0:
                print(f"  nu = {sim['param']['nu']:.2f}")
                print(f"  std_y = {stats_y['std']:.2e}")
                print(f"  mean_env_y = {stats_env_y['mean']:.2e}")
                print(f"  std_env_y = {stats_env_y['std']:.2e}")
                print(f"  mean_env_q = {stats_env_q['mean']:.2e}")
                print(f"  Leading frequency y = {f1:.2f}")

            self.save_data(self.sims)
        
            i += 1

    @staticmethod
    def ode_func(Y, t, param):
        # define the ODEs of the VIV model
        nu = param['nu']
        a, b, eps = param['a'], param['b'], param['epsilon']
        xi_s, M0 = param['xi_s'], param['M0']
        A0, A1, A2 = param['A0'], param['A1'], param['A2']
        nu1, sig1, nu2, sig2 = param['nu1'], param['sig1'], param['nu2'], param['sig2']
        
        y, dy, q, dq, eta1, eta2 = Y
        
        wn1 = np.interp(t, param['wnt'], param['wn1'])
        wn2 = np.interp(t, param['wnt'], param['wn2'])

        f1 = dy
        f2 = -2 * xi_s * dy - y + M0 * q + eta1
        f3 = dq
        f4 = -nu**2 * q - nu * (a * q**2 + b * dq**2 - eps) * dq + (A2 * f2 + A1 * dy + A0 * y) + eta2
        f5 = -nu1 * eta1 + sig1 * wn1
        f6 = -nu2 * eta2 + sig2 * wn2

        return [f1, f2, f3, f4, f5, f6]

    def simulation(self, param):
        """ We run simulations by batch sizes, which correspond to 'batchSize'' cycles. After the batch has run
        we evaluate the statistics of the response, and stop simulating if they are converged."""

        sim_param = {'fs': 50, 'batchSize': 100 ,'maxBatches': 100, 'storeBatches': 4, 'convTol': 1e-4}
        
        dt = 2 * np.pi / sim_param['fs']
        
        sol_t = np.empty(0)
        sol_y = np.empty((0, 6))
        batch_stats = []   # store the statistics of each batch, to check establishement of steady state
        
        Y0 = [1e-1, 1e-3, 1e-1, 1e-3, 0, 0]
        t0 = 0
        
        ibatch = 0
        is_transient = True
        while True:

            t = np.arange(0, 2 * np.pi * sim_param['batchSize'], dt) + t0

            # generate two white noises, which will be used to generate the noise in each equation
            param['wn1'] = np.random.normal(0, 1, len(t))
            param['wn2'] = np.random.normal(0, 1, len(t))
            param['wnt'] = t

            y = odeint(self.ode_func, Y0, t, args=(param,))

            # concatenate the results
            sol_t = np.concatenate((sol_t, t))
            sol_y = np.concatenate((sol_y, y))

            # compute statistics (mean and std) of batch
            batch_stats.append( np.array( [np.mean(y[:, 0]), np.std(y[:,0]) ] ) )

            if is_transient:
                if ibatch > 10:
                    # check if the transient has passed
                    convCheck = np.abs(batch_stats[-1] - batch_stats[-2])
                    print(f"Batch {ibatch}: convCheck mean: {convCheck[0]:.2e}, convCheck std: {convCheck[1]:.2e}")
                    if np.max(convCheck) < sim_param['convTol']:
                        is_transient = False
                        skip_transient = (ibatch+1) * len(t)
                        ikeep = 1
                else:
                    print (f"Batch {ibatch}: transient - {batch_stats[-1]}")
            else:
                print (f"Batch {ibatch}: in steady state - {batch_stats[-1]}")
                ikeep += 1
                if ikeep > sim_param['storeBatches']:
                    break

            if ibatch > sim_param['maxBatches']:
                print("Maximum number of batches reached. Diffcult to converge, statistically.")
                skip_transient = (ibatch-sim_param['storeBatches']) * len(t)
                break

            ibatch += 1
            t0 = t[-1] + dt
            Y0 = y[-1, :]
            
            
        return sol_t, sol_y, skip_transient

    def compute_stats(self, x, skip):
        # raise Exception if skip is too large
        if skip > len(x):
            raise ValueError("skip is larger than the length of the array")
        
        mean = np.mean(x[skip:])
        std = np.std(x[skip:])
        
        hist, bins = np.histogram(x[skip:], bins=int(4.*np.log(len(x))), density=True)
            
        return {'mean': mean, 'std': std, 'hist': hist, 'bins': bins}

    def analyze_phase_and_freq(self, t, x1, x2):
        # find leading frequency and phase of the response
        psd1, freq1, *_ = ssd.signal.vpsd(x1, len(x1), t[1]-t[0])
        psd2, freq2, *_ = ssd.signal.vpsd(x2, len(x2), t[1]-t[0])
        spectra = {'y': {'psd': psd1, 'freq': freq1}, 'q': {'psd': psd2, 'freq': freq2}}

        # find the peak of the PSD
        i1 = np.argmax(psd1)
        i2 = np.argmax(psd2)
        f_leading1 = freq1[i1]
        f_leading2 = freq2[i2]

        # compute the phase shift between the two signals
        # concatenate X1 and X2 to from a 2-by-n array

        phase_shift = ssd.signal.phase(np.vstack((x1, x2)))

        return f_leading1, f_leading2, phase_shift, spectra


    def plot_results(self, sim):
        t = sim['time']
        y = sim['time series']['y']
        env_y = sim['envelopes']['y']
        stats = sim['statistics']
        
        ssd.tools.plot_tf(t, y)
        plt.subplot(2, 1, 1)
        plt.plot(t, env_y, label='envelope')
        plt.show()

        plt.figure()
        plt.subplot(1, 2, 1)
        plt.bar(stats['env_y']['bins'][:-1], stats['env_y']['hist'], width=np.diff(stats['y']['bins']))
        plt.title('Envelope Histogram')
        
        plt.subplot(1, 2, 2)
        plt.bar(stats['y']['bins'][:-1], stats['y']['hist'], width=np.diff(stats['y']['bins']))
        plt.title('y Histogram')
        plt.show()

    def get_data_filename(self):
         # find next available file number
        if not os.path.exists(self.dir):
            os.makedirs(self.dir)

        files = os.listdir(self.dir)
        numbers = [int(file[-7:-4]) for file in files if re.match(self.filemask + '\d\d\d.pkl', file)]
        next_num = max(numbers) + 1 if numbers else 1
        new_filename = os.path.join(self.dir, f"{self.filemask}{next_num:03d}.pkl")

        return new_filename
    
    def save_data(self, sims, mask='VIV_results'):

        # if self.filename exists, delete the file and replace by new file (not easy to append to a pickle file)
        new_filename = self.filename
        if os.path.exists(new_filename):
            os.remove(new_filename)

        with open(new_filename, 'wb') as f:
            pkl.dump({'model': self.model, 'sims': sims}, f)
        print(f'File saved: {new_filename}')

    @staticmethod
    def load_data(num, mask='VIV_results', dir='results'):
        filename = os.path.join(dir, f"{mask}{num:03d}.pkl")
        with open(filename, 'rb') as f:
            return pkl.load(f)

    @staticmethod
    def list_files(mask='VIV_results', dir='results'):
        return [file for file in os.listdir(dir) if re.match(mask + r'\d\d\d.pkl', file)]


print('start.')

do_run = True

if do_run:
    
    # Define a model with parameters
    #model = VIVModel(nu = np.linspace(0.8,1.3,30), type = 'Tamura 1983')

    #model = VIVModel(nu = np.linspace(0.8,1.4,25), type = 'Tamura-Feng')
    #model = VIVModel(nu = np.linspace(0.9,1.2,30), type = 'Hartlen-Currie')
    
    #model = VIVModel(nu = np.linspace(0.8,1.3,30), type = 'Rigo-Andrianne-Denoel')
    model = VIVModel(nu = 1, type = 'Rigo-Andrianne-Denoel')

    
    

    # Run simulations and save data
    simulation = VIVSimulation(model)

else:
    # Load data from a previous run
    simulation = VIVSimulation.load_data(1)
    print(data)

# do some post-processing
if len(simulation.sims) > 1:
    # loop over all simulations and store nu, std_y, mean_env_y, std_env_y, leading_freq in lists
    nu = []
    std_y = []
    mean_env_y = []
    std_env_y = []
    mean_env_q = []

    leading_freq_y = []
    leading_freq_q = []

    mean_phase = []

    for sim in simulation.sims:
        nu.append(sim['param']['nu'])
        std_y.append(sim['statistics']['y']['std'])
        mean_env_y.append(sim['statistics']['env_y']['mean'])
        std_env_y.append(sim['statistics']['env_y']['std'])
        mean_env_q.append(sim['statistics']['env_q']['mean'])
        mean_phase.append(np.mod(np.mean(sim['phase']), 2 * np.pi))

        leading_freq_y.append(2 * np.pi * sim['leading_freq']['y'])
        leading_freq_q.append(2 * np.pi * sim['leading_freq']['q'])
        
    # plot the results
    plt.figure(figsize=(5, 4))
    # increase vertical space between subplots
    plt.subplots_adjust(hspace=0.5)

    plt.subplot(2,2,1)
    plt.plot(nu, mean_env_y, label='env_y')
    plt.xlabel(r'$\nu$')
    plt.grid()
    plt.legend(prop={'size': 8})

    plt.subplot(2,2,2)
    plt.plot(nu, mean_env_q, label='env_q')
    plt.xlabel(r'$\nu$')
    plt.grid()
    plt.legend(prop={'size': 8})

    plt.subplot(2,2,3)
    plt.plot(nu, leading_freq_y, label='Shed. freq.')
    if hasattr(model, 'own_param') and 'St' in model.own_param:
        plt.plot(nu, np.array(nu),'k:', label='Strouhal')
    plt.xlabel(r'$\nu$')
    plt.grid()
    plt.legend(prop={'size': 8})
    
    plt.subplot(2,2,4)
    plt.plot(nu, mean_phase, label='y-q phase shift')
    plt.xlabel(r'$\nu$')
    plt.grid()
    plt.legend(prop={'size': 8})
    plt.gca().set_ylim([0, 2 * np.pi])
    plt.gca().set_yticks([0, np.pi, 2 * np.pi])
    plt.gca().set_yticklabels(['0', '$\pi$', '$2\pi$'])

    ssd.tools.clipfig()

plt.show()
print('done.')


