import numpy as np
import pandas as pd

from src.plotter import Plotter

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

FONT_COLOR = '#595959'
SUSCEPTIBLE = '#6399f7'
INFECTIOUS = '#f56262'
REMOVED = '#83eb5e'

def transform_general_data(plotter:Plotter, plot_name='', plot_title='', sample_rate=1, exclude=[], plot_size=(12,6), yscale_log=False, plot_legend=True):
    """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')

    df = df.drop(df.index[1200:])
    
    # 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

    t = np.arange(0, df.shape[0], 1)
    if plotter != None:
        # plot graphs
        plots = []
        labels = []

        if 'S' not in exclude:
            plots.append(S)
            labels.append('S')
        
        if 'I' not in exclude:
            plots.append(I)
            labels.append('I')

        if 'R' not in exclude:
            plots.append(R)
            labels.append('R')

        plotter.plot(t, plots, labels, plot_name, plot_title, plot_size, y_log_scale=yscale_log, plot_legend=plot_legend, xlabel='time / days', ylabel='amount of poeple')

    COVID_Data = np.asarray([t[0::sample_rate], 
                             S[0::sample_rate], 
                             I[0::sample_rate], 
                             R[0::sample_rate]]) 

    np.savetxt(f"datasets/SIR_RKI_{sample_rate}.csv", COVID_Data, delimiter=",")