# pylint: disable=too-many-arguments, too-many-locals
""" Variational inference """
import math
import functools
from collections import namedtuple
from collections import OrderedDict
from scipy.special import gammaln
import numpy as np
import torch.nn.functional as F
import torch.nn as nn
from torch.autograd import Variable
from torch.nn.parameter import Parameter

VariationalParameter = namedtuple('VariationalParameter',
                                  ['mean', 'rho', 'eps', 'netvi'])


class NetIVI(nn.Module):
    def __init__(self, dims):
        super(NetIVI, self).__init__()
        self.dim1 = dims[0]
        if len(dims)==1:
            self.dim2 = 1
        else:
            self.dim2 = dims[1]
        dims = self.dim1 * self.dim2
        self.fc1 = nn.Linear(dims, dims)

    def forward(self, x):
        x = self.fc1(x.view(-1))
        return x


def evaluate(variational_parameter):
    """ Evaluates the current value of a variational parameter.

    Returns mean + log(1 + e^rho) * eps
    :args variational_parameter: the variational parameter
    :returns: the value of the variational parameter
    """
    assert isinstance(variational_parameter, VariationalParameter), \
        "Incorrect type."
    # return variational_parameter.mean + \
    #     (1 + variational_parameter.rho.exp()).log() * variational_parameter.eps
    out = variational_parameter.netvi(variational_parameter.eps)
    if variational_parameter.netvi.dim2 != 1:
        out = out.view(variational_parameter.netvi.dim1, variational_parameter.netvi.dim2)
    return out


def rebuild_parameters(dico, module, epsilon_setting):
    """ Rebuild parameters.

    Build the computational graph corresponding to
    the computations of the parameters of the given module,
    using the corresponding variational parameters in dico,
    and the rule used to sample epsilons. If the module has
    submodules, corresponding subcomputational graphs are also
    built.

    Typically, if a module has a parameter weight, weight should
    appear in dico, and the parameter will be rebuilt as
    module.weight = dico['weight'].mean + (1+dico['weight'].rho.exp()).log() *
            dico['weight'].eps

    :args dico: a 'tree' dictionnary that contains variational
    parameters for the current module, and subtrees for submodules
    :args module: the module whose parameters are to be rebuilt
    :args epsilon_settings: how epsilons ought to be drawn

    """

    for name, p in dico.items():
        if isinstance(p, VariationalParameter):
            if p.eps is None:
                dico[name] = p._replace(eps=Variable(p.mean.data.clone()))
            epsilon_setting(name, dico[name])
            setattr(module, name, evaluate(dico[name]))
        elif p is None:
            setattr(module, name, None)
        else:
            rebuild_parameters(p, getattr(module, name), epsilon_setting)


def prior_std(p):
    """ Compute a reasonable prior standard deviation for parameter p.

    :args p: the parameter

    :return: the resulting std

    """
    stdv = 1
    if p.dim() > 1:
        for i in range(p.dim() - 1):
            stdv = stdv * p.size()[i + 1]
        stdv = 1 / math.sqrt(stdv)
    else:
        stdv = 1e-2
    return stdv


def sub_prior_loss(dico):
    """ Compute the KL divergence between prior and parameters for
    all Variational Parameters in the tree dictionary dico.

    :args dico: tree dictionary
    :return: KL divergence between prior and current
    """
    loss = 0
    for p in dico.values():
        if isinstance(p, VariationalParameter):
            mean = p.mean
            std = (1 + p.rho.exp()).log()
            std_prior = prior_std(mean)
            loss += (-(std / std_prior).log() +
                     (std.pow(2) + mean.pow(2)) /
                     (2 * std_prior ** 2) - 1 / 2).sum()
        else:
            loss += sub_prior_loss(p)
    return loss


def sub_entropy(dico):
    """ Compute the entropy of the parameters for all Variational
    Parameters in the tree dictionary dico.
    :args dico: tree dictionary
    :returns: Entropy of the current distribution
    """
    entropy = 0.
    for _, p in dico.items():
        if isinstance(p, VariationalParameter):
            std = (1 + p.rho.exp()).log()
            n = np.prod(std.size())
            entropy += std.log().sum() + .5 * n * (1 + np.log(2 * np.pi))
        else:
            entropy += sub_entropy(p)
    return entropy


