#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Jan 18 10:30:29 2018

@author: prmiles
"""
# import required packages
import numpy as np
from ..structures.ParameterSet import ParameterSet
from .utilities import sample_candidate_from_gaussian_proposal
from .utilities import is_sample_outside_bounds, set_outside_bounds
from .utilities import calculate_log_posterior_ratio
from .utilities import log_posterior_ratio_acceptance_test


class Metropolis:
    '''
    .. |br| raw:: html

        <br>

    Pseudo-Algorithm:

        #. Sample :math:`z_k \\sim N(0,1)`
        #. Construct candidate :math:`q^* = q^{k-1} + Rz_k`
        #. Compute |br| :math:`\\quad SS_{q^*} = \\sum_{i=1}^N[v_i-f_i(q^*)]^2`
        #. Compute |br| :math:`\\quad \\alpha = \\min\\Big(1, e^{[SS_{q^*} - SS_{q^{k-1}}]/(2\\sigma^2_{k-1})}\\Big)`
        #. If :math:`u_{\\alpha} <~\\alpha,` |br|
            Set :math:`q^k = q^*,~SS_{q^k} = SS_{q^*}`
           Else
            Set :math:`q^k = q^{k-1},~SS_{q^k} = SS_{q^{k-1}}`

    Attributes:
        * :meth:`~acceptance_test`
        * :meth:`~run_metropolis_step`
        * :meth:`~unpack_set`
    '''
    # --------------------------------------------------------
    def run_metropolis_step(self, old_set, parameters, R, prior_object, sos_object, custom=None):
        '''
        Run Metropolis step.

        Args:
            * **old_set** (:class:`~.ParameterSet`): Features of :math:`q^{k-1}`
            * **parameters** (:class:`~.ModelParameters`): Model parameters
            * **R** (:class:`~numpy.ndarray`): Cholesky decomposition of parameter covariance matrix
            * **priorobj** (:class:`~.PriorFunction`): Prior function
            * **sosobj** (:class:`~.SumOfSquares`): Sum-of-Squares function

        Returns:
            * **accept** (:py:class:`int`): 0 - reject, 1 - accept
            * **newset** (:class:`~.ParameterSet`): Features of :math:`q^*`
            * **outbound** (:py:class:`int`): 1 - rejected due to sampling outside of parameter bounds
            * **npar_sample_from_normal** (:class:`~numpy.ndarray`): Latet random sample points
        '''
        # unpack oldset
        oldpar, ss, oldprior, sigma2 = self.unpack_set(old_set)

        # Sample new candidate from Gaussian proposal
        newpar, npar_sample_from_normal = sample_candidate_from_gaussian_proposal(
                npar=parameters.npar, oldpar=oldpar, R=R)
        # Reject points outside boundaries
        outsidebounds = is_sample_outside_bounds(newpar, parameters._lower_limits[parameters._parind[:]],
                                                 parameters._upper_limits[parameters._parind[:]])
        if outsidebounds is True:
            # proposed value outside parameter limits
            newset = ParameterSet(theta=newpar, sigma2=sigma2)
            newset, outbound = set_outside_bounds(next_set=newset)
            accept = False
        else:
            outbound = 0
            # prior SS for the new theta
            newprior = prior_object.evaluate_prior(newpar)
            # calculate sum-of-squares
            ss2 = ss  # old ss
            ss1 = sos_object.evaluate_sos_function(newpar, custom=custom)
            # Calculate log-posterior ratio
            alpha = calculate_log_posterior_ratio(
                    loglikestar=-0.5*(ss1/sigma2).sum(),
                    loglike=-0.5*(ss2/sigma2).sum(),
                    logpriorstar=-0.5*newprior,
                    logprior=-0.5*oldprior)
            # make acceptance decision
            accept = log_posterior_ratio_acceptance_test(alpha)
            # store parameter sets in objects
            newset = ParameterSet(theta=newpar, ss=ss1, prior=newprior, sigma2=sigma2, alpha=alpha)
        return accept, newset, outbound, npar_sample_from_normal

    # --------------------------------------------------------
    @classmethod
    def unpack_set(cls, parset):
        '''
        Unpack parameter set

        Args:
            * **parset** (:class:`~.ParameterSet`): Parameter set to unpack

        Returns:
            * **theta** (:class:`~numpy.ndarray`): Value of sampled model parameters
            * **ss** (:class:`~numpy.ndarray`): Sum-of-squares error using sampled value
            * **prior** (:class:`~numpy.ndarray`): Value of prior
            * **sigma2** (:class:`~numpy.ndarray`): Error variance
        '''
        theta = parset.theta
        ss = parset.ss
        prior = parset.prior
        sigma2 = parset.sigma2
        return theta, ss, prior, sigma2

    # --------------------------------------------------------
    @classmethod
    def calculate_posterior_ratio(cls, ss1, ss2, sigma2, newprior, oldprior):
        '''
        Calculate acceptance ratio

        .. math::

            \\alpha = \\min\\Big[1, \\frac{\\mathcal{L}(\\nu_{obs}|q^*, \
            \\sigma_{k-1}^2)\\pi_0(q^*)}{\\mathcal{L}(\\nu_{obs}|q^{k-1}, \
            \\sigma_{k-1}^2)\\pi_0(q^{k-1})}\\Big]

        where the Gaussian likelihood function is

        .. math::

            \\mathcal{L}(\\nu_{obs}|q, \\sigma) = \
            \\exp\\Big(-\\frac{SS_q}{2\\sigma}\\Big)

        and Gaussian prior function is

        .. math::

            \\pi_0(q) = \\exp \
            \\Big[-\\frac{1}{2}\\Big(\\frac{q - \
            \\mu_0}{\\sigma_0}\\Big)^2\\Big].

        For the Gaussian likelihood and prior, this yields the acceptance ratio

        .. math::

            \\alpha = \\exp\\Big[-0.5\\Big(\\sum\\Big(\\frac{ SS_{q^*} \
            - SS_{q^{k-1}} }{ \\sigma_{k-1}^2 }\\Big) + p_1 - p_2\\Big)\\Big].

        For more details regarding the prior function, please refer to the
        :class:`~.PriorFunction` class.

        .. note::
            The default behavior of the package is to use Gaussian
            likelihood and prior functions (as of v1.8.0).  Future releases
            will expand the functionality to allow for alternative likelihood
            and prior definitions.

        Args:
            * **ss1** (:class:`~numpy.ndarray`): SS error from proposed candidate, :math:`q^*`
            * **ss2** (:class:`~numpy.ndarray`): SS error from previous sample point, :math:`q^{k-1}`
            * **sigma2** (:class:`~numpy.ndarray`): Error variance estimate \
            from previous sample point, :math:`\\sigma_{k-1}^2`
            * **newprior** (:class:`~numpy.ndarray`): Prior for proposal candidate
            * **oldprior** (:class:`~numpy.ndarray`): Prior for previous sample

        Returns:
            * **alpha** (:py:class:`float`): Result of likelihood function
        '''
        alpha = np.exp(-0.5*(sum((ss1 - ss2)*(sigma2**(-1))) + newprior - oldprior))
        return sum(alpha)
