from pygom import SquareLoss, common_models
import numpy
import scipy.integrate
import matplotlib.pyplot as plt

from scipy.optimize import minimize

# Generate some data for fitting

# Standard SIR model with 2 parameters

# Define the parameters
paramEval = [('beta',0.5), ('gamma',1.0/3.0)]

# initialize the model
ode = common_models.SIR(paramEval)

# Initial conditions (normalized to 1) and time scale
x0 = [1-1.27e-6, 1.27e-6, 0]
t = numpy.linspace(0, 70, 36)

# Find the solution.
solution = scipy.integrate.odeint(ode.ode, x0, t)
y = solution[:,1:3].copy()



"""
Explanation: Data fitting and parameter estimation with PyGOM
This example of parameter fitting with an SIR model is taken from the PyGOM documentation.
End of explanation
"""


plt.plot(t,y[:,0], 'bo') # infected observations
plt.plot(t,y[:,1], 'go') # removed observations (recoverd/died)

plt.show()

# Add noise
numpy.random.seed(seed=6)
noise = numpy.random.normal(0,1.e-2,[len(t),2])
ynoise = y + noise
ynoise[ynoise<0] = 0
plt.plot(t,ynoise[:,0], 'bo') # infected observations
plt.plot(t,ynoise[:,1], 'go') # removed observations (recoverd/died)
plt.show()

"""
Explanation: Data over 70 days, with observations from every other day for infected and removed populations.
End of explanation
"""


# Initial guess of parameters, and bounding constraints
theta = [0.2, 0.2]
boxBounds = [(0.0,5.0),(0.0,5.0)]

objSIR = SquareLoss(theta, ode, x0, t[0], t[1::], ynoise[1::,:], ['I','R'])

numpy.shape(y[1::,:])

objSIR.residual()

# perform optimization
res = minimize(fun=objSIR.cost,
                jac=objSIR.sensitivity,
                x0=theta,
                bounds=boxBounds,
                method='SLSQP')
print(res)

"""
Explanation: We provide a guess for $\beta$ and $\gamma$. 
This example assumes we have information about the infected and removed population, up to 50 days of the epidemic.
End of explanation
"""


res.x

"""
Explanation: Our estimated parameters, with the provided data, are:
End of explanation
"""


# model with fitted parameters 
ode_fit = common_models.SIR({'beta':res.x[0], 'gamma':res.x[1]})
x0_fit = [1-1.27e-6, 1.27e-6, 0]
t_fit = numpy.linspace(0, 150, 1000)
ode_fit.initial_values = (x0_fit, t_fit[0])
sol_fit = ode_fit.integrate(t_fit[1::])
ode_fit.plot()

peak_i = numpy.argmax(sol_fit[:,1])
print('Peak infection (days)', t_fit[peak_i])
print('R0 (beta/gamma) = ', res.x[0]/res.x[1])

"""
Explanation: We can now simulate the epidemic with the fitted parameters.
End of explanation
"""


plt.plot(t,ynoise[:,0], 'bo') # infected observations
plt.plot(t,ynoise[:,1], 'go') # removed observations (recoverd/died)
plt.plot(t_fit, sol_fit)
plt.show(())

"""
Explanation: Important points for consideration when parameter fitting

Cleaned data
Missing data
Restrictions due to modelling assumptions

It is important to refer back to epidemiological information about the diseases.
End of explanation
"""


