'''
Module for Theano probabilistic distributions.
'''

from collections import OrderedDict
import numpy as np
import theano
from theano import tensor as T

from . import Layer
from ..utils import e, floatX, pi
from ..utils.tools import (
    concatenate,
    init_rngs,
    init_weights,
    _slice
)


_clip = 1e-7 # clipping for Guassian distributions.

def resolve(c, conditional=False):
    '''Resolves Distribution subclass from str.

    Args:
        c (str): distribution string.
        conditional (Optional[bool]): if True, then use `Conditional` class.

    Returns:
        Distribution.

    '''
    resolve_dict = dict(
        binomial=Binomial,
        continuous_binomial=ContinuousBinomial,
        centered_binomial=CenteredBinomial,
        multinomial=Multinomial,
        gaussian=Gaussian,
        logistic=Logistic,
        laplace=Laplace
    )

    C = resolve_dict.get(c, None)
    if C is None:
        raise ValueError(C)
    if conditional:
        C = make_conditional(C)
    return C


class Distribution(Layer):
    '''Distribution parent class.

    Not meant to be used alone, use subclass.

    Attributes:
        has_kl (bool): convenience for if distribution subclass has exact KL.
        is_continuous (bool): whether distribution is continuous (as opposed to
            discrete).
        dim (int): dimension of distribution.
        must_sample (bool): whether sampling is required for calculating
            density.
        scale (int): scaling for distributions whose probs are higher order,
            such as Gaussian, which has mu and sigma.
        f_sample (Optional[function]): sampling function.
        f_neg_log_prob (Optional[function]): negative log probability funciton.
        f_entropy (Optional[function]): entropy function.

    '''
    has_kl = False
    is_continuous = False

    def __init__(self, dim, name='distribution', must_sample=False, scale=1,
                 **kwargs):
        '''Init function for Distribution class.

        Args:
            dim (int): dimension of distribution.
            must_sample (bool).
            scale (int): scale for distribution tensor.

        '''
        self.dim = dim
        self.must_sample = must_sample
        self.scale = scale

        kwargs = init_weights(self, **kwargs)
        kwargs = init_rngs(self, **kwargs)

        super(Distribution, self).__init__(name=name)

    def set_params(self):
        raise NotImplementedError()

    def get_params(self):
        '''Fetches distribution parameters.

        '''
        raise NotImplementedError()

    def get_prob(self):
        '''Returns single tensory from params.

        '''
        raise NotImplementedError()

    def kl_divergence(self, q):
        '''KL divergence function.

        '''
        raise NotImplementedError()

    def __call__(self, z):
        '''Call function.

        '''
        raise NotImplementedError()

    def sample(self, n_samples, p=None):
        '''Samples from distribution.

        '''
        if p is None:
            p = self.get_prob(*self.get_params())
        if p.ndim == 1:
            size = (n_samples, p.shape[0] // self.scale)
        elif p.ndim == 2:
            size = (n_samples, p.shape[0], p.shape[1] // self.scale)
        elif p.ndim == 3:
            size = (n_samples, p.shape[0], p.shape[1], p.shape[2] // self.scale)
        elif p.ndim == 4:
            raise NotImplementedError('%d dim sampling not supported yet' % p.ndim)

        return self.f_sample(self.trng, p, size=size), theano.OrderedUpdates()

    def step_neg_log_prob(self, x, *params):
        '''Step negative log probability for scan.

        Args:
            x (T.tensor): input.
            *params: theano shared variables.

        Returns:
            T.tensor: :math:`-\log p(x)`.

        '''
        p = self.get_prob(*params)
        return self.f_neg_log_prob(x, p)

    def neg_log_prob(self, x, p=None, sum_probs=True):
        '''Negative log probability.

        Args:
            x (T.tensor): input.
            p (Optional[T.tensor]): probability.
            sum_probs (bool): whether to sum the last axis.

        Returns:
            T.tensor: :math:`-\log p(x)`.

        '''
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_neg_log_prob(x, p, sum_probs=sum_probs)

    def entropy(self, p=None):
        '''Entropy function.

        Args:
            p (T.tensor): probability.

        Returns:
            T.tensor: entropy.

        '''
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_entropy(p)

    def get_center(self, p):
        '''Center of the distribution.

        Args:
            p (T.tensor): probability.

        Returns:
            T.tensor: center of distribution.

        '''
        return p

    def get_energy_bias(self, x, z):
        '''For use in RBMs and other energy based models.

        Args:
            x (T.tensor): input.
            z (T.tensor): distribution tensor.

        '''
        raise NotImplementedError()

    def scale_for_energy_model(self, x, *params):
        '''Scales input for energy based models.

        Args:
            x (T.tensor): input.

        '''
        return x


def make_conditional(C):
    '''Conditional distribution.

    Conditional distributions do not own their parameters, they are given,
    such as from an MLP.

    Args:
        C (Distribution).

    Returns:
        Conditional.

    '''
    class Conditional(C):
        def set_params(self): self.params = OrderedDict()
        def get_params(self): return []
    Conditional.__name__ = Conditional.__name__ + '_' + C.__name__

    return Conditional


class Binomial(Distribution):
    '''Binomial distribution.

    '''
    def __init__(self, dim, name='binomial', **kwargs):
        self.f_sample = _binomial
        self.f_neg_log_prob = _cross_entropy
        self.f_entropy = _binary_entropy
        super(Binomial, self).__init__(dim, name=name, **kwargs)

    def set_params(self):
        z = np.zeros((self.dim,)).astype(floatX)
        self.params = OrderedDict(z=z)

    def get_params(self):
        return [self.z]

    def get_prob(self, z):
        return T.nnet.sigmoid(z) * 0.9999 + 0.000005

    def split_prob(self, p):
        return p

    def __call__(self, z):
        return T.nnet.sigmoid(z) * 0.9999 + 0.000005

    def step_sample(self, epsilon, p):
        return (epsilon <= p).astype(floatX)

    def prototype_samples(self, size):
        return self.trng.uniform(size, dtype=floatX)

    def generate_latent_pair(self):
        h0 = T.zeros((self.dim,)).astype(floatX)[None, :]
        h = T.eye(self.dim).astype(floatX)
        return h0, h

    def visualize(self, p0, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        p0 = T.addbroadcast(p0, 0)
        return p - p0

    def get_energy_bias(self, x, z):
        return T.dot(x, z)


class CenteredBinomial(Binomial):
    '''Centered binomial.

    '''
    def get_prob(self, z):
        return T.nnet.sigmoid(2.0 * z) * 0.9999 + 0.000005

    def __call__(self, z):
        return T.nnet.sigmoid(2.0 * z) * 0.9999 + 0.000005

    def step_sample(self, epsilon, p):
        return (2.0 * (epsilon <= p).astype(floatX) - 1.0)

    def sample(self, n_samples, p=None):
        '''Samples from distribution.

        '''
        if p is None:
            p = self.get_prob(*self.get_params())
        if p.ndim == 1:
            size = (n_samples, p.shape[0] // self.scale)
        elif p.ndim == 2:
            size = (n_samples, p.shape[0], p.shape[1] // self.scale)
        elif p.ndim == 3:
            size = (n_samples, p.shape[0], p.shape[2], p.shape[3] // self.scale)
        elif p.ndim == 4:
            raise NotImplementedError('%d dim sampling not supported yet' % p.ndim)

        return (2.0 * self.f_sample(self.trng, p, size=size) - 1), theano.OrderedUpdates()

    def neg_log_prob(self, x, p=None, sum_probs=True):
        if p is None:
            p = self.get_prob(*self.get_params())
        x = 0.5 * (x + 1.0)
        return self.f_neg_log_prob(x, p, sum_probs=sum_probs)


class ContinuousBinomial(Binomial):
    '''Continuous binomial.

    Note:
        Doesn't sample.

    '''
    def sample(self, n_samples, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        return T.shape_padleft(p), theano.OrderedUpdates()


class Multinomial(Distribution):
    '''Multinomial distribuion.

    '''
    def __init__(self, dim, name='multinomial', **kwargs):
        self.f_sample = _sample_multinomial
        self.f_neg_log_prob = _categorical_cross_entropy
        self.f_entropy = _categorical_entropy
        super(Multinomial, self).__init__(dim, name=name, **kwargs)

    def set_params(self):
        z = np.zeros((self.dim,)).astype(floatX)
        self.params = OrderedDict(z=z)

    def get_prob(self, z):
        return _softmax(z)

    def __call__(self, z):
        return _softmax(z)


class Gaussian(Distribution):
    '''Gaussian distribution.

    '''
    has_kl = True
    is_continuous = True

    def __init__(self, dim, name='gaussian', clip=-10, **kwargs):
        self.f_sample = _normal
        self.f_neg_log_prob = _neg_normal_log_prob
        self.f_entropy = _normal_entropy
        self.clip = clip
        super(Gaussian, self).__init__(dim, name=name, scale=2, **kwargs)

    def set_params(self):
        mu = np.zeros((self.dim,)).astype(floatX)
        log_sigma = np.zeros((self.dim,)).astype(floatX)

        self.params = OrderedDict(
            mu=mu, log_sigma=log_sigma)

    def get_params(self):
        return [self.mu, self.log_sigma]

    def get_prob(self, mu, log_sigma):
        return concatenate([mu, log_sigma], axis=mu.ndim-1)

    def __call__(self, z):
        return z

    def get_center(self, p):
        mu = _slice(p, 0, p.shape[p.ndim-1] // self.scale)
        return mu

    def split_prob(self, p):
        mu        = _slice(p, 0, p.shape[p.ndim-1] // self.scale)
        log_sigma = _slice(p, 1, p.shape[p.ndim-1] // self.scale)
        return mu, log_sigma

    def step_kl_divergence(self, q, mu, log_sigma):
        mu_q = _slice(q, 0, self.dim)
        log_sigma_q = _slice(q, 1, self.dim)
        log_sigma_q = T.maximum(log_sigma_q, self.clip)
        log_sigma = T.maximum(log_sigma, self.clip)

        kl = log_sigma - log_sigma_q + 0.5 * (
            (T.exp(2 * log_sigma_q) + (mu - mu_q) ** 2) /
            T.exp(2 * log_sigma)
            - 1)
        return kl.sum(axis=kl.ndim-1)

    def kl_divergence(self, q):
        return self.step_kl_divergence(q, *self.get_params())

    def step_sample(self, epsilon, p):
        dim = p.shape[p.ndim-1] // self.scale
        mu = _slice(p, 0, dim)
        log_sigma = _slice(p, 1, dim)
        return mu + epsilon * T.exp(log_sigma)

    def prototype_samples(self, size):
        return self.trng.normal(
            avg=0, std=1.0,
            size=size,
            dtype=floatX
        )

    def step_neg_log_prob(self, x, *params):
        p = self.get_prob(*params)
        return self.f_neg_log_prob(x, p=p, clip=self.clip)

    def neg_log_prob(self, x, p=None, sum_probs=True):
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_neg_log_prob(x, p, clip=self.clip, sum_probs=sum_probs)

    def standard_prob(self, x, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        return T.exp(-self.neg_log_prob(x, p))

    def entropy(self, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_entropy(p, clip=self.clip)

    def generate_latent_pair(self):
        h0 = self.mu
        sigma = T.nlinalg.AllocDiag()(T.exp(self.log_sigma)).astype(floatX)
        h = 2 * sigma + h0[None, :]
        return h0, h

    def visualize(self, p0, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())

        outs0 = self.split_prob(p0)
        outs = self.split_prob(p)
        y0_mu, y0_logsigma = outs0
        y_mu, y_logsigma = outs
        py = (y_mu - y0_mu) / T.exp(y0_logsigma)
        return py

    def scale_for_energy_model(self, x, mu, log_sigma):
        '''Scales input for energy based models.

        '''
        return x / T.exp(2 * log_sigma)

    def get_energy_bias(self, x, mu, log_sigma):
        '''For use in RBMs and other energy based models.

        '''
        return -((x - mu) ** 2 / (2. * T.exp(log_sigma)) ** 2).sum(axis=x.ndim-1)


class Logistic(Distribution):
    '''Logistic distribution.

    :math:`p(x)=\\frac{e^{\\frac{x - \mu}{s}}}{s(1+e^{\\frac{x - \mu}{s}})^2}`

    Note:
        Not to be confused with logistic function.

    '''
    is_continuous = True

    def __init__(self, dim, name='logistic', **kwargs):
        self.f_sample = _logistic
        self.f_neg_log_prob = _neg_logistic_log_prob
        self.f_entropy = _logistic_entropy
        super(Logistic, self).__init__(dim, name=name, scale=2, **kwargs)

    def set_params(self):
        mu = np.zeros((self.dim,)).astype(floatX)
        log_s = np.zeros((self.dim,)).astype(floatX)

        self.params = OrderedDict(
            mu=mu, log_s=log_s)

    def get_params(self):
        return [self.mu, self.log_s]

    def get_prob(self, mu, log_s):
        return concatenate([mu, log_s], axis=mu.ndim-1)

    def __call__(self, z):
        return z

    def get_center(self, p):
        mu = _slice(p, 0, p.shape[p.ndim-1] // self.scale)
        return mu

    def split_prob(self, p):
        mu    = _slice(p, 0, p.shape[p.ndim-1] // self.scale)
        log_s = _slice(p, 1, p.shape[p.ndim-1] // self.scale)
        return mu, log_s

    def step_sample(self, epsilon, p):
        dim = p.shape[p.ndim-1] // self.scale
        mu = _slice(p, 0, dim)
        log_s = _slice(p, 1, dim)
        return mu + T.log(epsilon / (1 - epsilon)) * T.exp(log_s)

    def prototype_samples(self, size):
        return self.trng.uniform(size=size, dtype=floatX)

    def step_neg_log_prob(self, x, *params):
        p = self.get_prob(*params)
        return self.f_neg_log_prob(x, p=p)

    def neg_log_prob(self, x, p=None, sum_probs=True):
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_neg_log_prob(x, p, sum_probs=sum_probs)

    def standard_prob(self, x, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        return T.exp(-self.neg_log_prob(x, p))

    def entropy(self, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_entropy(p)

    def generate_latent_pair(self):
        h0 = self.mu
        s = T.nlinalg.AllocDiag()(T.exp(self.log_s)).astype(floatX)
        h = 2 * s + h0[None, :]
        return h0, h

    def visualize(self, p0, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())

        outs0 = self.split_prob(p0)
        outs = self.split_prob(p)
        y0_mu, y0_logs = outs0
        y_mu, y_logs = outs
        py = (y_mu - y0_mu) / T.exp(y0_logs)
        return py


class Laplace(Distribution):
    '''Laplace distribution.

    :math:`p(x) = \\frac{1}{2 b} e^{-\\frac{|x - \mu|}{b}}`.

    '''
    is_continuous = True

    def __init__(self, dim, name='laplace', **kwargs):
        self.f_sample = _laplace
        self.f_neg_log_prob = _neg_laplace_log_prob
        self.f_entropy = _laplace_entropy
        super(Laplace, self).__init__(dim, name=name, scale=2, **kwargs)

    def set_params(self):
        mu = np.zeros((self.dim,)).astype(floatX)
        log_b = np.zeros((self.dim,)).astype(floatX)

        self.params = OrderedDict(
            mu=mu, log_b=log_b)

    def get_params(self):
        return [self.mu, self.log_b]

    def get_prob(self, mu, log_b):
        return concatenate([mu, log_b], axis=mu.ndim-1)

    def __call__(self, z):
        return z

    def get_center(self, p):
        mu = _slice(p, 0, p.shape[p.ndim-1] // self.scale)
        return mu

    def split_prob(self, p):
        mu    = _slice(p, 0, p.shape[p.ndim-1] // self.scale)
        log_s = _slice(p, 1, p.shape[p.ndim-1] // self.scale)
        return mu, log_s

    def step_sample(self, epsilon, p):
        dim = p.shape[p.ndim-1] // self.scale
        mu = _slice(p, 0, dim)
        log_b = _slice(p, 1, dim)
        return mu + T.exp(log_b) * T.sgn(epsilon) * T.log(1.0 - 2 * abs(epsilon))

    def prototype_samples(self, size):
        return self.trng.uniform(size=size, dtype=floatX) - 0.5

    def step_neg_log_prob(self, x, *params):
        p = self.get_prob(*params)
        return self.f_neg_log_prob(x, p=p)

    def neg_log_prob(self, x, p=None, sum_probs=True):
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_neg_log_prob(x, p, sum_probs=sum_probs)

    def standard_prob(self, x, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        return T.exp(-self.neg_log_prob(x, p))

    def entropy(self, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())
        return self.f_entropy(p)

    def generate_latent_pair(self):
        h0 = self.mu
        b = T.nlinalg.AllocDiag()(T.exp(self.log_b)).astype(floatX)
        h = 2 * b + h0[None, :]
        return h0, h

    def visualize(self, p0, p=None):
        if p is None:
            p = self.get_prob(*self.get_params())

        outs0 = self.split_prob(p0)
        outs = self.split_prob(p)
        y0_mu, y0_logs = outs0
        y_mu, y_logs = outs
        py = (y_mu - y0_mu) / T.exp(y0_logs)
        return py

# Various functions for distributions.
# BERNOULLI --------------------------------------------------------------------

def _binomial(trng, p, size=None):
    if size is None:
        size = p.shape
    return trng.binomial(p=p, size=size, n=1, dtype=p.dtype)

def _centered_binomial(trng, p, size=None):
    if size is None:
        size = p.shape
    return 2 * trng.binomial(p=0.5*(p+1), size=size, n=1, dtype=p.dtype) - 1.

def _cross_entropy(x, p, sum_probs=True):
    energy = -x * T.log(p) - (1 - x) * T.log(1 - p)
    if sum_probs:
        return energy.sum(axis=energy.ndim-1)
    else:
        return energy

def _binary_entropy(p):
    entropy = -p * T.log(p) - (1 - p) * T.log(1 - p)
    return entropy.sum(axis=entropy.ndim-1)

# SOFTMAX ----------------------------------------------------------------------

def _softmax(x):
    axis = x.ndim - 1
    e_x = T.exp(x - x.max(axis=axis, keepdims=True))
    out = e_x / e_x.sum(axis=axis, keepdims=True)
    return out

def _sample_multinomial(trng, p, size=None):
    if size is None:
        size = p.shape
    return trng.multinomial(pvals=p, size=size).astype('float32')

def _categorical_cross_entropy(x, p, sum_probs=True):
    p = T.clip(p, _clip, 1.0 - _clip)
    energy = T.nnet.binary_crossentropy(p, x)
    if sum_probs:
        return energy.sum(axis=x.ndim-1)
    else:
        return energy

def _categorical_entropy(p):
    p_c = T.clip(p, _clip, 1.0 - _clip)
    entropy = T.nnet.categorical_crossentropy(p_c, p)
    return entropy

# GAUSSIAN ---------------------------------------------------------------------

def _normal(trng, p, size=None):
    dim = p.shape[p.ndim-1] // 2
    mu = _slice(p, 0, dim)
    log_sigma = _slice(p, 1, dim)

    if size is None:
        size = mu.shape
    return trng.normal(avg=mu, std=T.exp(log_sigma), size=size, dtype=floatX)

def _normal_prob(p):
    dim = p.shape[p.ndim-1] // 2
    mu = _slice(p, 0, dim)
    return mu

def _neg_normal_log_prob(x, p, clip=None, sum_probs=True):
    dim = p.shape[p.ndim-1] // 2
    mu = _slice(p, 0, dim)
    log_sigma = _slice(p, 1, dim)
    if clip is not None:
        log_sigma = T.maximum(log_sigma, clip)
    energy = 0.5 * (
        (x - mu)**2 / (T.exp(2 * log_sigma)) + 2 * log_sigma + T.log(2 * pi))
    if sum_probs:
        return energy.sum(axis=energy.ndim-1)
    else:
        return energy

def _normal_entropy(p, clip=None):
    dim = p.shape[p.ndim-1] // 2
    log_sigma = _slice(p, 1, dim)
    if clip is not None:
        log_sigma = T.maximum(log_sigma, clip)
    entropy = 0.5 * T.log(2 * pi * e) + log_sigma
    return entropy.sum(axis=entropy.ndim-1)

# LOGISTIC ---------------------------------------------------------------------

def _logistic(trng, p, size=None):
    dim = p.shape[p.ndim-1] // 2
    mu = _slice(p, 0, dim)
    log_s = _slice(p, 1, dim)
    if size is None:
        size = mu.shape
    epsilon = trng.uniform(size=size, dtype=floatX)
    return mu + T.log(epsilon / (1 - epsilon)) * T.exp(log_s)

def _neg_logistic_log_prob(x, p, sum_probs=True):
    dim = p.shape[p.ndim-1] // 2
    mu = _slice(p, 0, dim)
    log_s = _slice(p, 1, dim)
    energy = -(x - mu) / T.exp(log_s) + log_s + 2 * T.log(1 + T.exp((x - mu) / T.exp(log_s)))
    if sum_probs:
        return energy.sum(axis=energy.ndim-1)
    else:
        return energy

def _logistic_entropy(p):
    dim = p.shape[p.ndim-1] // 2
    log_s = _slice(p, 1, dim)
    entropy = log_s + 2.0
    return entropy.sum(axis=entropy.ndim-1)

# Laplace ---------------------------------------------------------------------

def _laplace(trng, p, size=None):
    dim = p.shape[p.ndim-1] // 2
    mu = _slice(p, 0, dim)
    log_b = _slice(p, 1, dim)
    if size is None:
        size = mu.shape
    epsilon = trng.uniform(size=size, dtype=floatX) - 0.5
    return mu + T.exp(log_b) * T.sgn(epsilon) * T.log(1.0 - 2 * abs(epsilon))

def _neg_laplace_log_prob(x, p, sum_probs=True):
    dim = p.shape[p.ndim-1] // 2
    mu = _slice(p, 0, dim)
    log_b = _slice(p, 1, dim)
    energy = T.log(2.0) + log_b + abs(x - mu) / T.exp(log_b)
    if sum_probs:
        return energy.sum(axis=energy.ndim-1)
    else:
        return energy

def _laplace_entropy(p):
    dim = p.shape[p.ndim-1] // 2
    log_b = _slice(p, 1, dim)
    entropy = log_b + T.log(2.) + 1.0
    return entropy.sum(axis=entropy.ndim-1)
