# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>

from typing import Tuple

import torch
from torch import Tensor, nn

from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
from sbi.neural_nets.estimators.categorical_net import CategoricalMassEstimator


class MixedDensityEstimator(ConditionalDensityEstimator):
    """Class performing Mixed Neural Density Estimation.

    This estimator combines a categorical mass estimator and a density estimator to
    model variables with mixed types (discrete and continuous). It can be used for both
    likelihood estimation (e.g., for discrete decisions and continuous reaction times
    in decision-making models) or posterior estimation (e.g., for models that have both
    discrete and continuous parameters).
    """

    def __init__(
        self,
        discrete_net: CategoricalMassEstimator,
        continuous_net: ConditionalDensityEstimator,
        input_shape: torch.Size,
        condition_shape: torch.Size,
        embedding_net: nn.Module = nn.Identity(),
        log_transform_input: bool = False,
    ):
        """Initialize class for combining density estimators for mixed neural
        density estimation.

        Args:
            discrete_net: neural net to model discrete part of the data.
            continuous_net: neural net to model the continuous data.
            input_shape: Event shape of the input at which the density is being
                evaluated (and which is also the event_shape of samples).
            condition_shape: Shape of the condition. If not provided, it will assume a
                1D input.
            embedding_net: Embedding net for the condition.
            log_transform_input: whether to transform the continous part of the data
                into logarithmic domain before training. This is helpful for bounded
                data, e.g.,for reaction times.
        """
        super(MixedDensityEstimator, self).__init__(
            net=continuous_net, input_shape=input_shape, condition_shape=condition_shape
        )

        self.discrete_net = discrete_net
        self.continuous_net = continuous_net
        self.condition_embedding = embedding_net
        self.log_transform_input = log_transform_input

    def forward(self, input: Tensor):
        raise NotImplementedError(
            """The forward method is not implemented for mixed neural density
            estimation, use '.sample(...)' to generate samples though a forward
            pass."""
        )

    def sample(
        self, sample_shape: torch.Size, condition: Tensor, track_gradients: bool = False
    ) -> Tensor:
        """Return sample from mixed data distribution.

        Samples are generated by first sampling the discrete part of the data
        and then sampling the continuous part conditioned on the discrete part.

        Args:
            sample_shape: Shape of samples to generate.
            condition: Condition of shape `(batch_dim, *event_shape_condition)`

        Returns:
            Samples of shape `(*sample_shape, batch_dim, event_dim_input)`
        """
        num_samples = torch.Size(sample_shape).numel()
        batch_dim = condition.shape[0]
        embedded_condition = self.condition_embedding(condition)

        with torch.set_grad_enabled(track_gradients):
            # Sample discrete data given parameters.
            discrete_samples = self.discrete_net.sample(
                sample_shape=sample_shape,
                condition=condition,
            )
            num_variables = self.discrete_net.net.num_variables
            discrete_samples = discrete_samples.reshape(
                num_samples * batch_dim, num_variables
            )

            # repeat the batch of embedded condition to match number of choices.
            condition_event_dim = embedded_condition.dim() - 1
            ones_for_event_dims = (1,) * condition_event_dim
            repeated_condition = embedded_condition.repeat(
                num_samples, *ones_for_event_dims
            )
            # Combine discrete and continuous data and prepare for embedding.
            combined_condition = torch.cat(
                (discrete_samples, repeated_condition), dim=-1
            )

            # Sample continuous data conditioned on parameters and discrete data.
            # Pass num_samples=1 because the choices in the condition contains
            # num_samples elements already.
            # The continuous net has an embedding that combines the discrete and
            # the (already embedded) continuous data.
            continuous_samples = self.continuous_net.sample(
                sample_shape=torch.Size([]), condition=combined_condition
            )

            # In case input was log-transformed, move them to linear space.
            if self.log_transform_input:
                continuous_samples = continuous_samples.exp()

            joined_samples = torch.cat((continuous_samples, discrete_samples), dim=1)

            # `continuous_input` is of shape `(batch_dim * numel(sample_shape))`.
            return joined_samples.reshape(*sample_shape, batch_dim, -1)

    def log_prob(self, input: Tensor, condition: Tensor) -> Tensor:
        r"""Return the log probabilities of the inputs given a condition or multiple
        i.e. batched conditions.


        Args:
            input: Inputs to evaluate the log probability on. Of shape
                `(sample_dim, batch_dim, *event_shape)`.
            condition: Conditions of shape `(batch_dim, *event_shape)`.

        Raises:
            NotImplementedError: If `input` has a sample shape > 1. The reason
                is that in MLE, discrete and continuous data are combined into one
                condition, creating a mismatch if sample shapes are present because
                the condition cannot have a sample shape.
            AssertionError: If `input_batch_dim != condition_batch_dim`.

        Returns:
            Sample-wise log probabilities, shape `(input_sample_dim, input_batch_dim)`.
        """

        assert input.dim() > 2, (
            "Input must be of shape (sample_dim, batch_dim, *event_shape)."
        )
        input_sample_dim, input_batch_dim = input.shape[:2]
        condition_batch_dim = condition.shape[0]
        combined_batch_size = input_sample_dim * input_batch_dim

        assert condition_batch_dim == input_batch_dim, (
            f"Batch shape of condition {condition_batch_dim} and input "
            f"{input_batch_dim} do not match."
        )

        num_discrete_variables = self.discrete_net.net.num_variables
        cont_input, disc_input = _separate_input(input, num_discrete_variables)
        # Embed continuous condition
        embedded_condition = self.condition_embedding(condition)
        # expand and repeat to match batch of inputs.
        embedded_condition = embedded_condition.unsqueeze(0).repeat(
            input_sample_dim, 1, 1
        )
        combined_condition = torch.cat((disc_input, embedded_condition), dim=-1)
        # reshape to match requirement that condition has no sample dim.
        combined_condition = combined_condition.reshape(combined_batch_size, -1)

        # Get log probs from discrete data. Pass with singleton sample dim.
        disc_log_prob = self.discrete_net.log_prob(
            input=disc_input, condition=condition
        )

        cont_input_reshaped = cont_input.reshape((1, combined_batch_size, -1))

        cont_log_prob = self.continuous_net.log_prob(
            # Transform to log-space if needed.
            (
                torch.log(cont_input_reshaped)
                if self.log_transform_input
                else cont_input_reshaped
            ),
            condition=combined_condition,
        )
        cont_log_prob = cont_log_prob.reshape(disc_log_prob.shape)

        # Combine into joint lp.
        log_probs_combined = disc_log_prob + cont_log_prob

        # Maybe add log abs det jacobian of RTs: log(1/x) = - log(x)
        if self.log_transform_input:
            log_probs_combined -= torch.log(cont_input).sum(-1)

        return log_probs_combined.reshape((input_sample_dim, input_batch_dim))

    def loss(self, input: Tensor, condition: Tensor, **kwargs) -> Tensor:
        r"""Return the loss for training the density estimator.

        Args:
            input: Inputs of shape `(batch_dim, *input_event_shape)`.
            condition: Conditions of shape `(batch_dim, *condition_event_shape)`.

        Returns:
            Loss of shape `(batch_dim,)`
        """
        return -self.log_prob(input.unsqueeze(0), condition)[0]


def _separate_input(
    input: Tensor, num_discrete_columns: int = 1
) -> Tuple[Tensor, Tensor]:
    """Returns the continuous and discrete part of the given input.

    Assumes the discrete data to live in the last columns of input.
    """
    return input[..., :-num_discrete_columns], input[..., -num_discrete_columns:]


def _is_discrete(input: Tensor) -> Tensor:
    """Infer discrete columns in input data."""
    return torch.tensor([torch.allclose(col, col.round()) for col in input.T])
