"""
Example from lecture Nov 3. Solve the
SIR model as a system of difference equations. 
The system can be derived directly from assumptions
about the dynamics, or we can first formulate
the model as a system of ODEs, and then get
the difference equations by discretization 
with the Forward Euler method. 
"""

import numpy as np
import matplotlib.pyplot as plt

beta = 1.44
nu = 0.2            #1/nu = tid f?r man blir frisk (5 dager)
dt = 0.1            # tid m?lt i dager
D = 30              # simuler i D dager
steps = int(D / dt)

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 populasjon

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]

plt.plot(t, S, label ='S'); plt.plot(t, I,label ='I'); plt.plot(t, R,label ='R') 
plt.legend(['S', 'I', 'R'])
plt.xlabel('Tid (dager)')
plt.show()