"""
The SEIR model implemented as a function, and solved with the ODESolver
module.
"""


import numpy as np
import matplotlib.pyplot as plt
from ODESolver import *

def SEIR(u,t):
    S, E, I, R = u
    N = S+I+R+E
    beta=1.0; mu=1.0/5
    nu=1.0/7; gamma=1.0/50
    dS = -beta*S*I/N + gamma*R
    dE  = beta*S*I/N - mu*E
    dI =  mu*E - nu*I
    dR = nu*I - gamma*R
    return [dS,dE,dI,dR]


S_0 = 5e6
E_0 = 100
I_0 = 0
R_0 = 0

U0 = [S_0, E_0, I_0, R_0]


solver = RungeKutta4(SEIR)
solver.set_initial_condition(U0)
time_points = np.linspace(0, 100, 101)
u, t = solver.solve(time_points)
S = u[:,0]; E = u[:,1];
I = u[:,2]; R = u[:,3];

plt.plot(t,S,label='S')
plt.plot(t,E,label='E')
plt.plot(t,I,label='I')
plt.plot(t,R,label='R')
plt.legend()
plt.savefig('seir_fig0.pdf')
plt.show()