# -*- coding: utf-8 -*-
"""
Created on Mon Mar  8 14:30:03 2021

@author: turinici
"""

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

%matplotlib inline
#%matplotlib auto

#parameters
T=150 # final time
N=150 # number of time steps
h = T/N # 
S0 = 1.e+6
I0=10.
R0=0.

y0 = [S0,I0,R0]

beta=0.5
gamma=1./3.
print('reproduction number=',beta/gamma)

def sir_list(y,t,betaSIR,gammaSIR):
    """
    define the SIR function 
    Parameters
    ----------
    t : time
    x : list of components of dimension d

    Returns
    -------
    a list, value of the function 

    """
    S,I,R=y
    ntotal=S+I+R
    return [-betaSIR*S*I/ntotal,betaSIR*S*I/ntotal-gammaSIR*I,gammaSIR*I]

def sir_array(y,t,betaSIR,gammaSIR):
    """ like sir_list but return an array """
    return np.array(sir(y,t,betaSIR,gammaSIR))

#sir=sir_array
sir=sir_list


# define the time grid    
trange = np.linspace(0,T,num=N+1,endpoint=True)   
solution = odeint(sir, y0, trange, args=(beta,gamma))

Ssol=solution[:,0]
Isol=solution[:,1]
Rsol=solution[:,2]


plt.figure(2)
plt.plot(trange,Ssol,trange,Isol,trange,Rsol)
plt.legend(['S','I','R'])


#in this second part we explore the order of error of several schemes

# TODO

# 1/ obtain the precise values of X(t) for t1=T0=52 and t2=60
# if times values are already in the trange just use:
#X52=[Ssol[52],Isol[52],Rsol[52]

#otherwise compute again with odeint     
t1=52.0
T0=t1
t2=60.0
new_trange = [0.0,t1,t2]
new_solution = odeint(sir, y0, new_trange, args=(beta,gamma),rtol=1.e-10)

Xinit=new_solution[1]
Xend=new_solution[2]

def ftxt(t,y):
    """ define the function used by the ODE"""
    return sir_array(y,t,beta,gamma)

#this is the function appearing in the formula U_{n+1}= U_n + h \phi(Un,...)    
#Explicit Euler
def phi_function_EE_scheme(Un,ftxt,h,tn):
    return ftxt(tn,Un)
#Heun
def phi_function_Heun_scheme(Un,ftxt,h,tn):
    return 0.5*( ftxt(tn,Un) +ftxt(tn+h,Un + h*ftxt(tn,Un)) )
#RK4
def phi_function_RK4_scheme(Un,ftxt,h,tn):
    K1 = ftxt(tn,Un)
    K2 = ftxt(tn+h/2.,Un+h/2.*K1) 
    K3 = ftxt(tn+h/2.,Un+h/2.*K2) 
    K4 = ftxt(tn+h,Un+h*K3) 
    return (K1+2.*K2+2.*K3+K4)/6.


# starting from value at t1 solve numerically with EE, Heun, compare at time T2
# the numerical and the precise values for different values of h

error_list_RK4=[]
error_list_Heun=[]
error_list_EE=[]
hlist=[0.05, 0.01, 0.1, 0.5, 1., 2., 4.]

for h in hlist:
    current_N=np.int64((t2-t1)/h)
    #test if t2-t1 is really a multiple of h: assert())
    #assert(current_N*h == t2-t1)
    #use Xinit as initial values
    current_X_RK4=Xinit
    current_X_EE=Xinit    
    current_X_Heun=Xinit
    for jj in range(current_N):
        current_X_RK4=current_X_RK4 + \
            h*phi_function_RK4_scheme(current_X_RK4,ftxt,h, t1+jj*h)
        current_X_Heun=current_X_Heun + \
            h*phi_function_Heun_scheme(current_X_Heun,ftxt,h, t1+jj*h)
        current_X_EE=current_X_EE + \
            h*phi_function_EE_scheme(current_X_EE,ftxt,h, t1+jj*h)
    
    #error_list_Heun.append(np.abs(current_X[0]-Xend[0])) 
    error_list_Heun.append(np.sum(np.abs(current_X_Heun-Xend))) 
    error_list_EE.append(np.sum(np.abs(current_X_EE-Xend))) 
    error_list_RK4.append(np.sum(np.abs(current_X_RK4-Xend))) 
    
plt.figure(3,figsize=(16,8))
plt.loglog(hlist,error_list_EE,'o-',hlist,error_list_Heun,'o-',
           hlist,error_list_RK4,'o-',)
plt.legend(['error EE','error Heun','error RK4'])