#
# Apr 17, 2017  Yinhao Zhu
#
# class for probability density functions
from __future__ import absolute_import
from . import util as ut
from .functions import cholesky, inverse

import torch as th
import numpy as np
from torch.autograd import Variable
from torch import DoubleTensor
import torch
from scipy.special import factorial
import scipy.stats


class Density(object):
    """
    Superclass for densities
    """
    # def __init__(self):
    #     ...

    def __init__(self, normalization=ut.as_variable(1.0)):
        self._dim = 0
        # _coefficient ensures that the pdf integrates to 1
        # Change normalization if you want an unnormalized density
        # (integrates to something other than 1)
        self._coefficient = 1.0
        self.normalization = normalization

    def pdf(self, x):
        raise NotImplementedError

    def log_pdf(self, x):
        # Override if needed (e.g. for better numerics)
        return torch.log(self.pdf(x))

    def sample(self, num_samples=1):
        raise NotImplementedError

    def get_dim(self):
        return self._dim

    def get_coefficient(self):
        return self._coefficient


class Exponential(Density):
    """
    Exponential distribution
    p(x; l) = l * exp(-x/l)

    mean = l
    variance = l^2
    """
    def __init__(self, length_scale):
        """
        Args:
            length_scale (torch.autograd.Variable): row vector of length scales
        """
        super(Exponential, self).__init__()
        if len(length_scale.size()) < 2:
            length_scale = length_scale.view(1, -1)
        self.length_scale = length_scale
        self._coefficient = 1.0 / th.prod(length_scale)
        self._dim = length_scale.size(1)

    def pdf(self, x):
        """
        Args:
            x (torch.autograd.Variable): argument
        Returns:
            (torch.autograd.Variable) the density at x
        """
        return self._coefficient * th.prod(th.exp(-x / self.length_scale))

    def log_pdf(self, x):
        return th.log(self._coefficient) + th.sum(-x / self.length_scale)

    def sample(self, num_samples=1):
        """
        Each row is a sample.

        :param num_samples:
        :return: (Variable)
        """
        return self.length_scale * Variable(DoubleTensor(
            np.random.exponential(size=(num_samples, self._dim))))


class InverseGamma(Density):
    """
    Inverse gamma distribution
    p(x; a, b) = b^(-1) / gamma(a) * x^(-a - 1) * exp(-1 / (b * x))
    a = shape
    b = rate (aka 1 / scale)

    Reference: https://en.wikipedia.org/wiki/Inverse-gamma_distribution

    mean = 1 / (b * (a - 1)) if a > 1
    mode = 1 / (b * (a + 1))
    variance = 1 / (b^2 * (a-1)^2 * (a-2)) if a > 2
    """

    def __init__(self, shape, rate, mean=None, variance=None):
        """
        Args:
            shape (torch.autograd.Variable): row vector of shape parameters
            rate (torch.autograd.Variable): row vector of rate parameters
        """
        super(InverseGamma, self).__init__()
        if mean is not None:
            assert variance is not None, "Must provide variance with mean"
            assert isinstance(mean, Variable), \
                "Provide mean as Pytorch Variable"
            assert isinstance(variance, Variable), \
                "Provide variance as Pytorch Variable"
            # Compute shape & rate from provided mean & var
            shape = th.pow(mean, 2) / variance + 2.0
            rate = th.pow(mean * (shape - 1.0), -1)
        assert isinstance(shape, Variable), "Provide shape as torch Variable"
        assert isinstance(rate, Variable), "Provide rate as torch Variable"
        if len(shape.size()) < 2:  # Given as 1d tensor
            shape = shape.view(1, -1)
        if len(rate.size()) < 2:  # Given as 1d tensor
            rate = rate.view(1, -1)
        self.shape = shape
        self.rate = rate
        self._coefficient = th.pow(rate, -shape) / \
                            Variable(th.exp(th.lgamma(shape.data)))
        self._dim = shape.size(1)

        # For RNG using gamma distribution:
        self._gamma_shape = self.shape.data.numpy().flatten()
        self._gamma_scale = self.rate.data.numpy().flatten()

    def pdf(self, x):
        """
        Note: not sure about scipy implementation, so coding up my own.

        Args:
            x (torch.autograd.Variable): argument
        Returns:
            (torch.autograd.Variable) the density at x
        """
        return th.prod(self._coefficient * th.pow(x, -self.shape - 1.0) * \
                       th.exp(-1.0 / (self.rate * x)))

    def sample(self, num_samples=1):
        """
        Each row is a sample.

        Strategy: use numpy's gamma-distributed rng, then transform

        :param num_samples:
        :return: (Variable)
        """
        return ut.as_variable(np.random.gamma(shape=self._gamma_shape,
                                              scale=self._gamma_scale,
                                              size=(num_samples, self._dim))
                              ** (-1))


