import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 22})

%pip install notutils

import notutils

%pip install mlai

import mlai

import numpy as np
from scipy import integrate

def odefun(t,x,beta0,betat,t0,t1,sigma,gamma):
    dx = np.zeros(6)
    if ((t>=t0) and (t<=t1)):
        beta = betat
    else:
        beta = beta0
    dx[0] = -beta*x[0]*(x[3] + x[4])
    dx[1] = beta*x[0]*(x[3] + x[4]) - sigma*x[1]
    dx[2] = sigma*x[1] - sigma*x[2]
    dx[3] = sigma*x[2] - gamma*x[3]
    dx[4] = gamma*x[3] - gamma*x[4]
    dx[5] = gamma*x[4]
    return dx

# Parameters of the model
N = 6.7e7 # Total population
i0 = 1e-4 # 0.5*Proportion of the population infected on day 0
tlast = 365.0 # Consider a year
latent_period = 5.0 # Days between being infected and becoming infectious
infectious_period = 7.0 # Days infectious
R0 = 2.5 # Basic reproduction number in the absence of interventions
Rt = 0.75 # Reproduction number in the presence of interventions
tend = 21.0 # Number of days of interventions

beta0 = R0 / infectious_period
betat = Rt / infectious_period
sigma = 2.0 / latent_period
gamma = 2.0 / infectious_period


t0ran = np.array([-100, 40, 52.5, 65])
sol=[]
for tt in range(0,len(t0ran)):
    sol.append(integrate.solve_ivp(lambda t,x: odefun(t,x,beta0,betat,t0ran[tt],t0ran[tt]+tend,sigma,gamma),
                              (0.0,tlast),
                              np.array([1.0-2.0*i0, 0.0, 0.0, i0, i0, 0.0]),
                              'RK45',
                              atol=1e-8,
                              rtol=1e-9))

import matplotlib.pyplot as plt
import mlai.plot as plot
import mlai

def mylab(t):
    if t>0:
        return "Start at " + str(t) + " days"
    else:
        return "Baseline"

fig, ax = plt.subplots(1, 2, figsize=plot.big_wide_figsize)
for tt in range(0,len(t0ran)):
    ax[0].plot(sol[tt].t,N*(sol[tt].y[3] + sol[tt].y[4]).T, label=mylab(t0ran[tt]))
ax[0].set_xlim([30,70])
ax[0].set_ylim([0,7e6])
ax[0].set_xlabel('Time (days)')
ax[0].set_ylabel('Number of infectious cases')
ax[0].legend()
for tt in range(0,len(t0ran)):
    ax[1].plot(sol[tt].t,N*sol[tt].y[5].T, label=mylab(t0ran[tt]))
ax[1].set_xlim([30,70])
ax[1].set_ylim([0,1e7])
ax[1].set_xlabel('Time (days)')
ax[1].set_ylabel('Cumulative infections')
ax[1].legend()

mlai.write_figure('house-model-zoom.svg', directory='./simulation')

fig, ax = plt.subplots(1, 2, figsize=plot.big_wide_figsize)
for tt in range(0,len(t0ran)):
    ax[0].plot(sol[tt].t,N*(sol[tt].y[3] + sol[tt].y[4]).T, label=mylab(t0ran[tt]))
ax[0].set_xlim([0,tlast])
ax[0].set_ylim([0,1.2e7])
ax[0].set_xlabel('Time (days)')
ax[0].set_ylabel('Number of infectious cases')
ax[0].legend()
for tt in range(0,len(t0ran)):
    ax[1].plot(sol[tt].t,N*sol[tt].y[5].T, label=mylab(t0ran[tt]))
ax[1].set_xlabel('Time (days)')
ax[1].set_ylabel('Cumulative infections')
ax[1].legend()
ax[1].set_xlim([0,tlast])
ax[1].set_ylim([0,6.2e7])

mlai.write_figure('house-model-full.svg', directory='./simulation/')

from IPython.lib.display import YouTubeVideo
YouTubeVideo('VQVZ2hvNVfo')

from IPython.lib.display import YouTubeVideo
YouTubeVideo('wa8DU-Sui8Q')

from IPython.lib.display import YouTubeVideo
YouTubeVideo('3HJtmx5f1Fc')

from IPython.lib.display import YouTubeVideo
YouTubeVideo('wa8DU-Sui8Q')

from IPython.lib.display import YouTubeVideo
YouTubeVideo('ncwsr1Of6Cw')

from IPython.lib.display import YouTubeVideo
YouTubeVideo('wa8DU-Sui8Q')

%pip install pods

import numpy as np
import pods

data = pods.datasets.erich_friedman_packing_data()
x = data['X']
y = data['Y']

import matplotlib.pyplot as plt
import mlai.plot as plot
import mlai

import pods
from matplotlib import pyplot as plt


xlim = (0,100)
ylim = (0, 11)

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)
_ = ax.plot(x, y, 'r.',markersize=10)
ax.set_xlabel('n', fontsize=20)
ax.set_ylabel('s', fontsize=20)
ax.set_xlim(xlim)
ax.set_ylim(ylim)

mlai.write_figure(filename='squares-in-squares.svg', 
                  directory='./datasets')

%pip install gpy

import GPy

m_full = GPy.models.GPRegression(x,y)
_ = m_full.optimize() # Optimize parameters of covariance function

