# Path: pyda/ensemble_generator/ensemble_generator_class.py
# class EnsembleGeneratorClass(object):
#     # This serves as a template for the user of pyda to implement
#     # their own ensemble generation. It usually calls some simulation
#     # routine provided by the user. The initial conditions and
#     # parameters used in the ensemble generation are contained in one
#     # array. Hence the initialization conditions are treated as
#     # parameters.
# 
#     def __init__(self):
#         # Initialization just defines the object and this will be used
#         # by the data assimilation function in the
#         # DataAssimilationClass.
#         self.Name = 'Ensemble Name'
# 
#     # Returns the ensemble array propagated forward from the
#     # initialization and parametrization in Param. 
#     # EnsArray = (Ntimestep*SimulationDimension)x(EnSize) numpy array    
#     def fwd_propagate(self,Param,start_time,stop_time,Ntimestep):
#         # For an example implementation look at
#         # SIRensemble.py.
# 
#         # Initialization includes the parameter array, start time, stop time, and the number of timesteps
#         # Param = (param size + initialization size)x(Ensemble Size) numpy array
#         # start_time, stop_time = float
#         # Ntimestep = integer
#         self.Param = Param
#         self.start_time = start_time
#         self.stop_time = stop_time
#         self.Ntimestep = Ntimestep
#         self.EnSize = Param.shape[1]
#         self.EnsTime = np.linspace(start_time,stop_time,Ntimestep)
# 
#         # Usually a function is to be called a number of times equal
#         # to self.EnSize corresponding to the different
#         # initializations/parametrizations in self.ParamArray. The
#         # simulation is then run from self.start_time to
#         # self.stop_time for a number of timesteps equal to
#         # self.Ntimestep. Regardless of the shape of the array
#         # returned by the simulation each simulation output is turned
#         # into a column vector of succesive time points and inserted
#         # into a column of EnsArray. This is then attached to the
#         # EnsembleGeneratorClass object using self.EnsArray =
#         # EnsArray.
#         raise NotImplementedError(inspect.stack()[0][3])
# 
# Path: pyda/utilities/KSsimulator.py
# def fullfreq(U0, time, Lparam):
# def lowfreq(U0, time, Lparam, FmodeMax):
#     PI = np.pi
#     N = U0.shape[0]
#     L = np.power(k, 2.0) - np.power(k, 4.0)
#     E = np.exp(h*L)
#     E2 = np.exp(h*L/2.0)
#     M = 32
#     LR = h*np.tile(L, (M, 1)).transpose() + np.tile(r, (FmodeMax+1, 1));
#     Q = h*(np.mean((np.exp(LR/2.0) - 1.0)/LR, 1))
#     Q = Q.real
#     PI = np.pi
#     N = U0.shape[0]
#     L = np.power(k, 2.0) - np.power(k, 4.0)
#     E = np.exp(h*L)
#     E2 = np.exp(h*L/2.0)
#     M = 32
#     LR = h*np.tile(L, (M, 1)).transpose() + np.tile(r, (FmodeMax+1, 1));
#     Q = h*(np.mean((np.exp(LR/2.0) - 1.0)/LR, 1))
#     Q = Q.real


import inspect
import numpy as np
from .ensemble_generator_class import EnsembleGeneratorClass
from ..utilities import KSsimulator as KS

# Import the EnsembleGeneratorClass

# Import Epidemic ODE library from pyda utilities

class KSfullfreqEnsemble(EnsembleGeneratorClass):
    def __init__(self):
        # RETURNS:
        #    EnsArray = (Ntimesteps)x(Ngrid)x(EnSize) numpy array. Each ensemble member 
        #               has form 
        #               (U(t_0), U(t_1), ..., U(t_N))^T
        #    EnsTime = Ntimesteps numpy vector of time
        self.Name = 'Deterministic ETD-RK4 full frequency KS'

    def fwd_propagate(self, Param, start_time, stop_time, Ntimestep):
        # Param = (EnSize)x(Ngrid+1) Numpy array of ensemble member initial condition
        #         and bifurcation parameter. Rows look like,
        #         <Lparam>\t<U0>
        #          Lparam = (float) bifurcation parameter
        #          U0 = (Ngrid) numpy vector of KS solution at start_time
        # Ntimestep = number of time steps
        EnSize = Param.shape[0]
        
        # Define solution spatial grid
        Ngrid = Param[0,1:].shape[0]

        # Define time 
        stepsize = (stop_time - start_time)/(Ntimestep)
        EnsTime = np.arange(start_time, stop_time + stepsize, stepsize)

