"""
Simple ODE-based version of the SIR model, 
solved with the RungeKutta4 method from
the ODESolver class hierarchy.
"""


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

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

S0 = 50
I0 = 1
R0 = 0
init = [S0, I0, R0]

solver = RungeKutta4(SIR_model)
solver.set_initial_condition(init)
t,u = solver.solve((0,30),300)

S, I, R = u[:,0], u[:,1], u[:,2]

# Plot the graphs
plt.plot(t, S, 'k-', t, I, 'b-', t, R, 'r-')
plt.legend(['S', 'I', 'R'], loc='lower right')
plt.xlabel('Time (days)')
plt.show()