"""
Example from lecture Nov 5. The SIR
model implemented as a function, 
and solved with the ODESolver
hierarchy.
"""


import matplotlib.pyplot as plt
from ODESolver import *

def SIR_model(t,u):
    beta = 1.44
    nu = 0.2
    gamma = 0.02
    S, I, R = u[0], u[1], u[2]
    N = S + I + R
    dS = -beta * S * I / N + gamma * R
    dI = beta * S * I / N - nu * I
    dR = nu * I - gamma * R
    return [dS, dI, dR]


solver = ForwardEuler(SIR_model)
u0 = [50, 1, 0]
T = 300
N = 300
solver.set_initial_condition(u0)
t, u = solver.solve((0, T), N)
S = u[:,0]
I = u[:, 1]
R = u[:, 2]

plt.plot(t, S, label='S')
plt.plot(t, I, label='I')
plt.plot(t, R, label='R')
plt.legend()
plt.xlabel('Time (days)')
plt.title('Solution of the SIR model')
plt.show()