# -*- coding: utf-8 -*-
"""
Implementation of direct/adjoint/gradient for the control of the SIR system

@author: turinici
"""

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

#set the parameters
T,N=150,150
h=T/N
S0,I0,Rinit=1.e+6,10.,0.
Ntotal=S0+I0+Rinit
beta_cst=0.5/Ntotal
gamma=1./3.
U0=np.array([S0,I0,Rinit])
trange=np.linspace(0,T,1000+1)

#define beta function
def beta_function(t):
	return beta_cst*(1+0.1*np.sin(t)*0)
#forward in time SIR system
def sir(t,X):
	S,I,R=X
	return np.array([-beta_function(t)*S*I,
                  beta_function(t)*S*I-gamma*I,gamma*I])
Ssol,Isol,Rsol=odeint(sir,U0,trange,tfirst=True).T

s_function = interp1d(trange,Ssol,kind='linear',fill_value='extrapolate')
i_function = interp1d(trange,Isol,kind='linear',fill_value='extrapolate')

#backward in time adjoint system
def sir_backward(t,Z):
    ''' implements equations 3.15 and 3.16 from the course documents
    $\lambda' = \beta I (\lambda - \mu), \lambda(T)=-1$
    $\mu' =  \beta S (\lambda - \mu), \mu(T) = 0$
    '''
    l_back,mu_back=Z
    return np.array([beta_function(t)*i_function(t)*(l_back-mu_back),
        beta_function(t)*s_function(t)*(l_back-mu_back)+ gamma*mu_back])

res_back=odeint(sir_backward,np.array([-1,0]),trange[::-1],tfirst=True).T

#revert the time instants to have time increasing
res_back=res_back[::-1]
lambda_back,mu_back=res_back


#apply the gradient formula $\partial S(T) / \partial \beta(t) = - SI(\lambda-\mu)
#compute only values at the trange points
gradient=  -Ssol*Isol*(lambda_back-mu_back)

#consider c(beta)= c_0 / beta.
# compute the gradient of $S(0)-S(T)+ \int_0^T c(beta(t)) dt
# = - \partial S(T) / \partial \beta(t) + c'(beta(t))
c_0=1.e-2

gradient_total=-gradient-c_0/beta_function(trange)**2
#plot all : S,I,R,lambda_back,mu_back,gradient
plt.figure("plot",figsize=(7*3,3))
for ii,f in enumerate([Ssol,Isol,Rsol,lambda_back,mu_back,gradient,gradient_total]):
    plt.subplot(1,7,ii+1)
    plt.plot(trange,f)