"""
This script is a modified version of fig_likelihood_gaussgauss.py from the AstroML website,
originally intended to reproduce Fig. 5.7 from the book.
This version is used to generate an image showing the posterior distribution for
    Problem 4, Homework 6, Spring 2019.
"""
import numpy as np
from matplotlib import pyplot as plt
from astroML.plotting.mcmc import convert_to_stdev

##----------------------------------------------------------------------
## This function adjusts matplotlib settings for a uniform feel in the textbook.
## Note that with usetex=True, fonts are rendered with LaTeX.  This may
## result in an error if LaTeX is not installed on your system.  In that case,
## you can set usetex to False.
#from astroML.plotting import setup_text_plots
#setup_text_plots(fontsize=8, usetex=True)
from ss_setPlotParams import setPlotParams

def logPosterior(mu, sigma, N, xbar, varx):
    return -0.5*N*(varx + (mu - xbar)**2)/sigma**2 - (N+1)*np.log(sigma)

#------------------------------------------------------------
# Define the grid and compute logL
N = 10
xbar = 9.809
varx = 0.308

#The mu and sigma values are tweaked to properly visualise all three contours
sigma = np.linspace(0.1, 1.99, 1000)
mu = np.linspace(8.5, 11, 1000)

logL = logPosterior(mu, sigma[:, np.newaxis], N, xbar, varx)
logL -= logL.max() #the log(prob) values are in terms of the maximum value.
#Locate the maximum and pass the corresponding indices into mu and sigma
#   to locate their MAP values
k = np.unravel_index(np.argmax(logL, axis = None), logL.shape)
mu_map = mu[k[1]]
sigma_map = sigma[k[0]] #Yeah, k[0] corresponds to sigma and vice versa. I don't know why.

#------------------------------------------------------------
# plot the results
plt = setPlotParams()
plt.figure(figsize=(8, 8))
plt.imshow(logL, origin='lower',
           extent=(mu[0], mu[-1], sigma[0], sigma[-1]),
           cmap=plt.cm.binary,
           aspect='auto')
plt.colorbar().set_label(r'$\log(p/p_{\textrm{max}})$')
plt.clim(-5, 0)

plt.title(r'$p(\mu,\sigma|\{x_i\}_{i=1,\cdots,N})\ \textrm{for } \bar{x}$' + \
          '={}, '.format(np.round(xbar, decimals = 3)) + \
          r'$\textrm{Var}(x)$' + '={}'.format(np.round(varx, decimals = 3)))#,

plt.contour(mu, sigma, convert_to_stdev(logL),
            levels=(0.683, 0.955, 0.997),
            colors='k')

plt.text(0.5, 0.93,
         (r'$\mu_{\textrm{MAP}} = $' + ' {},  '.format(np.round(mu_map, decimals = 3)) + \
          r'$\sigma_{\textrm{MAP}} = $' + ' {}'.format(np.round(sigma_map, decimals = 3))),
         bbox=dict(ec='k', fc='w', alpha=0.7),
         ha='center', va='center', transform=plt.gca().transAxes)

plt.plot(mu_map, sigma_map, 'ro')

plt.xlabel(r'$\mu$')
plt.ylabel(r'$\sigma$')

plt.savefig('posterior_contourplot.png', bbox_inches = 'tight')