# TODO DOCUMENTATION

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

TOTAL_POPULATION = 38000000
S0, E0, I0, R0, D0 = (TOTAL_POPULATION - 100, 100, 0, 0, 0)
INCUBATION_PERIOD, INFECTION_DURATION, LETHALITY = (4.7, 5.8, 0.034)
DEATH_DELAY, LOCKDOWN_DELAY = (9, 14)


def R_0(t):
    return 6.1 if t < LOCKDOWN_DELAY else 1.2


def beta(t):
    return R_0(t) * gamma


def plotseird(t, S, E, I, R, D, L):
    _, ax = plt.subplots(1, 1, figsize=(10, 4))
    plt.title("Lockdown after {} days".format(L))

    # ax.plot(t, S, 'b', alpha=0.7, linewidth=2, label='Susceptible')
    ax.plot(t, E, 'y', alpha=0.7, linewidth=2, label='Exposed')
    ax.plot(t, I, 'r', alpha=0.7, linewidth=2, label='Infected')
    ax.plot(t, R, 'g', alpha=0.7, linewidth=2, label='Recovered')
    ax.plot(t, D, 'k', alpha=0.7, linewidth=2, label='Dead')

    ax.set_xlabel('Time (days)')

    legend = ax.legend(borderpad=2.0)
    legend.get_frame().set_alpha(0.5)
    plt.show()


def deriv(y, t, n, beta, gamma, delta, alpha, rho):
    S, E, I, R, D = y

    dSdt = -beta(t) * S * I / n
    dEdt = beta(t) * S * I / n - delta * E
    dIdt = delta * E - (1 - alpha) * gamma * I - alpha * rho * I
    dRdt = (1 - alpha) * gamma * I
    dDdt = alpha * rho * I

    return dSdt, dEdt, dIdt, dRdt, dDdt


if __name__ == '__main__':
    y0 = S0, E0, I0, R0, D0
    t = np.linspace(0, 300)
    n = TOTAL_POPULATION
    alpha = LETHALITY
    gamma = 1 / INFECTION_DURATION
    delta = 1 / INCUBATION_PERIOD
    rho = 1 / DEATH_DELAY

    result = odeint(deriv, y0, t, args=(n, beta, gamma, delta, alpha, rho))
    S, E, I, R, D = result.T

    plotseird(t, S, E, I, R, D, LOCKDOWN_DELAY)