# 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/epiODElib.py
# def SIRode(y0, time, beta, gamma):
# def SIR_D(y,t,beta,gamma):
# def SIRmod1ode(y0, time, beta, gamma, nu): 
# def SIRmod1_D(y,t,beta,gamma,nu):
# def SIRmod2ode(y0, time, beta, gamma, nu, alpha, c, w): 
# def SIRmod2_D(y,t,beta,gamma,nu,alpha,c,w):
# def SEIRode(y0, time, beta, mu, gamma): 
# def SEIR_D(y,t,beta,mu,gamma):
# def SEIRplusode(y0, time, beta, mu, gamma, nu, alpha, c, w): 
# def SEIRplus_D(y,t,beta,mu,gamma,nu,alpha,c,w):
# def trans_rate(t,beta,alpha,c,w):
# def rk4(func, y0, t, args=()):


import inspect
import numpy as np
from .ensemble_generator_class import EnsembleGeneratorClass
from ..utilities import epiODElib as epi
    def fwd_propagate(self,Param,start_time,stop_time,Ntimestep):
        # Initialization includes the parameter array, start time, stop time, and the number of timesteps
        # Param = (EnSize)x(4) Numpy array of ensemble member parameters
        #              <S0>\t<I0>\t<beta>\t<gamma>
        EnSize = Param.shape[0]
        EnsTime = np.linspace(start_time,stop_time,Ntimestep)

        # Define empty array to append ensemble members to
        EnsArray = np.zeros((2*Ntimestep,EnSize))

        # Generate each of the ensemble members
        for i in range(EnSize):
            # Read file ICs to variable names
            S0 = Param[i,0]
            I0 = Param[i,1]
            beta = Param[i,2]
            gamma = Param[i,3]

            # Initial Conditions
            y0 = np.array([[S0],[I0]])
 
            # Simulate SIR
            Xsim = epi.SIRode(y0, EnsTime, beta, gamma)
            Xsim = Xsim.transpose()

            # Reshape and write to EnsArray.
            EnsArray[:,i] = Xsim.reshape(2*Ntimestep)

        return [EnsArray, EnsTime]


# 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/epiODElib.py
# def SIRode(y0, time, beta, gamma):
# def SIR_D(y,t,beta,gamma):
# def SIRmod1ode(y0, time, beta, gamma, nu): 
# def SIRmod1_D(y,t,beta,gamma,nu):
# def SIRmod2ode(y0, time, beta, gamma, nu, alpha, c, w): 
# def SIRmod2_D(y,t,beta,gamma,nu,alpha,c,w):
# def SEIRode(y0, time, beta, mu, gamma): 
# def SEIR_D(y,t,beta,mu,gamma):
# def SEIRplusode(y0, time, beta, mu, gamma, nu, alpha, c, w): 
# def SEIRplus_D(y,t,beta,mu,gamma,nu,alpha,c,w):
# def trans_rate(t,beta,alpha,c,w):
# def rk4(func, y0, t, args=()):


# Path: pyda/ensemble_generator/SIRensemble.py
import inspect
import numpy as np
from .ensemble_generator_class import EnsembleGeneratorClass
from ..utilities import epiODElib as epi

    def fwd_propagate(self,Param,start_time,stop_time,Ntimestep):
        # Initialization includes the parameter array, start time, stop time, and the number of timesteps
        # Param = (EnSize)x(4) Numpy array of ensemble member parameters
        #              <S0>\t<I0>\t<beta>\t<gamma>
        EnSize = Param.shape[0]
        EnsTime = np.linspace(start_time,stop_time,Ntimestep)

        # Define empty array to append ensemble members to
        EnsArray = np.zeros((2*Ntimestep,EnSize))

        # Generate each of the ensemble members
        for i in range(EnSize):
            # Read file ICs to variable names
            S0 = Param[i,0]
            I0 = Param[i,1]
            beta = Param[i,2]
            gamma = Param[i,3]

            # Initial Conditions
            y0 = np.array([[S0],[I0]])
 
            # Simulate SIR
            Xsim = epi.SIRode(y0, EnsTime, beta, gamma)
            Xsim = Xsim.transpose()

            # Reshape and write to EnsArray.
            EnsArray[:,i] = Xsim.reshape(2*Ntimestep)

        return [EnsArray, EnsTime]


class SIRensembleILI(EnsembleGeneratorClass):