import os
import re
import pickle as pkl
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from SSDpy.tools.tools import plot_tf, envelopes
import time

class VIVModel:
    
    def __init__(self, Omega=np.linspace(0.1, 2.0, 20), xi_s=0.003, M0=1e-3, param=None, overwrite_param=None, type=None):
                
        if type is None:
            type = 'User-defined'
            param = {'epsilon': 0.1, 'a': 0.1, 'b': 0,
                        'A0': 1, 'A1': 0., 'A2': 0.,
                        'nu1': 1., 'sig1': 0., 'nu2': 1., 'sig2': 0.}
        
        elif type == 'Hartlen-Currie':
            param = {'epsilon': 0.1, 'a': 0.1, 'b': 0,
                        'A0': 1, 'A1': 0., 'A2': 0.,
                        'nu1': 1., 'sig1': 0., 'nu2': 1., 'sig2': 0.}
        
        elif type == 'Tamura-Feng':
                param = {'St': 0.2, 'CL0': 0.4, 'f': 1.16, 'm*': 0.625, 'CD': 1.2, 'zeta': 0.038}
                param['n'] =  0.00257
                param['eta'] = 0.00181
                xi_s = param['eta']
        
        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]
        
        # if Omege or xi_s are numpy array, convert them to list
        if isinstance(Omega, np.ndarray):
            Omega = Omega.tolist()
        if isinstance(xi_s, np.ndarray):
            xi_s = xi_s.tolist()
        # make sure all items in Omega and xi_s are iterable (lists)
        if not isinstance(Omega, list):
            Omega = [Omega]
        if not isinstance(xi_s, list):
            xi_s = [xi_s]
        
        # make sure all items in param are float
        for key in param:
            if isinstance(param[key], int):
                param[key] = float(param[key])
            if not isinstance(param[key], float):
                raise ValueError(f"Parameter {key} should be a float")

        self.type = type
        self.Omega = Omega
        self.xi_s = xi_s
        self.param = param

    def get_parameter(self, Omega, xi_s):
        param = self.param
        if self.type == 'User-defined':
            pass
        
        elif self.type == 'Hartlen-Currie':
            pass
        
        elif self.type == 'Tamura-Feng':
            xi_s = xi_s + Omega * param['n'] * (param['f'] + param['CD']) / ( 4* np.pi * param['St'])
            a = 8. * param['zeta'] * param['f']**2 / param['CL0']**2
            b = 0.
            eps = 2 * param['zeta']
            A0 = 0
            A1 = -2 * np.pi * param['St'] * Omega
            A2 = - param['m*']
            M0 = - param['f'] * param['n'] * Omega**2 / (2 * np.pi * param['St'])**2
            nu1, sig1, nu2, sig2 = 1.0, 0.0, 1.0, 0.0
    
        return {'Omega': Omega, '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, t=np.linspace(0, 1000, 5000), verbosity=1,filemask='VIV_results'):
        self.model = model
        self.skip_transient = int(50 * 2 * np.pi / (t[1] - t[0]))
        self.verbosity = verbosity
        self.sims = []

        # create a list of simulations to be performed
        for Omega in model.Omega:
            for xi_s in model.xi_s:
                self.sims.append({'param': model.get_parameter(Omega, 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
        t0 = time.time()
        for sim in self.sims:
            print (f"Running simulation {i} of {len(self.sims)}")
            
            t, y = self.simulation(sim['param'], t)

            _, env_y = envelopes(y[:,0])
            _, env_q = envelopes(y[:,2])
            
            # compute statistics
            stats_y = self.compute_stats(y[:, 0],  self.skip_transient)
            stats_q = self.compute_stats(y[:, 2],  self.skip_transient)
            stats_env_y = self.compute_stats(env_y, self.skip_transient)
            stats_env_q = self.compute_stats(env_q, self.skip_transient)
            
            # 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}
            
            if len(self.sims) <= 10:
                self.plot_results(self.sims[i-1])
            
            if self.verbosity > 0:
                print(f"  Omega = {sim['param']['Omega']:.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}")                

            remaining_time = (time.time() - t0) / i * (len(self.sims) - i)
            print(f"Simulation {i} done. Estimated remaining time: {remaining_time:.2f} s")
            self.save_data(self.sims)
        
            i += 1

    @staticmethod
    def ode_func(Y, t, param):
        Omega = param['Omega']
        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['t'], param['wn1'])
        wn2 = np.interp(t, param['t'], param['wn2'])

        f1 = dy
        f2 = -2 * xi_s * dy - y + M0 * q + eta1
        f3 = dq
        f4 = -Omega**2 * q - Omega * (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, t=np.linspace(0, 2000, 8000)):
        
        # 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['t'] = t

        Y0 = [1e-6, 1e-6, 1e-6, 1e-6, 0, 0]
        y = odeint(self.ode_func, Y0, t, args=(param,))
        return t, y

    def compute_stats(self, x, skip=1000):
        skip = min(self.skip_transient, int(len(x)/2))
        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 plot_results(self, sim):
        t = sim['time']
        y = sim['time series']['y']
        env_y = sim['envelopes']['y']
        stats = sim['statistics']
        
        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)]
