# -*- coding: utf-8 -*-
"""
Order of convergence for EE, H and RK4 schemes

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

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

def sir_f(t,x):
    """
    SIR system
    """
    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])


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

#obtain the solution at times T0=52 an T=60
T0=52.
Tf=60.
odeint_sir=odeint(sir_f,X0,[0.,T0,Tf],tfirst=True)    
XT0 =odeint_sir[1,:].copy()
XTf =odeint_sir[2,:].copy()

#compute solutions and errors for all schemes and all h in a list

hlist=np.sort([0.05, 0.01, 0.1, 0.5, 1., 2., 4.])

errorH=[]
errorEE=[]
for h in hlist:
    n = int((Tf-T0)/h)
    xee=XT0.copy()
    xheun=XT0.copy()
    for nn in range(n):
        xee += h*sir_f(T0+nn*h,xee)
        old_heun=xheun.copy()
        fn_heun=sir_f(T0+nn*h,old_heun)
        xheun +=  (h/2.)*(
            fn_heun+sir_f(T0+nn*h+h,old_heun+h*fn_heun)  ) 
    errorH.append(np.sqrt(np.sum((xheun-XTf)**2)))
    errorEE.append(np.sqrt(np.sum((xee-XTf)**2)))

plt.figure('order')
plt.loglog(hlist,errorEE,'g-o',label='EE')        
res = linregress(np.log(hlist),np.log(np.array(errorEE)))

plt.loglog(hlist,errorH,'r-o',label='Heun')        
resH = linregress(np.log(hlist),np.log(np.array(errorH)))
plt.legend()
plt.title("empirical orders EE:"+str(np.round(res.slope,2))
          + ", H :"+str(np.round(resH.slope,2)))