import torch
import torch.nn as nn

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


class ProposalUniformTruncatedNormalMixture(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, activation_last=None,
                                        hidden_dim=hidden_dim)
        self._total_train_iterations = 0

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

        # extract variable shape from prior (which also includes potential batching)
        variable_shape = prior_distribution.low.shape
        x = self._ff(x)

        prior_lows = prior_distribution.low
        prior_highs = prior_distribution.high

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

        prior_range = (prior_highs.repeat(1, self._mixture_components) - prior_lows.repeat(1, self._mixture_components))

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

        means = prior_lows + (means * prior_range)
        stddevs = (prior_range / 1000) + (stddevs * prior_range * 10)

        distributions = [TruncatedNormal(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),
                                         low=prior_distribution.low,
                                         high=prior_distribution.high)
                         for i in range(self._mixture_components)]

        return Mixture(distributions, coeffs)
