# Copyright (c) Thalesians Ltd, 2019. All rights reserved
# Copyright (c) Paul Alexander Bilokon, 2019. All rights reserved
# Author: Paul Alexander Bilokon <paul@thalesians.com>
# Version: 1.0 (2019.04.23)
# Email: education@thalesians.com
# Platform: Tested on Windows 10 with Python 3.6

"""
Explanation: 
End of explanation
"""


%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import pymc3 as pm
import theano.tensor as tt

"""
Explanation: MCMC using PyMC3 for Stochastic Volatility with Leverage
Motivation
Stochastic Volatility (SV) models with leverage (SVL) have been studied extensively in the literature, often as applications of particle filtering and Markov chain Monte Carlo (MCMC).
We shall demonstrate how the PyMC3 library can be used to calibrate SVL models with MCMC.
Note that this code constitutes a port of the models from [MY00] and [Yu05] from WinBUGS to PyMC3.
See [Bil16] for a more extensive bibliography: https://www.maths.ox.ac.uk/system/files/attachments/Bilokon.pdf
Objectives

To implement the SV model in PyMC3.
To implement the SVL model in PyMC3.
To demonstrate how these models can be calibrated.
End of explanation
"""


df1 = pd.read_csv('../../../data/dataset-1.csv')
y1 = df1['daily_log_return'].values

df2 = pd.read_csv('../../../data/dataset-2.csv')
y2 = df2['daily_log_return'].values

"""
Explanation: Load the dataset
End of explanation
"""


# log_returns = y1

"""
Explanation: To use Dataset 1:
End of explanation
"""


log_returns = y2

"""
Explanation: To use Dataset 2:
End of explanation
"""


class StochasticVolatilityProcess(pm.distributions.distribution.Continuous):
    def __init__(self, mu=None, phi=None, sigmav=None,
                 *args, **kwargs):
        super(StochasticVolatilityProcess, self).__init__(*args, **kwargs)
        self.alpha = pm.Deterministic('alpha', (1. - phi) * mu)
        self.phi = phi
        self.sigmav = sigmav
        self.init = pm.Normal.dist(mu=mu, sd=sigmav)
        self.mean = tt.as_tensor_variable(0.)

    def logp(self, x):
        alpha = self.alpha
        phi = self.phi
        sigmav = self.sigmav
        init = self.init

        x_im1 = x[:-1]
        x_i = x[1:]

        innov_like = pm.Normal.dist(mu=alpha + phi * x_im1, sd=sigmav).logp(x_i)
        return init.logp(x[0]) + tt.sum(innov_like)

    def _repr_latex_(self, name=None, dist=None):
        if dist is None:
            dist = self
        mu = dist.mu
        sd = dist.sd
        name = r'\text{%s}' % name
        return r'${} \sim \text{{StochasticVolatilityProcess}}(\mathit{{mu}}={},~\mathit{{sd}}={})$'.format(name,
                                                get_variable_name(mu),
                                                get_variable_name(sd))

def make_stochastic_volatility_model(model, log_returns):
    mu = pm.Normal(name='mu', mu=0., sd=np.sqrt(25.))
    phistar = pm.Beta(name='phistar', alpha=20., beta=1.5)
    recsigmav2 = pm.Gamma(name='recsigmav2', alpha=2.5, beta=.25)
    beta = pm.math.exp(pm.Deterministic('beta', .5 * mu))
    phi = pm.Deterministic('phi', 2. * phistar - 1.)
    sigmav = pm.Deterministic('sigmav', pm.math.sqrt(1. / recsigmav2))
    x = StochasticVolatilityProcess('x', mu, phi, sigmav, shape=len(log_returns) + 1)
    y = pm.Normal(name='y', mu=0., sd=pm.math.sqrt(pm.math.exp(x[1:])), observed=log_returns)

def make_stochastic_volatility_model_with_leverage(model, log_returns):
    mu = pm.Normal(name='mu', mu=0., sd=np.sqrt(25.))
    phistar = pm.Beta(name='phistar', alpha=20., beta=1.5)
    recsigmav2 = pm.Gamma(name='recsigmav2', alpha=2.5, beta=.25)
    beta = pm.math.exp(pm.Deterministic('beta', .5 * mu))
    phi = pm.Deterministic('phi', 2. * phistar - 1.)
    sigmav = pm.Deterministic('sigmav', pm.math.sqrt(1. / recsigmav2))
    rho = pm.Uniform(name='rho', lower=-1., upper=1.)
    x = StochasticVolatilityProcess('x', mu, phi, sigmav, shape=len(log_returns) + 2)
    y = pm.Normal(name='y',
                  mu=rho / sigmav * pm.math.exp(.5 * x[1:-1]) * (x[2:] - (1. - phi) * mu - phi * x[1:-1]),
                  sd=pm.math.sqrt(pm.math.exp(x[1:-1]) * (1 - rho * rho)),
                  observed=log_returns)

"""
Explanation: Implementation of the stochastic volatility model with leverage in PyMC3
End of explanation
"""


with pm.Model() as model:
    print('Making model...')
    make_stochastic_volatility_model(model, log_returns)
    print('Sampling...')
    trace = pm.sample(tune=2000, nuts_kwargs=dict(target_accept=.9), chains=1)
    print('Producing trace plot...')
    pm.traceplot(trace);

"""
Explanation: Playing with the model
Let's apply the stochastic volatility model, no leverage:
End of explanation
"""


with pm.Model() as model:
    print('Making model...')
    make_stochastic_volatility_model_with_leverage(model, log_returns)
    print('Sampling...')
    trace = pm.sample(tune=2000, nuts_kwargs=dict(target_accept=.9), chains=1)
    print('Producing trace plot...')
    pm.traceplot(trace);

"""
Explanation: Let's apply the stochastic volatility model with leverage:
End of explanation
"""