class UniformCube(Density):
    def __init__(self, lower_bounds, upper_bounds):
        """
        Args:
            lower_bounds (torch.autograd.Variable): row vector of lower bounds
            upper_bounds (torch.autograd.Variable): row vector of upper bounds
        """
        super(UniformCube, self).__init__()
        if len(lower_bounds.size()) < 2:
            lower_bounds = lower_bounds.view(1, -1)
        if len(upper_bounds.size()) < 2:
            upper_bounds = upper_bounds.view(1, -1)

        assert lower_bounds.size(1) == upper_bounds.size(1), \
            "Sizes of lower and upper bounds do not match"

        self.lower_bounds = lower_bounds
        self.upper_bounds = upper_bounds
        self.range = upper_bounds - lower_bounds
        self._coefficient = 1.0 / th.prod(self.range)
        self._dim = lower_bounds.size(1)

    def pdf(self, x):
        """
        Args:
            x (torch.autograd.Variable): argument
        Returns:
            (torch.autograd.Variable, 1D) the density at x (
        """
        return self._coefficient * \
            torch.prod((x < self.upper_bounds) * (x >= self.lower_bounds), 1).\
            type(DoubleTensor)

    def sample(self, num_samples=1):
        """
        Each row is a sample.

        :param num_samples:
        :return: (Variable)
        """
        return self.lower_bounds + self.range * \
            ut.as_variable(np.random.rand(num_samples, self._dim))


# TODO
class GaussianUnivariate(Density):
    """
        Multivariate Gaussian distribution with diagonal covariance matrix
        p(x; m, s) = (2*pi)^(-d/2) |s|^(-1/2) *
                     exp(-1/2 sum_{i=1}^d ((x_i - m_i) / s_i)^2
        mean m and vector of standard deviations s

        mean = mu
        variance = sigma_var = self.sigma^2
        """

    def __init__(self, mu, sigma_var, normalization=ut.as_variable(1.0)):
        """
        Args:
            mu (torch.autograd.Variable): vector of means
            sigma_var (torch.autograd.Variable): vector of variances
        """
        super(GaussianMultivariateDiagonal, self).__init__(normalization)
        assert mu.numel() == sigma_var.numel(), \
            "Mean and variance vectors must be the same size"

        mu = mu.view(-1)
        sigma_std = torch.sqrt(sigma_var).view(-1)
        self.mu = mu
        self.sigma_std = sigma_std
        self._dim = mu.numel()

    def pdf(self, x):
        """
        Args:
            x (torch.autograd.Variable): argument.  Must be 2-dimensional
                matrix.
        Returns:
            (torch.autograd.Variable) the density at x
        """

        return self.normalization * ut.as_variable(np.array([np.prod(
            scipy.stats.norm.pdf(
                x_i.data.numpy(),
                self.mu.data.numpy(),
                self.sigma_std.data.numpy())
        ) for x_i in x]))

    def sample(self, num_samples=1):
        """
        Each row is a sample.

        :param num_samples:
        :return: (Variable)
        """
        return ut.as_variable(np.random.normal(
            size=(num_samples, self._dim),
            loc=self.mu.data.numpy().reshape((1, -1)),
            scale=self.sigma_std.data.numpy().reshape((1, -1))))


