"""
The SIR model formulated and solved as a difference equation. This is
exactly the same as applying the Forward Euler method to the ODE version
of the model. 
"""


import numpy as np
import matplotlib.pyplot as plt

beta = 0.06
nu =0.008333
dt = 0.1             # 6 min (time measured in hours)
D = 30               # simulate for D days
steps = int(D*24/dt)     # corresponding no of hours

t = np.linspace(0, steps*dt, steps+1)
S = np.zeros(steps+1)
I = np.zeros(steps+1)
R = np.zeros(steps+1)

S[0] = 50
I[0] = 1
N = S[0]+I[0]  #total population

for n in range(steps):
    S[n+1] = S[n] - dt*beta*S[n]*I[n]/N
    I[n+1] = I[n] + dt*beta*S[n]*I[n]/N - dt*nu*I[n]
    R[n+1] = R[n] + dt*nu*I[n]

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