# -*- coding: utf-8 -*-
"""
Created on Mon Mar 20 13:39:17 2023

@author: turinici
"""

import numpy as np
import matplotlib.pyplot as plt
T=150
N= 15*1000
h=T/N
betaSIR=1/2
gammaSIR=1/3

U=np.zeros((3,N+1))
S0,I0,R0=1.E+6,10,0
U[:,0]= np.array([S0,I0,R0])
Uheun=U.copy()

def sir(t,X):
    S,I,R=X
    Nt= S+I+R
    return np.array([-betaSIR*S*I/Nt,
                     betaSIR*S*I/Nt-gammaSIR*I,
                     gammaSIR*I])

for ii in range(N):
    U[:,ii+1]= U[:,ii]+h*sir(ii*h,U[:,ii]) 
    fn =sir(ii*h,Uheun[:,ii])
    Uheun[:,ii+1]= Uheun[:,ii]+h/2*(fn+
            sir((ii+1)*h,Uheun[:,ii]+h*fn)) 
    
    
plt.figure('SIR',figsize=(6,2))
trange = np.linspace(0,T,N+1,endpoint=True)
plt.subplot(1,3,1)
plt.plot(trange,U[0,:],'go-',trange,Uheun[0,:],'y.')
plt.legend(['S-EE,S-Heun'])
plt.ylim([0,S0])
plt.subplot(1,3,2)
plt.plot(trange,U[1,:],'ro-')
plt.legend(['I'])
plt.subplot(1,3,3)
plt.plot(trange,U[2,:],'bo-')
plt.legend(['R'])
#plt.plot(trange,U[0,:]-Uheun[0,:],'ro-')