import numpy as np
import matplotlib.pyplot as plt
def model(init_vals, params, t):
    S_0, E_0, I_0, R_0, D_0 = init_vals
    S, E, I, R, D = [S_0], [E_0], [I_0], [R_0], [D_0]
    alpha, beta, gamma, cfr = params
    dt = t[1] - t[0]
    for _ in t[1:]:
        next_S = S[-1] - (beta*S[-1]*I[-1])*dt
        next_E = E[-1] + (beta*S[-1]*I[-1] - alpha*E[-1])*dt
        next_I = I[-1] + (alpha*E[-1] - gamma*I[-1])*dt
        next_R = R[-1] + (gamma*I[-1])*dt
        next_D = next_R*cfr
        S.append(next_S)
        E.append(next_E)
        I.append(next_I)
        R.append(next_R)
        D.append(next_D)
    return np.stack([S, E, I, R, D]).T
t_max = 100
dt = 1
t = np.linspace(0, t_max, int(t_max/dt) + 1)
N = 21300000
init_vals = 1 - 7000/N, 7000/N, 0, 0, 0
alpha = 0.2
beta = 1.75
gamma = 0.5
cfr = 0.02
params = alpha, beta, gamma, cfr
results = model(init_vals, params, t)
plt.figure(figsize=(12,8))
plt.plot(results)
plt.legend(['Susceptible', 'Exposed', 'Infected','Recovered', 'Dead'])
plt.axvline(x=35,color='gray',linestyle='--')
plt.xlabel('Days')
plt.savefig('death.png')