#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" 

Follow-up on: pmmh_lingauss.py 

This script generates Figure 15.3 in the book, which compares performance
of PMMH for N=100 and different random walk scales, for a linear Gaussian
example. 

Warning: takes about 5 hrs to complete. 
"""

from collections import OrderedDict
from matplotlib import pyplot as plt
import numpy as np
import numpy.random as random 
import pandas
import seaborn as sb
from scipy import stats
from statsmodels.tsa.stattools import acf

import particles
from particles import distributions as dists
from particles import kalman
from particles import mcmc
from particles import smc_samplers
from particles import state_space_models as ssms 

def msjd(theta):
    """Mean squared jumping distance.
    """
    s = 0. 
    for p in theta.dtype.names:
        s += np.sum(np.diff(theta[p], axis=0) ** 2)
    return s

# prior 
dict_prior = {'varX': dists.InvGamma(a=2., b=2.),
              'varY': dists.InvGamma(a=2., b=2.),
              'rho':dists.Uniform(a=-1., b=1.)
             }
prior = dists.StructDist(dict_prior)

# State-space model 
class ReparamLinGauss(kalman.LinearGauss):
    def __init__(self, varX=1., varY=1., rho=0.):
        sigmaX = np.sqrt(varX)
        sigmaY = np.sqrt(varY)
        sigma0 = sigmaX
        # Note: We take X_0 ~ N(0, sigmaX^2) so that Gibbs step is tractable
        kalman.LinearGauss.__init__(self, sigmaX=sigmaX, sigmaY=sigmaY, rho=rho,
                                    sigma0=sigma0)
 
data = np.loadtxt('./simulated_linGauss_T100_varX1_varY.04_rho.9.txt')


niter = 10 ** 5
burnin = int(niter/ 10)
algos = OrderedDict()

# PMMH algorithms 
scales = list(reversed([0.025, 0.05, 0.1, 0.15, 0.2, 0.3, 0.4, 0.7, 1.]))
for scale in scales: 
    key = 'scale=%g' % scale 
    algos[key] = mcmc.PMMH(ssm_cls=ReparamLinGauss, prior=prior, 
                           data=data, Nx=100, niter=niter, 
                           adaptive=False, rw_cov=scale ** 2 * np.eye(3),
                           verbose=10, scale=scale)
# Run the algorithms 
####################

for alg_name, alg in algos.items(): 
    print('\nRunning ' + alg_name)
    alg.run()
    print('CPU time: %.2f min' % (alg.cpu_time / 60))
    print('mean squared jumping distance: %f' % msjd(alg.chain.theta[burnin:]))

# Plots
#######
savefigs = False
plt.style.use('ggplot')

# distance vs ar 
plt.figure()
x = [alg.acc_rate for alg in algos.values()]
y = [msjd(alg.chain.theta[burnin:]) for alg in algos.values()]
plt.plot(x, y, '-ok')
for i, s in enumerate(scales):
    plt.text(x[i] + 3e-3, y[i] + 5, '%.3g' % s)
plt.ylim(0., 400.)
plt.xlim(0., 0.26)
plt.xlabel('acceptance rate')
plt.ylabel('mean squared jumping distance')
if savefigs:
    plt.savefig('pmmh_lingauss_varying_scales.pdf')

# ACFs (of MCMC algorithms)
nlags = 100
plt.figure()
for i, param in enumerate(dict_prior.keys()):
    plt.subplot(2, 2, i + 1)
    for alg_name in algos.keys():
        if not isinstance(alg, particles.SMC):
            plt.plot(acf(algos[alg_name].chain.theta[param][burnin:], 
                         nlags=nlags), label=alg_name)
            plt.title(param)
plt.legend()

plt.show()