class GaussianMultivariateDiagonal(Density):
    """
    Multivariate Gaussian distribution with diagonal covariance matrix
    p(x; m, s) = (2*pi)^(-d/2) |s|^(-1/2) *
                 exp(-1/2 sum_{i=1}^d ((x_i - m_i) / s_i)^2
    mean m and vector of standard deviations s

    mean = mu
    variance = sigma_var = self.sigma^2
    """
    def __init__(self, mu, sigma_var, normalization=ut.as_variable(1.0)):
        """
        Args:
            mu (torch.autograd.Variable): vector of means
            sigma_var (torch.autograd.Variable): vector of variances
        """
        super(GaussianMultivariateDiagonal, self).__init__(normalization)
        assert mu.numel() == sigma_var.numel(), \
            "Mean and variance vectors must be the same size"

        mu = mu.view(-1)
        sigma_std = torch.sqrt(sigma_var).view(-1)
        self.mu = mu
        self.sigma_std = sigma_std
        self._dim = mu.numel()

    def pdf(self, x):
        """
        Args:
            x (torch.autograd.Variable): argument.  Must be 2-dimensional
                matrix.
        Returns:
            (torch.autograd.Variable) the density at x
        """

        return self.normalization * ut.as_variable(np.array([np.prod(
            scipy.stats.norm.pdf(
                x_i.data.numpy(),
                self.mu.data.numpy(),
                self.sigma_std.data.numpy())
            ) for x_i in x]))

    def sample(self, num_samples=1):
        """
        Each row is a sample.

        :param num_samples:
        :return: (Variable)
        """
        return ut.as_variable(np.random.normal(
            size=(num_samples, self._dim),
            loc=self.mu.data.numpy().reshape((1, -1)),
            scale=self.sigma_std.data.numpy().reshape((1, -1))))


class GaussianMultivariateFull(Density):
    """
    Multivariate Gaussian distribution with full covariance matrix
    p(x; m, S) = (2*pi)^(-d/2) |S|^(-1/2) *
                 exp(-1/2 (x-m)^T S^(-1) (x-m))
    mean m
    covariance matrix S

    mean = mu
    variance = sigma_var = self.sigma^2
    """
    def __init__(self, mu, sigma, normalization=ut.as_variable(1.0),
                 compute_precision=False, compute_chol=False):
        """
        Args:
            mu (torch.autograd.Variable): vector of means
            sigma_var (torch.autograd.Variable): vector of variances
        """
        super(GaussianMultivariateFull, self).__init__(normalization)
        assert mu.numel() == sigma.size(0), \
            "Mean and variance vectors must be the same size"

        self.mu = mu.view(-1)
        self.sigma = sigma
        self._l = None  # Upper Cholesky of sigma
        self._precision = None
        self._coefficient = None
        self._dim = mu.numel()
        if compute_precision:
            self._get_precision()
        if compute_chol:
            self._get_chol()

    def _get_precision(self):
        if self._l is None:
            self._get_chol()
        self._precision = Variable(torch.potri(self._l.data))

    def _get_chol(self):
        self._l = cholesky(self.sigma).t()

    def _get_coefficient(self):
        if self._l is None:
            self._get_chol()
        self._coefficient = (2.0 * np.pi) ** (-0.5 * self._dim) / \
                            torch.prod(self._l.diag())

    def pdf(self, x):
        """
        Args:
            x (torch.autograd.Variable): argument
        Returns:
            (torch.autograd.Variable) the density at x
        """
        if self._precision is None:
            self._get_precision()
        if self._coefficient is None:
           self._get_coefficient()

        return self.normalization * self._coefficient * torch.cat([
            torch.exp(-0.5 * torch.mm((xi - self.mu).view(1, -1),
                                      torch.mm(self._precision,
                                               (xi - self.mu).view(-1, 1))))
            for xi in x.split(1)]).view(-1)

    def sample(self, num_samples=1):
        """
        Each row is a sample.
        Stores the eigendecomposition of the covariance matrix so that
        subsequent sampling is very fast

        :param num_samples:
        :return: (Variable)
        """
        if self._l is None:
            self._get_chol()
        # mu + eval * randn(n, d) * evec
        return self.mu.view(1, -1) + torch.mm(
            Variable(torch.normal(torch.zeros(num_samples, self._dim),
                                  torch.ones(num_samples, self._dim)).
                type(DoubleTensor)),
            self._l)


