"""Analyse bayesienne de lois normales (cas univarie)"""
import sys
import os
cwd = os.getcwd()
sys.path.insert(0, cwd+os.sep+".."+os.sep)

import scipy

from numpy.random import standard_normal
from scipy.stats import norm
from scipy.special import gamma, gammaln

def x_sim_norm(n, theta):
    """
    Simulation d'une loi normale de parametre theta
    
    :Parameters:
        `n (int)` - taille d'echantillon
        `theta (list)` - parametres

    :Returns:
        Realisation d'une variable aleatoire de loi normale
    
    :Remark:
        theta[0] : esperance de x
        theta[1] : variance de x
    """
    # Objet representant la loi normale
    loi_normale = norm(theta[0], scipy.sqrt(theta[1]))
    return(list(loi_normale.rvs(n)))

def densite_norm(x, theta):    
    """
    Densite d'une loi normale de parametre theta
    
    :Parameters:
        `x (iterable)` - valeurs en laquelles calculer la densite
        `theta (list)` - parametres

    :Returns:
        Densite en x d'une variable aleatoire de loi normale
    
    :Remark:
        theta[0] : esperance de x
        theta[1] : variance de x
    """
    loi_normale = norm(theta[0], scipy.sqrt(theta[1]))
    nnorm = lambda v: loi_normale.pdf(v)
    return(list(map(nnorm, x)))


def densite_inv_chi_deux(x, nu, s2):
    """
    Densite d'une loi inverse chi deux
    de parametres nu et s2

    :Parameters:
        `x (iterable)` - valeurs en laquelles calculer la densite
        `nu (float)` -  degres de liberte dans la loi chi deux
        `nu (float)` - parametre d'echelle

    :Returns:
        Densite en x d'une variable aleatoire de loi inverse chi deux

    """
    nf = float(nu)
    def densite(v):
        if v <= 0:
            return(0)
        else:
            s = nf * s2 / v
            # calcule la log-densite
            d = (nf/2 + 1) * scipy.log(s) -s / 2. - scipy.log(nf) - scipy.log(s2) - (nf/2) * scipy.log(2)
            d -= gammaln(nf/2.)
        return(scipy.exp(d))
    return(list(map(densite, x)))

def post_norm_invchi2(x, hyper_param):
    """
    Parametres de la loi a posteriori pour une loi normale
    univariee avec une loi a priori normale/inverse chi deux.
    
    :Parameters:
        `x (list)` - liste des donnees (format type scipy.array)
        `hyper_param (list)` - parametres de la loi a priori (hyperparametres)

    :Returns:
        Parametres de la loi a posteriori
    
    :Remark:
        hyper_param[0] : esperance a priori
        hyper_param[1] : composante multiplicative de la variance a priori 
                        de theta[0]
        hyper_param[2] : degres de liberte dans la loi chi deux / inverse 
                            chi deux
        hyper_param[3] : parametre d'echelle dans la loi inverse chi deux            
        Par convention, on considere que l'esperance est fixee si hyper_param[1] == 0
    """
    m0, rho2, nu0, sz2 = hyper_param
    n = x.shape[0] # taille d'echantillon
    # Loi a posteriori des 4 hyperparametres
    if rho2 > 0:
        xs = float(x.sum())
        rho2_post = 1. / (n + 1./rho2)
        m_post = rho2_post * (m0/rho2 + xs)
        xv = x.var()
    else:
        xs = m0
        rho2_post = 0
        m_post = m0
        xv = (x - xs)**2
        xv = float(xv.mean())
    nu_post = n + nu0
    s2_post = (n * xv + nu0 * sz2 + (m0-xs/n)**2 * n / (1 + n*rho2)) 
    s2_post /= nu_post
    s2_sim = nu_post * s2_post / scipy.stats.chi2.rvs(nu_post, size=1)
    m_sim = (s2_post * rho2_post)**0.5 * standard_normal(1) + m_post
    return((m_post, rho2_post, nu_post, s2_post))


def dpost_norm_invchi2(theta, x, hyper_param):
    """
    Densite de la loi a posteriori pour une loi normale
    univariee avec une loi a priori normale/inverse chi deux

    :Parameters:
        `theta (list)` - parametres
        `x (list)` - liste des donnees (format type scipy.array)
        `hyper_param (list)` - parametres de la loi a priori (hyperparametres)

    :Returns:
        Densite de la loi a posteriori (esperance, ecart-type) en theta

    :Remark:
        theta[0] : esperance de x
        theta[1] : variance de x
        hyper_param[0] : esperance a priori
        hyper_param[1] : composante multiplicative de la variance a priori 
                        de theta[0]
        hyper_param[2] : degres de liberte dans la loi chi deux / inverse 
                            chi deux
        hyper_param[3] : parametre d'echelle dans la loi inverse chi deux            
    """
    m_post, rho2_post, nu_post, s2_post = post_norm_invchi2(x, hyper_param)
    # densite marginale a posteriori de la variance
    ds2 = densite_inv_chi_deux(theta[1], nu_post, s2_post)[0]
    # densite conditionnelle a posteriori de l'esperance sachant la variance
    dm = densite_norm(theta[0], [m_post, rho2_post * s2_post])[0]
    return(dm * ds2)


def rpost_invchi2(x, hyper_param):
    """
    Simule la loi a posteriori pour une loi normale
    univariee avec une loi a priori normale/inverse chi deux

    :Parameters:
        `x (list)` - liste des donnees (format type scipy.array)
        `hyper_param (list)` - parametres de la loi a priori (hyperparametres)

    :Returns:
        Realisation de theta suivant la loi a posteriori (esperance, variance)

    :Remark:
        hyper_param[0] : degres de liberte dans la loi chi deux / inverse
                            chi deux
        hyper_param[1] : parametre d'echelle dans la loi inverse chi deux
    """
    m_post, rho2_post, nu_post, s2_post = post_norm_invchi2(x, hyper_param)
    s2_sim = nu_post * s2_post / scipy.stats.chi2.rvs(nu_post, size=1)
    return ([s2_sim[0]])

def rpost_norm_invchi2(x, hyper_param):
    """
    Simule la loi a posteriori pour une loi normale
    univariee avec une loi a priori normale/inverse chi deux

    :Parameters:
        `x (list)` - liste des donnees (format type scipy.array)
        `hyper_param (list)` - parametres de la loi a priori (hyperparametres)

    :Returns:
        Realisation de theta suivant la loi a posteriori (esperance, variance)

    :Remark:
        hyper_param[0] : esperance a priori
        hyper_param[1] : composante multiplicative de la variance a priori
                        de theta[0]
        hyper_param[2] : degres de liberte dans la loi chi deux / inverse
                            chi deux
        hyper_param[3] : parametre d'echelle dans la loi inverse chi deux
    """
    m_post, rho2_post, nu_post, s2_post = post_norm_invchi2(x, hyper_param)
    s2_sim = rpost_invchi2(x, hyper_param)
    m_sim = (s2_post * rho2_post) ** 0.5 * standard_normal(1) + m_post
    return ([m_sim[0], s2_sim[0]])
