import torch
import torch.nn as nn

from . import EmbeddingFeedForward
from .. import util
from ..distributions import Normal, Mixture


class ProposalNormalNormalMixture(nn.Module):
    def __init__(self, input_shape, output_shape, num_layers=2, hidden_dim=None, mixture_components=10):
        super().__init__()
        # Currently only supports event_shape=torch.Size([]) for the mixture components
        self._mixture_components = mixture_components
        input_shape = util.to_size(input_shape)
        self.output_dim = util.prod(output_shape)
        self._ff = EmbeddingFeedForward(input_shape=input_shape,
                                        output_shape=torch.Size([(2*self.output_dim + 1) * self._mixture_components]),
                                        num_layers=num_layers, activation=torch.relu, hidden_dim=hidden_dim,
                                        activation_last=None)
        self._total_train_iterations = 0
        # initialise standard deviations at 0 for stability
        self._ff._layers[1].weight[mixture_components:2*mixture_components, :].data *= 0

    def forward(self, x, prior_distribution):
        """ Proposal forward function

        """
        variable_shape = prior_distribution.mean.shape
        x = self._ff(x)

        slice_size = self.output_dim*self._mixture_components
        means = x[..., :slice_size]
        stddevs = x[..., slice_size:2*slice_size]
        coeffs = x[..., 2*slice_size:]
        stddevs = torch.exp(stddevs)
        coeffs = torch.softmax(coeffs, dim=1)

        prior_means = prior_distribution.loc
        prior_stddevs = prior_distribution.scale

        if prior_means.dim() < 2:
            while prior_means.dim() < 2:
                prior_means = prior_means.unsqueeze(-1)
        else:
            prior_means = prior_means.flatten(start_dim=1)
        if prior_stddevs.dim() < 2:
            while prior_stddevs.dim() < 2:
                prior_stddevs = prior_stddevs.unsqueeze(-1)
        else:
            prior_stddevs = prior_stddevs.flatten(start_dim=1)

        prior_means = prior_means.repeat(1, self._mixture_components)
        prior_stddevs = prior_stddevs.repeat(1, self._mixture_components)

        means = prior_means + (means * prior_stddevs)
        stddevs = stddevs * prior_stddevs

        distributions = [Normal(means[..., i*self.output_dim:(i+1)*self.output_dim].view(variable_shape),
                                stddevs[..., i*self.output_dim:(i+1)*self.output_dim].view(variable_shape))
                         for i in range(self._mixture_components)]

        return Mixture(distributions, coeffs)