def sub_conjprior(dico, alpha_0, beta_0, mu_0, kappa_0):
    """ Compute an estimation of the KL divergence between the conjugate
    prior and parameters for all Variational Parameters in the tree
    dictionary dico.

    :args dico: tree dictionary
    :args alpha_0: hyperparameter of the conjugate prior
    :args beta_0: hyperparameter of the conjugate prior
    :args mu_0: hyperparameter of the conjugate prior
    :args kappa_0: hyperparameter of the conjugate prior
    :return: estimation of the KL divergence between prior and current
    """
    logprior = 0.
    for _, p in dico.items():
        if isinstance(p, VariationalParameter):
            theta = evaluate(p)
            S = (theta.mean() - mu_0).norm() ** 2
            V = (theta - theta.mean()).norm() ** 2
            n = np.prod(theta.size())
            alpha_n = alpha_0 + n / 2
            kappa_n = kappa_0 + n
            beta_n = beta_0 + V / 2 + S * (kappa_0 * n) / (2 * kappa_n)
            logprior += - beta_n.log() * alpha_n + alpha_0 * np.log(beta_0) + \
                gammaln(alpha_n) - gammaln(alpha_0) + \
                .5 * np.log(kappa_0 / kappa_n) - .5 * n * np.log(2 * np.pi)

        else:
            logprior += sub_conjprior(
                p, alpha_0, beta_0, mu_0, kappa_0)
    return logprior


def sub_conjpriorknownmean(dico, mean, alpha_0, beta_0):
    """ Compute an estimation of the KL divergence between the conjugate
    prior when the mean is known and parameters for all Variational
    Parameters in the tree dictionary dico.

    :args dico: tree dictionary
    :args mean: known mean for the conjugate prior
    :args alpha_0: hyperparameter of the conjugate prior
    :args beta_0: hyperparameter of the conjugate prior
    :return: estimation of the KL divergence between prior and current
    """
    logprior = 0.
    for _, p in dico.items():
        if isinstance(p, VariationalParameter):
            theta = evaluate(p)
            S = (theta - mean).norm() ** 2
            n = np.prod(theta.size())
            alpha_n = alpha_0 + n / 2
            beta_n = beta_0 + S / 2
            logprior += - beta_n.log() * alpha_n + \
                gammaln(alpha_n) - gammaln(alpha_0) + \
                alpha_0 * np.log(beta_0) - .5 * n * np.log(2 * np.pi)
        else:
            logprior += sub_conjpriorknownmean(
                p, mean, alpha_0, beta_0)
    return logprior


def sub_mixtgaussprior(dico, sigma_1, sigma_2, pi):
    """ Compute an estimation of the KL divergence between the prior
    defined by the mixture of two gaussian distributions
    for all Variational Parameters in the tree dictionary dico.
    More details on this prior and the notations can be found in :
    "Weight Uncertainty in Neural Networks" Blundell et al, 2015
    https://arxiv.org/pdf/1505.05424.pdf

    :args dico: tree dictionary
    :args sigma_1: std of the first gaussian in the mixture
    :args sigma_2: std of the second gaussian in the mixture
    :args pi: probability of the first gaussian in the mixture
    :return: estimation of the KL divergence between prior and current
    """
    logprior = 0.
    for _, p in dico.items():
        if isinstance(p, VariationalParameter):
            theta = evaluate(p)
            n = np.prod(theta.size())
            theta2 = theta ** 2
            pgauss1 = (- theta2 / (2. * sigma_1 ** 2)).exp() / sigma_1
            pgauss2 = (- theta2 / (2. * sigma_2 ** 2)).exp() / sigma_2
            logprior += (pi * pgauss1 + (1 - pi) * pgauss2 + 1e-8).log().sum()
            logprior -= n / 2 * np.log(2 * np.pi)
        else:
            logprior += sub_mixtgaussprior(
                p, sigma_1, sigma_2, pi)
    return logprior


