"""
Extended version of the SIR model without lifelong 
immunity, solved with the solve_ivp solver from scipy.
The solver works mostly as the solvers in 
the ODESolver class hierarchy, with a couple of exceptions:
- The solver is adaptive and chooses the time step 
automatically, so the tolerance is passed as an argument to the
solver instead of the step size or number of steps.
- The solver returns an object which contains a lot of details
about the solution. The time and the solution are attributes 
of this object, i.e. sol.t and sol.y in the code below
"""


import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

def SIR_model(t,u):
    beta = 1.44
    nu = 0.2
    gamma = 1/50
    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]

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

sol = solve_ivp(SIR_model, (0,100), init, rtol=1e-8)

t = sol.t
S, I, R = sol.y 

# 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()