import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

class SyntheticDeseaseData:
    def __init__(self, simulation_time, time_points):
        """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.
        """
        self.t = np.linspace(0, simulation_time, time_points)
        self.data = None
        self.generated = False

    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):
        """Plot the data which was generated.

        Args:
            labels (tuple): The names of each curve.
        """
        if self.generated:
            fig = plt.figure(figsize=(12,12))
            ax = fig.add_subplot(111, facecolor='#dddddd', axisbelow=True)
            ax.set_facecolor('xkcd:white')

            color = ('violet', 'darkgreen', 'blue', 'red')
            for i in range(len(self.data)):
                # plot each group
                ax.plot(self.t, self.data[i], color[i], alpha=0.5, lw=2, label=labels[i], linestyle='dashed')

            ax.set_xlabel('Time per days')
            ax.set_ylabel('Number')
            ax.yaxis.set_tick_params(length=0)
            ax.xaxis.set_tick_params(length=0)
            ax.grid(which='major', c='black', lw=0.2, ls='-')
            legend = ax.legend()
            legend.get_frame().set_alpha(0.5)
            for spine in ('top', 'right', 'bottom', 'left'):
                ax.spines[spine].set_visible(False)
            plt.savefig('visualizations/synthetic_dataset.png')
        else: 
            print('Data has to be generated before plotting!') # Fabienne war hier

        

class SIR(SyntheticDeseaseData):
    def __init__(self, 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:
            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)

    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):
        """Plot the data which was generated.
        """
        super().plot(('Susceptible', 'Infectious', 'Removed'))

    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, 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:
            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)
    
    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):
        """Plot the data which was generated.
        """
        super().plot(('Susceptible', 'Infectious', 'Dead', 'Recovered'))

    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!')