class IVariationalize(nn.Module):
    """ Build a Implicit Variational model over the model given as input.

    Variationalize changes all parameters of the given model
    to allow learning of a gaussian distribution over the
    parameters using Variational inference. For more information,
    see e.g. https://papers.nips.cc/paper/4329-practical-variational
    -inference-for-neural-networks.pdf.

    :args model: the model on which VI is to be performed
    :args zero_mean: if True, sets initial mean to 0, else
        keep model initial mean
    :args learn_mean: if True, learn the posterior mean
    :args learn_rho: if True, learn the posterior rho
    """
    def __init__(self, model, zero_mean=True, learn_mean=True, learn_rho=True):
        super().__init__()
        self.model = model
        self.model_ivi = nn.ModuleList()
        self.dico = OrderedDict()
        self._variationalize_module(self.dico, self.model, '', zero_mean,
                                    learn_mean, learn_rho)
        self._prior_loss_function = functools.partial(
            sub_prior_loss,
            dico=self.dico)

    def _variationalize_module(self, dico, module, prefix, zero_mean,
                               learn_mean, learn_rho):
        to_erase = []
        paras = module._parameters.items()  # pylint: disable=protected-access
        for name, p in paras:
            if p is None:
                dico[name] = None
            else:
                stdv = prior_std(p)
                init_rho = math.log(math.exp(stdv) - 1)

                init_mean = p.data.clone()
                if zero_mean:
                    init_mean.fill_(0)

                dico[name] = VariationalParameter(
                    Parameter(init_mean),
                    Parameter(p.data.clone().fill_(init_rho)),
                    None, NetIVI(p.shape))

                self.model_ivi.append(dico[name].netvi)

                # if learn_mean:
                #     self.register_parameter(prefix + '_' + name + '_mean',
                #                             dico[name].mean)
                # if learn_rho:
                #     self.register_parameter(prefix + '_' + name + '_rho',
                #                             dico[name].rho)

            to_erase.append(name)

        for name in to_erase:
            delattr(module, name)

        for mname, sub_module in module.named_children():
            sub_dico = OrderedDict()
            self._variationalize_module(sub_dico, sub_module,
                                        prefix + ('_' if prefix else '') +
                                        mname, zero_mean,
                                        learn_mean, learn_rho)
            dico[mname] = sub_dico

    def set_prior(self, prior_type, **prior_parameters):
        """ Change the prior to be used.

        Available priors are 'gaussian', 'conjugate', 'mixtgauss' and
        'conjugate_known_mean'. For each prior, you must
        specify the corresponding parameter:
          - For the gaussian prior, no parameter is required.
          - For the conjugate prior, you must specify
            - n_mc_samples, the number of samples used in the Monte Carlo
              estimation of the prior loss and its gradient.
            - mu_0, the prior sample mean
            - kappa_0, the number of samples used to estimate the
              prior sample mean
            - alpha_0 and beta_0, such that variance was estimated from 2
             alpha_0 observations with sample mean mu_0 and sum of squared
             deviations 2 beta_0
          - For the conjugate prior with known mean,
            - n_mc_samples, the number of samples used in the Monte Carlo
              estimation of the prior loss and its gradient.
            - mean, the known mean
            - alpha_0 and beta_0 defined as above
          - For the mixture of two gaussians,
            - n_mc_samples, the number of samples used in the Monte Carlo
              estimation of the prior loss and its gradient.
            - sigma_1 and sigma_2 the std of the two gaussians
            - pi the probability of the first gaussian
        For further information, see:
            https://en.wikipedia.org/wiki/Conjugate_prior.
        Acts inplace by modifying the value of _prior_loss_function
        :args prior_type: one of 'gaussian', 'conjugate',
            'conjugate_known_mean', 'mixtgauss'
        :args prior_parameters: the parameters for the associated prior
        """
        if prior_type == 'gaussian':
            self._prior_loss_function = functools.partial(
                sub_prior_loss,
                dico=self.dico)
        else:
            n_mc_samples = prior_parameters.pop("n_mc_samples")
            if prior_type == 'conjugate':
                mc_logprior_function = functools.partial(
                    sub_conjprior,
                    **prior_parameters
                )
            if prior_type == 'conjugate_known_mean':
                mc_logprior_function = functools.partial(
                    sub_conjpriorknownmean,
                    **prior_parameters
                )
            if prior_type == 'mixtgauss':
                mc_logprior_function = functools.partial(
                    sub_mixtgaussprior,
                    **prior_parameters
                )

            def prior_loss_function():
                """Compute the prior loss"""
                logprior = 0.
                for _ in range(n_mc_samples):
                    rebuild_parameters(
                        self.dico, self.model,
                        lambda name, p: p.eps.data.normal_()
                    )
                    logprior += mc_logprior_function(self.dico)
                logprior = logprior / n_mc_samples
                H = sub_entropy(self.dico)
                prior_loss = - logprior - H
                return prior_loss
            self._prior_loss_function = prior_loss_function

    def forward(self, *inputs):
        def _epsilon_setting(name, p):  # pylint: disable=unused-argument
            if self.training:
                return p.eps.data.normal_()
            return p.eps.data.zero_()

        rebuild_parameters(self.dico, self.model, _epsilon_setting)
        return self.model(*inputs)

    def prior_loss(self):
        """ Returns the prior loss """
        return self._prior_loss_function()


class ISample(nn.Module):
    """ Utility to sample a single model from a Variational Model.

    Sample is a decorator that wraps a variational model, sample
    a model from the current parameter distribution and make the
    model usable as any other pytorch model. The sample can be
    redrawn using the draw() method. Draw needs to be called
    once before the model can be used.

    :args var_model: Variational model from which the sample models
    are to be drawn
    """
    def __init__(self, var_model):
        super().__init__()
        self.var_model = var_model

        self.association = []

    def draw(self, association=None, var_dico=None):
        """ Draw a single model from the posterior variationally learned """
        if association is None:
            self.association = []
            association = self.association
            var_dico = self.var_model.dico

        for name, p in var_dico.items():
            if isinstance(p, VariationalParameter):
                if p.eps is None:
                    var_dico[name] = p._replace(eps=Variable(
                        p.mean.data.clone()))
                association.append((var_dico[name].eps,
                                    var_dico[name].eps.data.clone().normal_()))
            else:
                self.draw(association, p)

    def forward(self, *inputs):
        for p, drawn_value_p in self.association:
            p.data.copy_(drawn_value_p)

        def _epsilon_setting(name, p):  # pylint: disable=unused-argument
            return 1
        rebuild_parameters(self.var_model.dico, self.var_model.model,
                           _epsilon_setting)
        return self.var_model.model(*inputs)
