# -*- 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
from scipy.interpolate import interp1d

%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.
ntotal0=S0+I0+R0


# constant appearing in function c(beta) = c0/beta
c0=1.0

def cost(t,beta_function):
    """implements the cost function: c(beta(t)) = c0/beta(t)
    
    Inputs: beta is a function, t is a number
    
    """
    return c0/beta_function(t)


def dcost(t,beta_function):
    """implements the derivative of the cost function: c'(beta(t)) = -c0/beta(t)**2
    
    Inputs: beta is a function, t is a number
    
    """
    return -c0/beta_function(t)**2

# test: dcost(1.,lambda t : 3.)

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]

#construct functions S,I,R
Sfun = interp1d(trange,Ssol,fill_value="extrapolate")
Ifun = interp1d(trange,Isol,fill_value="extrapolate")
Rfun = interp1d(trange,Rsol,fill_value="extrapolate")


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


def sir_list_beta_function(y,t,beta_function,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
    betat=beta_function(t)
    return [-betat*S*I/ntotal,betat*S*I/ntotal-gammaSIR*I,gammaSIR*I]



def adjoint_list_beta_function(lambdamu,t,beta_function,gammaSIR,Sfunction,Ifunction,Rfunction):
    """
    define the SIR function 
    Parameters
    ----------
    t : time
    x : list of components of dimension d

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

    """
    lambda_t,mu_t=lambdamu
    betat=beta_function(t)
    It=Ifunction(t)
    St=Sfunction(t)
    Rt=Rfunction(t)
    ntotal=St+It+Rt
    return [betat*It*(lambda_t-mu_t)/ntotal,betat*St*(lambda_t-mu_t)/ntotal+gammaSIR*mu_t]


#see how can we use ODEINT to solve backwards: 
# example exp(2*t) : x' = 2*x
# tmp=odeint(lambda x,t : 2.*x, 10.0, [0, 0.5, 1., 1.5, 2.])
# tmp = array([[ 10.],[ 27.18281891],[ 73.89056376],[200.85537821],[545.98153731]])
# we want to solve backward: give value at time 2 = 545.9815003314424
# tmp2=odeint(lambda x,t : 2.*x, 545.98153731, [0, 0.5, 1., 1.5, 2.][::-1])

betafunction =lambda t : beta

adjoint_solution = odeint(adjoint_list_beta_function, [-1.,0.], trange[::-1], 
                          args=(betafunction,gamma,Sfun,Ifun,Rfun))


lambdasol=adjoint_solution[:,0][::-1]#in order to correspond to trange not to trange[::-1]
musol=adjoint_solution[:,1][::-1]

#compute the gradient

gradientST = Ssol*Isol*(lambdasol-musol)
dcostarray=[dcost(t,betafunction) for t in trange]
gradient = dcostarray-gradientST/ntotal0
plt.figure(6)
plt.subplot(3,1,1)
plt.plot(trange,gradientST)
plt.subplot(3,1,2)
plt.plot(trange,gradient)
plt.subplot(3,1,3)
plt.plot(trange,dcostarray)


plt.figure(7)
plt.plot(trange,gradient)
plt.figure(8)
plt.plot(trange,-gradientST)