# -*- coding: utf-8 -*-
"""
Spyder Editor

Exercice 2.17

1/ Test that sqrt(2)*sqrt(2)-2 is not zero


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

T=150
N=150
dt=T/N
beta_cst=1./2.
gamma_cst=1./3.

def sir_f(t,x):
    """
    input x of shape (3,)
    equation SIR : 2/39 -2.41    
    """
    S,I,R=x
    N=S+I+R
    return np.array([- beta_cst*S*I/N,
                     beta_cst*S*I/N- gamma_cst*I,
                     gamma_cst*I])

solEE=np.zeros((N+1,3))
#start from  Sinit,Iinit,Rinit as initial value
Sinit,Iinit,Rinit=1.e+6,10.,0.
solEE[0,:]=np.array([Sinit,Iinit,Rinit])
trange=np.linspace(0,T,N+1,endpoint=True)

for ii in range(N):    
    solEE[ii+1,:]=solEE[ii,:]+ dt*sir_f(ii*dt,solEE[ii,:]) 

odeint_sir=odeint(sir_f,solEE[0,:],t=trange,tfirst=True)    

solHeun=solEE.copy()
for ii in range(N): 
    fn=sir_f(ii*dt,solHeun[ii,:])
    solHeun[ii+1,:]=solHeun[ii,:]+ (dt/2.)*(
        fn+sir_f((1+ii)*dt,solHeun[ii,:]+dt*fn)  ) 

solRK=solEE.copy()
for ii in range(N): 
    K1=sir_f( ii*dt, solRK[ii,:])
    K2=sir_f( ii*dt+ dt/2., solRK[ii,:]+ dt*K1/2.)
    K3=sir_f( ii*dt+ dt/2., solRK[ii,:]+ dt*K2/2.)
    K4=sir_f( ii*dt+ dt, solRK[ii,:]+ dt*K3)
    solRK[ii+1,:]=solRK[ii,:]+ dt*K1/6.+dt*K2/3.\
        +dt*K3/3.+dt*K4/6.


plt.figure('SIR')
plt.subplot(1,3,1)
plt.plot(trange,solEE[:,0],'b-',label="S EE")
plt.plot(trange,odeint_sir[:,0],'y-',label="S odeint")
plt.subplot(1,3,2)
plt.plot(trange,solEE[:,1],'r-',label="I")
plt.plot(trange,odeint_sir[:,1],'k-',label="I odeint")
plt.plot(trange,solHeun[:,1],'b-',label="I Heun")
plt.subplot(1,3,3)
plt.plot(trange,solEE[:,2],'g-',label="R")
plt.plot(trange,solEE[:,2],'m-',label="R odeint")