# Path: pyda/ensemble_generator/ensemble_generator_class.py
# class EnsembleGeneratorClass(object):
#     # This serves as a template for the user of pyda to implement
#     # their own ensemble generation. It usually calls some simulation
#     # routine provided by the user. The initial conditions and
#     # parameters used in the ensemble generation are contained in one
#     # array. Hence the initialization conditions are treated as
#     # parameters.
# 
#     def __init__(self):
#         # Initialization just defines the object and this will be used
#         # by the data assimilation function in the
#         # DataAssimilationClass.
#         self.Name = 'Ensemble Name'
# 
#     # Returns the ensemble array propagated forward from the
#     # initialization and parametrization in Param. 
#     # EnsArray = (Ntimestep*SimulationDimension)x(EnSize) numpy array    
#     def fwd_propagate(self,Param,start_time,stop_time,Ntimestep):
#         # For an example implementation look at
#         # SIRensemble.py.
# 
#         # Initialization includes the parameter array, start time, stop time, and the number of timesteps
#         # Param = (param size + initialization size)x(Ensemble Size) numpy array
#         # start_time, stop_time = float
#         # Ntimestep = integer
#         self.Param = Param
#         self.start_time = start_time
#         self.stop_time = stop_time
#         self.Ntimestep = Ntimestep
#         self.EnSize = Param.shape[1]
#         self.EnsTime = np.linspace(start_time,stop_time,Ntimestep)
# 
#         # Usually a function is to be called a number of times equal
#         # to self.EnSize corresponding to the different
#         # initializations/parametrizations in self.ParamArray. The
#         # simulation is then run from self.start_time to
#         # self.stop_time for a number of timesteps equal to
#         # self.Ntimestep. Regardless of the shape of the array
#         # returned by the simulation each simulation output is turned
#         # into a column vector of succesive time points and inserted
#         # into a column of EnsArray. This is then attached to the
#         # EnsembleGeneratorClass object using self.EnsArray =
#         # EnsArray.
#         raise NotImplementedError(inspect.stack()[0][3])
# 
# Path: pyda/utilities/KSsimulator.py
# def fullfreq(U0, time, Lparam):
# def lowfreq(U0, time, Lparam, FmodeMax):
#     PI = np.pi
#     N = U0.shape[0]
#     L = np.power(k, 2.0) - np.power(k, 4.0)
#     E = np.exp(h*L)
#     E2 = np.exp(h*L/2.0)
#     M = 32
#     LR = h*np.tile(L, (M, 1)).transpose() + np.tile(r, (FmodeMax+1, 1));
#     Q = h*(np.mean((np.exp(LR/2.0) - 1.0)/LR, 1))
#     Q = Q.real
#     PI = np.pi
#     N = U0.shape[0]
#     L = np.power(k, 2.0) - np.power(k, 4.0)
#     E = np.exp(h*L)
#     E2 = np.exp(h*L/2.0)
#     M = 32
#     LR = h*np.tile(L, (M, 1)).transpose() + np.tile(r, (FmodeMax+1, 1));
#     Q = h*(np.mean((np.exp(LR/2.0) - 1.0)/LR, 1))
#     Q = Q.real


# Path: pyda/ensemble_generator/KSfullfreqEnsemble.py
import inspect
import numpy as np
from .ensemble_generator_class import EnsembleGeneratorClass
from ..utilities import KSsimulator as KS


# Import the EnsembleGeneratorClass

# Import Epidemic ODE library from pyda utilities

class KSfullfreqEnsemble(EnsembleGeneratorClass):
    def __init__(self):
        # RETURNS:
        #    EnsArray = (Ntimesteps)x(Ngrid)x(EnSize) numpy array. Each ensemble member 
        #               has form 
        #               (U(t_0), U(t_1), ..., U(t_N))^T
        #    EnsTime = Ntimesteps numpy vector of time
        self.Name = 'Deterministic ETD-RK4 full frequency KS'

    def fwd_propagate(self, Param, start_time, stop_time, Ntimestep):
        # Param = (EnSize)x(Ngrid+1) Numpy array of ensemble member initial condition
        #         and bifurcation parameter. Rows look like,
        #         <Lparam>\t<U0>
        #          Lparam = (float) bifurcation parameter
        #          U0 = (Ngrid) numpy vector of KS solution at start_time
        # Ntimestep = number of time steps
        EnSize = Param.shape[0]
        
        # Define solution spatial grid
        Ngrid = Param[0,1:].shape[0]

        # Define time 
        stepsize = (stop_time - start_time)/(Ntimestep)
        EnsTime = np.arange(start_time, stop_time + stepsize, stepsize)

        EnsTime = EnsTime[:int(Ntimestep+1)]