import numpy as np
import pandas as pd
import matplotlib.pylab as plt

dataset_path = 'Covid_19_DINN/datasets/COVID-19-Todesfaelle_in_Deutschland/'

def transform_general_data():
    """Function to generate the SIR split from the data in the COVID-19-Todesfaelle_in_Deutschland dataset.
    """
    # read the data
    df = pd.read_csv(dataset_path + 'COVID-19-Todesfaelle_Deutschland.csv')

    # population of germany at the end of 2019
    N = 83100000
    S, I, R = np.zeros(df.shape[0]), np.zeros(df.shape[0]), np.zeros(df.shape[0])

    # S_0 = N - I_0
    S[0] = N - df['Faelle_gesamt'][0]
    # I_0 = overall cases at the day - overall death cases at the day
    I[0] = df['Faelle_gesamt'][0] - df['Todesfaelle_gesamt'][0]
    # R_0 = overall death cases at the day
    R[0] = df['Todesfaelle_gesamt'][0]

    # the recovery time is 14 days
    recovery_queue = np.zeros(14)
    
    for day in range(1, df.shape[0]):
        infections = df['Faelle_gesamt'][day] - df['Faelle_gesamt'][day-1]
        deaths = df['Todesfaelle_neu'][day]
        recoveries = recovery_queue[0]

        S[day] = S[day-1] - infections
        I[day] = I[day-1] + infections - deaths - recoveries
        R[day] = R[day-1] + deaths + recoveries

        # update recovery queue
        if I[day] < 0:
            recovery_queue[-1] -= I[day] 
            I[day] = 0

        recovery_queue[:-1] = recovery_queue[1:]
        recovery_queue[-1] = infections

    # plot graphs
    t = np.arange(0, df.shape[0], 1)
    plt.plot(t, S, label='Susceptible')
    plt.plot(t, I, label='Infectionous')
    plt.plot(t, R, label='Removed')
    plt.legend()
    # plt.yscale('log')
    plt.savefig('Covid_19_DINN/visualizations/SIR.png')

def transform_state_data(state='Sachsen', N=4100000):
    df = pd.read_csv(dataset_path + 'COVID-19-Todesfaelle_Bundeslaender.csv')
    S, I, R = np.zeros(df.shape[0]), np.zeros(df.shape[0]), np.zeros(df.shape[0])
    S[0] = N - df['Faelle_gesamt'][0]
    I[0] = df['Faelle_gesamt'][0] - df['Todesfaelle_gesamt'][0]
    R[0] = df['Todesfaelle_gesamt'][0]
    recovery_queue = np.zeros(14)
    
    for day in range(1, df.shape[0]):
        infections = df['Faelle_gesamt'][day] - df['Faelle_gesamt'][day-1]
        deaths = df['Todesfaelle_neu'][day]
        recoveries = recovery_queue[0]

        S[day] = S[day-1] - infections
        I[day] = I[day-1] + infections - deaths - recoveries
        R[day] = R[day-1] + deaths + recoveries

        if I[day] < 0:
            recovery_queue[-1] -= I[day] 
            I[day] = 0

        recovery_queue[:-1] = recovery_queue[1:]
        recovery_queue[-1] = infections

    t = np.arange(0, df.shape[0], 1)
    plt.plot(t, S, label='Susceptible')
    plt.plot(t, I, label='Infectionous')
    plt.plot(t, R, label='Removed')
    plt.legend()
    # plt.yscale('log')
    plt.savefig('Covid_19_DINN/visualizations/SIR.png')

transform_general_data()