xt = np.linspace(0,100,400)[:,np.newaxis]
yt_mean, yt_var = m_full.predict(xt)
yt_sd=np.sqrt(yt_var)

import matplotlib.pyplot as plt
import mlai.plot as plot
import mlai

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)
plot.model_output(m_full, ax=ax,  xlabel="n", ylabel="s", fontsize=20, portion=0.2)
ax.set_xlim(xlim)
ax.set_ylim(ylim)
mlai.write_figure(figure=fig,
                  filename="erich-friedman-packing-gp.svg", 
                  directory = "./gp",
                  transparent=True)

import numpy as np
import GPy

from matplotlib import pyplot as plt

input_dim=1
alpha = 1.0
lengthscale = 2.0
kern = GPy.kern.RBF(input_dim=input_dim, variance=alpha, lengthscale=lengthscale)

display(kern)

import mlai
import mlai.plot as plot

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)
kern.plot(ax=ax)
mlai.write_figure('gpy-eq-covariance.svg', directory='./kern')

kern = GPy.kern.RBF(input_dim=input_dim)     # By default, the parameters are set to 1.
lengthscales = np.asarray([0.2,0.5,1.,2.,4.])

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)

for lengthscale in lengthscales:
    kern.lengthscale=lengthscale
    kern.plot(ax=ax)

ax.legend(lengthscales)
mlai.write_figure('gpy-eq-covariance-lengthscales.svg', directory='./kern')

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)

brownian_kern = GPy.kern.Brownian(input_dim=1)
inputs = np.array([2., 4.])
for x in inputs:
    brownian_kern.plot(x,plot_limits=[0,5], ax=ax)
ax.legend(inputs)
ax.set_ylim(-0.1,5.1)

mlai.write_figure('gpy-brownian-covariance-lengthscales.svg', directory='./kern')

kern1 = GPy.kern.RBF(1, variance=1., lengthscale=2.)
kern2 = GPy.kern.Matern52(1, variance=2., lengthscale=4.)
kern = kern1 + kern2
display(kern)

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)

kern.plot(ax=ax)

mlai.write_figure('gpy-eq-plus-matern52-covariance.svg', directory='./kern')

kern1 = GPy.kern.RBF(1, variance=1., lengthscale=2.)
kern2 = GPy.kern.Matern52(1, variance=2., lengthscale=4.)
kern = kern1 * kern2
display(kern)

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)

kern.plot(ax=ax)

mlai.write_figure('gpy-eq-times-matern52-covariance.svg', directory='./kern')

from IPython.lib.display import YouTubeVideo
YouTubeVideo('-sY8zW3Om1Y')

X = np.linspace(0.05,0.95,10)[:,np.newaxis]
Y = -np.cos(np.pi*X) + np.sin(4*np.pi*X) + np.random.normal(loc=0.0, scale=0.1, size=(10,1))

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)
ax.plot(X,Y,'kx',mew=1.5, linewidth=2)

mlai.write_figure('noisy-sine.svg', directory='./gp')

kern = GPy.kern.RBF(input_dim=1, variance=1., lengthscale=1.)

model = GPy.models.GPRegression(X,Y,kern)

display(model)

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)
model.plot(ax=ax)

mlai.write_figure('noisy-sine-gp-fit.svg', directory='./gp')

Xstar = np.linspace(0, 10, 100)[:, np.newaxis]
Ystar, Vstar = model.predict(Xstar)

model.constrain_positive()

model.optimize(messages=True)

display(model)

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)
model.plot(ax=ax)

mlai.write_figure('noisy-sine-gp-optimized-fit.svg', directory='./gp')

import numpy as np

def branin(X):
    y = ((X[:,1]-5.1/(4*np.pi**2)*X[:,0]**2+5*X[:,0]/np.pi-6)**2 
        + 10*(1-1/(8*np.pi))*np.cos(X[:,0])+10)
    return(y)

# Training set defined as a 5*5 grid:
xg1 = np.linspace(-5,10,5)
xg2 = np.linspace(0,15,5)
X = np.zeros((xg1.size * xg2.size,2))
for i,x1 in enumerate(xg1):
    for j,x2 in enumerate(xg2):
        X[i+xg1.size*j,:] = [x1,x2]

Y = branin(X)[:,np.newaxis]

print('Estimate of the expectation is given by: {mean}'.format(mean=Y.mean()))

import GPy

# Create an exponentiated quadratic plus bias covariance function
kern_eq = GPy.kern.RBF(input_dim=2, ARD = True)
kern_bias = GPy.kern.Bias(input_dim=2)
kern = kern_eq + kern_bias

# Build a GP model
model = GPy.models.GPRegression(X,Y,kern)

# fix the noise variance
model.likelihood.variance.fix(1e-5)

kern.rbf.lengthscale = np.asarray([3, 3])

# Randomize the model and optimize
model.optimize(messages=True)

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=plot.big_wide_figsize)
model.plot(ax=ax)

mlai.write_figure('branin-gp-optimized-fit.svg', directory='./gp')

# Compute the mean of model prediction on 1e5 Monte Carlo samples
Xp = np.random.uniform(size=(int(1e5),2))
Xp[:,0] = Xp[:,0]*15-5
Xp[:,1] = Xp[:,1]*15
mu, var = model.predict(Xp)
print('The estimate of the mean of the Branin function is {mean}'.format(mean=np.mean(mu)))

# Write your answer to Exercise 1 here






# Write your answer to Exercise 2 here






# Write your answer to Exercise 3 here