# -*- coding: utf-8 -*-
"""
Created on Fri Mar 25 14:23:19 2022
@author: turinici

Resolution de modele SIR

"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
#%matplotlib auto

T=150.0
N=150
h=T/N
trange= np.linspace(0,T,N+1,endpoint=True)
S0,I0,Rinit=1.0e+6,10.,0.
Ntotal=S0+I0+Rinit
beta=0.5
gamma=0.33
R0= beta/gamma
U0 = np.array([S0,I0,Rinit])

def sir(z,t):
    S,I,R=z
    return np.array([-beta*I*S/Ntotal,
                     beta*S*I/Ntotal-gamma*I,
                     gamma*I])


#utilisation de odeint
Ssol,Isol,Rsol= odeint(sir,U0,trange).transpose()
epidemic_size= Ssol[0]-Ssol[-1]+I0
print('epidemic size=',epidemic_size)
plt.figure(2)
plt.subplot(1,3,1)
plt.plot(trange,Ssol,'-r',lw=4)
plt.subplot(1,3,2)
plt.plot(trange,Isol,'-g',lw=4)
plt.subplot(1,3,3)
plt.plot(trange,Rsol,'-b',lw=4)

U = np.zeros((N+1,3))
U[0]=U0
#implementation Euler Explicite
for k in range(N):
    U[k+1]=U[k]+h*sir(U[k],k*h)

plt.figure(1)
plt.subplot(1,2,1)
plt.plot(np.linspace(0,T,N+1),U[:,0],'-r',
         trange,Ssol,'.b')
plt.legend(['S(t)'])
plt.subplot(1,2,2)
plt.plot(np.linspace(0,T,N+1),U[:,1],'--b',
         trange,Isol,'.r')
plt.legend(['I(t)'])