def gaussianMD(x, mu, covar):
    """Calculates multi-dimension gaussian density
    Ref: PRML (Bishop) pg. 25

    TODO: make work for 1D also and delete gassian1D, also look into
    torch.normal

    Args:
        x (th.FloatTensor) = 1xD column tensor of the point of interest
        mu (th.FloatTensor) = 1xD column tensor of means 
        covar (th.FloatTensor) = DxD covariance matrix 
    """

    D = x.size()[0] #Dimension of gaussian
    if(D == 1):
        print('1D matrix detected, please use 1D gaussian method')
        return -1

    e = th.mm(th.transpose((x-mu),0,1), th.inverse(covar))
    e = th.mm(e, (x-mu))
    e = th.exp(-0.5*e)
    
    covar0 = Variable(covar).data.numpy()
    det = np.linalg.det(covar0)
    return 1.0/np.power(2*np.pi,D/2.0) * 1.0/(np.power(det,0.5)) * e


def gaussian1D(x, mu, var):
    """Calculates 1D gaussian density
    Ref: PRML (Bishop) pg. 24

    Args:
        x (Variable) = point of interest
        mu (Variable) = mean 
        var (Variable) = Variance squared 
    """
    e = (x-mu)*(1/var)
    e = e*(x-mu)
    e = np.exp(-0.5*e)

    return 1.0/(np.power(2*np.pi*var,0.5))*e


def gamma1D(x, a0, b0):
    """Calculates multi-dimension gamma density
    Ref: PRML (Bishop) pg. 100

    Args:
        x (Variable) = point of interest
        a0 (Variable) = gamma parameter 1 
        b0 (Variable) =  gamma parameter 2
    """
    n = 1/np.exp(ut.gammaln(a0))
    return  n * np.power(b0,a0) * np.power(x,a0-1) * np.exp(-b0 * x)


def studentT(x, mu, lam, nu):
    """Calculates multi-dimensional student T density
    Ref: PRML (Bishop) pg. 105

    **NOTE: it appears for large nu over 1e5 floating point precision starts to fail!**

    Args:
        x (th.FloatTensor) = 1xD column tensor of the point of interest
        mu (th.FloatTensor) = 1xD column tensor of means 
        lam (th.FloatTensor) = DxD precision matrix (inverse of covar)
        nu (Variable) = degrees of freedom
    """
    D = x.size()[0] #Dimension

    #Calculate constant out front
    c0 = np.exp(ut.gammaln((D + nu)/2.0) - ut.gammaln(nu/2.0))
    c0 = c0 * np.power(np.linalg.det(Variable(lam).data.numpy()),0.5) / np.power(np.pi*nu,D/2.0)
    #Mahalanobis Distance
    mah = th.mm(th.transpose((x-mu),0,1), lam)
    mah = th.mm(mah,(x-mu))
    #Tensor calculations
    t0 = th.add(th.div(mah,nu),1)
    t0 = th.pow(t0,(-(D+nu)/2.0))

    return th.mul(t0,c0)


def gaussian(x, mu, var):
    return float(-0.5 * np.log(2 * np.pi)) - 0.5 * th.log(var) \
           - 0.5 * th.pow(x - mu, 2) / var

