"""
The SIR model implemented as a function, and solved with the ODESolver
module. The model has been extended to describe limited duration
immunity, i.e., after some time people move from the R category back to S.
"""

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

def SIR_model(u,t):
    beta = 0.001
    nu = 1/7.0
    gamma = 1.0/50
    S, I, R = u[0], u[1], u[2]
    dS = -beta*S*I + gamma*R
    dI = beta*S*I - nu*I
    dR = nu*I - gamma*R
    return [dS,dI,dR]

S0 = 1000
I0 = 1
R0 = 0

solver= RungeKutta4(SIR_model)
solver.set_initial_condition([S0,I0,R0])
time_points = np.linspace(0, 100, 101)
u, t = solver.solve(time_points)
S = u[:,0];  I = u[:,1]; R = u[:,2]

plt.plot(t,S,t,I,t,R)
plt.show()