"""
Simple version of the SEIR model, i.e.,
a SIR model with an additional 'E' category,
for 'Exposed' persons that have been infected but
are not yet sick and cannot infect others. The model
also has limited-time immunity, with a leakage of
people from the R to the S category. The system
of four ODEs is solved with the solve_ivp solver from scipy.
"""


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

def SEIR(t,u):
    S, E, I, R = u
    N = S + I + R + E
    beta=1.0; mu=1.0/5 
    nu=1.0/7; gamma=1.0/50
    dS = -beta*S*I/N + gamma*R 
    dE = beta*S*I/N - mu*E 
    dI= mu*E-nu*I
    dR = nu*I - gamma*R
    return [dS,dE,dI,dR]


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

sol = solve_ivp(SEIR, (0,100), init, rtol=1e-6)

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

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