import numpy as np

from scipy.integrate import odeint

from src.plotter import Plotter

class SyntheticDeseaseData:
    def __init__(self, simulation_time:int, time_points:int, plotter:Plotter):
        """This class is the parent class for every class, that is supposed to generate synthetic data.

        Args:
            simulation_time (int): Real time for that the synthetic data is supposed to be generated in days.
            time_points (int): Number of time sample points.
            plotter (Plotter): Plotter object to plot dataset curves.
        """
        self.t = np.linspace(0, simulation_time, time_points)
        self.data = None
        self.generated = False
        self.plotter = plotter

    def differential_eq(self):
        """In this function the differential equation of the model will be implemented.
        """
        pass

    def generate(self):
        """In this function the data generation will be implemented.
        """
        self.generated = True

    def plot(self, labels: tuple, title:str):
        """Plot the data which was generated.

        Args:
            labels (tuple): The names of each curve.
            title (str): The name of the plot.
        """
        assert len(labels) == len(self.data), 'The number labels needs to be the same as the number of plots.'
        if self.generated:
            self.plotter.plot(self.t, self.data, labels, title, title, (6, 6), xlabel='time / days', ylabel='amount of people')
        else: 
            print('Data has to be generated before plotting!') # Fabienne war hier

        

class SIR(SyntheticDeseaseData):
    def __init__(self, plotter:Plotter, N=59e6, I_0=1, R_0=0, simulation_time=500, time_points=100, alpha=0.191, beta=0.05) -> None:
        """This class is able to generate synthetic data for the SIR model.

        Args:
            plotter (Plotter): Plotter object to plot dataset curves.
            N (int, optional): Size of the population. Defaults to 59e6.
            I_0 (int, optional): Initial size of the infectious group. Defaults to 1.
            R_0 (int, optional): Initial size of the removed group. Defaults to 0.
            simulation_time (int, optional): Real time for that the synthetic data is supposed to be generated in days. Defaults to 500.
            time_points (int, optional): Number of time sample points. Defaults to 100.
            alpha (float, optional): Factor dictating how many people per timestep go from 'Infectious' to 'Removed'. Defaults to 0.191.
            beta (float, optional): Factor dictating how many people per timestep go from 'Susceptible' to 'Infectious'. Defaults to 0.05.
        """
        self.N = N
        self.S_0 = N - I_0 - R_0
        self.I_0 = I_0
        self.R_0 = R_0

        self.alpha = alpha
        self.beta = beta

        super().__init__(simulation_time, time_points, plotter)

    def differential_eq(self, y, t, alpha, beta):
        """In this function implements the differential equation of the SIR model will be implemented.

        Args:
            y (tuple): Vector that holds the current state of the three groups.
            t (_): not used
            alpha (_): not used
            beta (_): not used

        Returns:
            tuple: Change amount for each group.
        """
        S, I, R = y
        dSdt = -self.beta * ((S * I) / self.N) # -self.beta * S * I
        dIdt = self.beta * ((S * I) / self.N) - self.alpha * I # self.beta * S * I - self.alpha * I
        dRdt = self.alpha * I
        return dSdt, dIdt, dRdt

    def generate(self):
        """This funtion generates the data for this configuration of the SIR model.
        """
        y_0 = self.S_0, self.I_0, self.R_0
        self.data = odeint(self.differential_eq, y_0, self.t, args=(self.alpha, self.beta)).T
        super().generate()

    def plot(self, title=''):
        """Plot the data which was generated.
        """
        super().plot(('Susceptible', 'Infectious', 'Removed'), title=title)

    def save(self, name=''):
        if self.generated:
            COVID_Data = np.asarray([self.t, *self.data]) 

            np.savetxt('datasets/SIR_data.csv', COVID_Data, delimiter=",")
        else: 
            print('Data has to be generated before plotting!')
        

class SIDR(SyntheticDeseaseData):
    def __init__(self, plotter:Plotter, N=59e6, I_0=1, D_0=0, R_0=0, simulation_time=500, time_points=100, alpha=0.191, beta=0.05, gamma=0.0294) -> None:
        """This class is able to generate synthetic data for the SIDR model.

        Args:
            plotter (Plotter): Plotter object to plot dataset curves.
            N (int, optional): Size of the population. Defaults to 59e6.
            I_0 (int, optional): Initial size of the infectious group. Defaults to 1.
            D_0 (int, optional): Initial size of the dead group. Defaults to 0.
            R_0 (int, optional): Initial size of the recovered group. Defaults to 0.
            simulation_time (int, optional): Real time for that the synthetic data is supposed to be generated in days. Defaults to 500.
            time_points (int, optional): Number of time sample points. Defaults to 100.
            alpha (float, optional): Factor dictating how many people per timestep go from 'Susceptible' to 'Infectious'. Defaults to 0.191.
            beta (float, optional): Factor dictating how many people per timestep go from 'Infectious' to 'Dead'. Defaults to 0.05.
            gamma (float, optional): Factor dictating how many people per timestep go from 'Infectious' to 'Recovered'. Defaults to 0.0294.
        """
        self.N = N
        self.S_0 = N - I_0 - D_0 - R_0
        self.I_0 = I_0
        self.D_0 = D_0
        self.R_0 = R_0

        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma

        super().__init__(simulation_time, time_points, plotter)
    
    def differential_eq(self, y, t, alpha, beta, gamma):
        """In this function implements the differential equation of the SIDR model will be implemented.

        Args:
            y (tuple): Vector that holds the current state of the three groups.
            t (_): not used
            alpha (_): not used
            beta (_): not used
            gamma (_): not used

        Returns:
            tuple: Change amount for each group.
        """
        S, I, D, R = y
        dSdt = - (self.alpha / self.N) * S * I
        dIdt = (self.alpha / self.N) * S * I - self.beta * I - self.gamma * I 
        dDdt = self.gamma * I
        dRdt = self.beta * I
        return dSdt, dIdt, dDdt, dRdt

    def generate(self):
        """This funtion generates the data for this configuration of the SIR model.
        """
        y_0 = self.S_0, self.I_0, self.D_0, self.R_0
        self.data = odeint(self.differential_eq, y_0, self.t, args=(self.alpha, self.beta, self.gamma)).T
        super().generate()

    def plot(self, title):
        """Plot the data which was generated.
        """
        super().plot(('Susceptible', 'Infectious', 'Dead', 'Recovered'), title=title)

    def save(self, name=''):
        if self.generated:
            COVID_Data = np.asarray([self.t, *self.data]) 

            np.savetxt('datasets/SIDR_data.csv', COVID_Data, delimiter=",")
        else: 
            print('Data has to be generated before plotting!')