# 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/>

import warnings
from typing import Callable, List, Optional, Tuple, Union

import torch
from torch import Tensor
from torch.distributions import Distribution

from sbi.inference.potentials.base_potential import BasePotential
from sbi.neural_nets.estimators import (
    ConditionalDensityEstimator,
    MixedDensityEstimator,
)
from sbi.neural_nets.estimators.shape_handling import (
    reshape_to_batch_event,
    reshape_to_sample_batch_event,
)
from sbi.sbi_types import TorchTransform
from sbi.utils.sbiutils import mcmc_transform
from sbi.utils.torchutils import BoxUniform
from sbi.utils.user_input_checks import process_x
from sbi.utils.user_input_checks_utils import (
    MultipleIndependent,
    move_distribution_to_device,
)


def likelihood_estimator_based_potential(
    likelihood_estimator: ConditionalDensityEstimator,
    prior: Distribution,  # type: ignore
    x_o: Optional[Tensor],
    enable_transform: bool = True,
) -> Tuple["LikelihoodBasedPotential", TorchTransform]:
    r"""Returns potential :math:`\log(p(x_o|\theta)p(\theta))` for likelihood estimator.

    It also returns a transformation that can be used to transform the potential into
    unconstrained space.

    Args:
        likelihood_estimator: The density estimator modelling the likelihood.
        prior: The prior distribution.
        x_o: The observed data at which to evaluate the likelihood.
        enable_transform: Whether to transform parameters to unconstrained space.
             When False, an identity transform will be returned for `theta_transform`.

    Returns:
        The potential function $p(x_o|\theta)p(\theta)$ and a transformation that maps
        to unconstrained space.
    """

    device = str(next(likelihood_estimator.parameters()).device)

    potential_fn = LikelihoodBasedPotential(
        likelihood_estimator, prior, x_o, device=device
    )
    theta_transform = mcmc_transform(
        prior, device=device, enable_transform=enable_transform
    )

    return potential_fn, theta_transform


class LikelihoodBasedPotential(BasePotential):
    def __init__(
        self,
        likelihood_estimator: ConditionalDensityEstimator,
        prior: Distribution,
        x_o: Optional[Tensor] = None,
        device: str = "cpu",
    ):
        r"""Returns the potential function for likelihood-based methods.

        Args:
            likelihood_estimator: The density estimator modelling the likelihood.
            prior: The prior distribution.
            x_o: The observed data at which to evaluate the likelihood.
            device: The device to which parameters and data are moved before evaluating
                the `likelihood_nn`.

        Returns:
            The potential function $p(x_o|\theta)p(\theta)$.
        """

        super().__init__(prior, x_o, device)
        self.likelihood_estimator = likelihood_estimator
        self.likelihood_estimator.eval()

    def to(self, device: Union[str, torch.device]) -> "LikelihoodBasedPotential":
        """Move likelihood estimator, prior and x_o to the given device.

        Args:
            device: Device to move the likelihood_estimator, prior and x_o to.

        Returns:
            Self for method chaining.
        """
        super().to(device)
        self.likelihood_estimator.to(device)
        return self

    def bind(
        self, x_o: Optional[Tensor], x_is_iid: bool = True
    ) -> "LikelihoodBasedPotential":
        """Create new potential with x bound, without mutable state.

        Args:
            x_o: Observed data to bind.
            x_is_iid: Whether x_o is a batch of iid observations.

        Returns:
            A new LikelihoodBasedPotential with x_o bound.
        """
        bound = LikelihoodBasedPotential(
            likelihood_estimator=self.likelihood_estimator,
            prior=self.prior,
            x_o=None,
            device=self.device,
        )
        if x_o is not None:
            x_o = process_x(x_o).to(self.device)
        bound._x_o = x_o
        bound._x_is_iid = x_is_iid
        return bound

    def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
        r"""Returns the potential $\log(p(x_o|\theta)p(\theta))$.

        Args:
            theta: The parameter set at which to evaluate the potential function.
            track_gradients: Whether to track the gradients.

        Returns:
            The potential $\log(p(x_o|\theta)p(\theta))$.
        """
        with torch.set_grad_enabled(track_gradients):
            log_prior = self.prior.log_prob(theta)  # type: ignore
        if self.x_is_iid:
            # For each theta, calculate the likelihood sum over all x in batch.
            log_likelihood_trial_sum = _log_likelihoods_over_trials(
                x=self.x_o,
                theta=theta.to(self.device),
                estimator=self.likelihood_estimator,
                track_gradients=track_gradients,
            )
            return log_likelihood_trial_sum + log_prior
        else:
            # Calculate likelihood for each (theta,x) pair separately
            theta_batch_size = theta.shape[0]
            x_batch_size = self.x_o.shape[0]
            assert theta_batch_size == x_batch_size, (
                f"Batch size mismatch: {theta_batch_size} and {x_batch_size}.\
                When performing batched sampling for multiple `x`, the batch size of\
                `theta` must match the batch size of `x`."
            )
            x = self.x_o.unsqueeze(0)
            with torch.set_grad_enabled(track_gradients):
                log_likelihood_batches = self.likelihood_estimator.log_prob(
                    x, condition=theta
                )
            return log_likelihood_batches + log_prior

    def condition_on_theta(
        self,
        local_theta: Tensor,
        dims_global_theta: List[int],
        prior: Optional[Distribution] = None,
    ) -> Callable:
        r"""Returns a potential function conditioned on a subset of theta dimensions.

        The goal of this function is to divide the original `theta` into a
        `global_theta` $\theta_g$ we do inference over, and a `local_theta`
        $\theta_{l,i}$ per trial that we condition on (in addition to conditioning on
        `x_o`). Thus, the returned potential function will calculate
        $\log p(\theta_g) + \sum_{i=1}^{N} \log p(x_i \mid \theta_g, \theta_{l,i})$,
        where `x_i` and $\theta_{l,i}$ are fixed and $\theta_g$ varies at inference
        time.

        Args:
            local_theta: The condition values to be conditioned.
            dims_global_theta: The indices of the columns in `theta` that will be
                sampled, i.e., that *not* conditioned. For example, if original theta
                has shape `(batch_dim, 3)`, and `dims_global_theta=[0, 1]`, then the
                potential will set `theta[:, 2] = local_theta` at inference time.
            prior: Prior over the `global_theta`. If None, the marginal of the prior
                of this potential over `dims_global_theta` is used. This requires a
                `BoxUniform` prior, or a `MultipleIndependent` prior whose components
                cover exactly these dimensions, given in increasing order. The prior
                is moved to the device of this potential.

        Returns:
            A potential function conditioned on the `local_theta`. It returns the log
            likelihood summed over the trials plus the log prior of `global_theta`.

        Raises:
            ValueError: If `prior` is None and the prior of this potential has no
                marginal over `dims_global_theta`.
        """

        assert self.x_is_iid, "Conditioning is only supported for iid data."
        global_prior = (
            _marginal_prior(self.prior, dims_global_theta) if prior is None else prior
        )
        global_prior_device = None

        def conditioned_potential(
            theta: Tensor, x_o: Optional[Tensor] = None, track_gradients: bool = True
        ) -> Tensor:
            nonlocal global_prior, global_prior_device
            assert len(dims_global_theta) == theta.shape[-1], (
                "dims_global_theta must match the number of parameters to sample."
            )
            if theta.dim() > 2:
                assert theta.shape[0] == 1, (
                    "condition_on_theta does not support sample shape for theta."
                )
                theta = theta.squeeze(0)
            global_theta = theta[:, dims_global_theta].to(self.device)
            x_o = x_o if x_o is not None else self.x_o
            # x needs shape (sample_dim (iid), batch_dim (xs), *event_shape)
            if x_o.dim() < 3:
                x_o = reshape_to_sample_batch_event(
                    x_o, event_shape=x_o.shape[1:], leading_is_sample=self.x_is_iid
                )

            log_likelihood = _log_likelihood_over_iid_trials_and_local_theta(
                x=x_o.to(self.device),
                global_theta=global_theta,
                local_theta=local_theta.to(self.device),
                estimator=self.likelihood_estimator,
                track_gradients=track_gradients,
            )
            # Follow the device of this potential, also after `.to()`.
            if global_prior_device != self.device:
                global_prior = move_distribution_to_device(global_prior, self.device)
                global_prior_device = self.device
            # Priors with batch_shape [1] return one log prob per dimension.
            with torch.set_grad_enabled(track_gradients):
                log_prior = global_prior.log_prob(global_theta)
            return log_likelihood + log_prior.reshape(len(global_theta), -1).sum(-1)

        return conditioned_potential


def _marginal_prior(prior: Distribution, dims: List[int]) -> Distribution:
    """Return the marginal of a `BoxUniform` or `MultipleIndependent` prior over `dims`.

    Args:
        prior: Prior over all dimensions of `theta`.
        dims: Dimensions of the marginal, in increasing order.

    Raises:
        ValueError: If `prior` is neither a `BoxUniform` nor a `MultipleIndependent`
            whose components cover exactly `dims`.
    """
    if isinstance(prior, BoxUniform) and prior.base_dist.low.dim() == 1:
        # `prior.low` and `prior.high` keep their input shapes; the base
        # distribution holds the broadcast bounds.
        low, high = prior.base_dist.low, prior.base_dist.high
        return BoxUniform(low[dims], high[dims], device=prior.device)
    if isinstance(prior, MultipleIndependent):
        components, covered_dims, start = [], [], 0
        for dist, num_dims in zip(prior.dists, prior.dims_per_dist, strict=True):
            component_dims = list(range(start, start + num_dims))
            start += num_dims
            if set(component_dims) <= set(dims):
                components.append(dist)
                covered_dims += component_dims
        if covered_dims == list(dims):
            if len(components) > 1:
                return MultipleIndependent(components, device=prior.device)
            return components[0]
    raise ValueError(
        f"condition_on_theta needs the prior of the parameters {list(dims)}. It "
        "derives it only from a BoxUniform prior, or a MultipleIndependent prior "
        "whose components cover exactly these dimensions in increasing order. This "
        f"potential has a {type(prior).__name__} prior. Pass `prior` with the prior "
        "of these parameters."
    )


def _log_likelihoods_over_trials(
    x: Tensor,
    theta: Tensor,
    estimator: ConditionalDensityEstimator,
    track_gradients: bool = False,
) -> Tensor:
    r"""Return log likelihoods summed over iid trials of `x`.

    Note: `x` can be a batch with batch size larger 1. Batches in `x` are assumed
    to be iid trials, i.e., data generated based on the same paramters /
    experimental conditions.

    Repeats `x` and $\theta$ to cover all their combinations of batch entries.

    Args:
        x: Batch of iid data of shape `(iid_dim, *event_shape)`.
        theta: Batch of parameters of shape `(batch_dim, *event_shape)`.
        estimator: DensityEstimator.
        track_gradients: Whether to track gradients.

    Returns:
        log_likelihood_trial_sum: log likelihood for each parameter, summed over all
            batch entries (iid trials) in `x`.
    """
    # Shape of `x` is (iid_dim, *event_shape).
    x = reshape_to_sample_batch_event(
        x, event_shape=x.shape[1:], leading_is_sample=True
    )

    # Match the number of `x` to the number of conditions (`theta`). This is important
    # if the potential is simulataneously evaluated at multiple `theta` (e.g.
    # multi-chain MCMC).
    theta_batch_size = theta.shape[0]
    trailing_minus_ones = [-1 for _ in range(x.dim() - 2)]
    x = x.expand(-1, theta_batch_size, *trailing_minus_ones)

    assert (
        next(estimator.parameters()).device == x.device and x.device == theta.device
    ), f"""device mismatch: estimator, x, theta: \
        {next(estimator.parameters()).device}, {x.device},
        {theta.device}."""

    # Shape of `theta` is (batch_dim, *event_shape). Therefore, the call below should
    # not change anything, and we just have it as "best practice" before calling
    # `DensityEstimator.log_prob`.
    theta = reshape_to_batch_event(theta, event_shape=theta.shape[1:])

    # Calculate likelihood in one batch.
    with torch.set_grad_enabled(track_gradients):
        log_likelihood_trial_batch = estimator.log_prob(x, condition=theta)
        # Sum over trial-log likelihoods.
        log_likelihood_trial_sum = log_likelihood_trial_batch.sum(0)

    return log_likelihood_trial_sum


def _log_likelihood_over_iid_trials_and_local_theta(
    x: Tensor,
    global_theta: Tensor,
    local_theta: Tensor,
    estimator: ConditionalDensityEstimator,
    track_gradients: bool = False,
) -> Tensor:
    """Returns $\\prod_{i=1}^N \\log(p(x_i|\theta, local_theta_i)$.

    `x` is a batch of iid data, and `local_theta` is a matching batch of condition
    values that were part of `theta` but are treated as local iid variables at inference
    time.

    This function is different from `_log_likelihoods_over_trials` in that it moves the
    iid batch dimension of `x` onto the batch dimension of `theta`. This is needed when
    the likelihood estimator is conditioned on a batch of conditions that are iid with
    the batch of `x`. It avoids the evaluation of the likelihood for every combination
    of `x` and `local_theta`.

    Args:
        x: data with shape `(sample_dim, x_batch_dim, *x_event_shape)`, where sample_dim
            holds the i.i.d. trials and batch_dim holds a batch of xs, e.g., non-iid
            observations.
        global_theta: Batch of parameters `(theta_batch_dim,
            num_parameters)`.
        local_theta: Batch of conditions of shape `(sample_dim, num_local_thetas)`, must
            match x's `sample_dim`.
        estimator: DensityEstimator.
        track_gradients: Whether to track gradients.

    Returns:
        log_likelihood: log likelihood for each x in x_batch_dim, for each theta in
            theta_batch_dim, summed over all iid trials. Shape `(x_batch_dim,
            theta_batch_dim)`.
    """
    assert x.dim() > 2, "x must have shape (sample_dim, batch_dim, *event_shape)."
    assert local_theta.dim() == 2, (
        "condition must have shape (sample_dim, num_conditions)."
    )
    assert global_theta.dim() == 2, "theta must have shape (batch_dim, num_parameters)."
    num_trials, num_xs = x.shape[:2]
    num_thetas = global_theta.shape[0]
    assert local_theta.shape[0] == num_trials, (
        "Condition batch size must match the number of iid trials in x."
    )
    if num_xs > 1:
        raise NotImplementedError(
            "Batched sampling for multiple `x` is not supported for iid conditions."
        )

    # move the iid batch dimension onto the batch dimension of theta and repeat it there
    x_repeated = torch.transpose(x, 0, 1).repeat_interleave(num_thetas, dim=1)

    # construct theta and condition to cover all trial-theta combinations
    theta_with_condition = torch.cat(
        [
            global_theta.repeat(num_trials, 1),  # repeat ABAB
            local_theta.repeat_interleave(num_thetas, dim=0),  # repeat AABB
        ],
        dim=-1,
    )

    with torch.set_grad_enabled(track_gradients):
        # Calculate likelihood in one batch. Returns (1, num_trials * num_theta)
        log_likelihood_trial_batch = estimator.log_prob(
            x_repeated, condition=theta_with_condition
        )
        # Reshape to (x-trials x parameters), sum over trial-log likelihoods.
        log_likelihood_trial_sum = log_likelihood_trial_batch.reshape(
            num_xs, num_trials, num_thetas
        ).sum(1)

    # remove xs batch dimension
    return log_likelihood_trial_sum.squeeze(0)


def mixed_likelihood_estimator_based_potential(
    likelihood_estimator: MixedDensityEstimator,
    prior: Distribution,
    x_o: Optional[Tensor],
) -> Tuple[Callable, TorchTransform]:
    r"""Returns $\log(p(x_o|\theta)p(\theta))$ for mixed likelihood-based methods.

    It also returns a transformation that can be used to transform the potential into
    unconstrained space.

    Args:
        likelihood_estimator: The neural network modelling the likelihood.
        prior: The prior distribution.
        x_o: The observed data at which to evaluate the likelihood.

    Returns:
        The potential function $p(x_o|\theta)p(\theta)$ and a transformation that maps
        to unconstrained space.
    """

    warnings.warn(
        "This function is deprecated and will be removed in a future release. Use "
        "`likelihood_estimator_based_potential` instead.",
        DeprecationWarning,
        stacklevel=2,
    )

    device = str(next(likelihood_estimator.discrete_net.parameters()).device)

    potential_fn = MixedLikelihoodBasedPotential(
        likelihood_estimator, prior, x_o, device=device
    )
    theta_transform = mcmc_transform(prior, device=device)

    return potential_fn, theta_transform


class MixedLikelihoodBasedPotential(LikelihoodBasedPotential):
    def __init__(
        self,
        likelihood_estimator: MixedDensityEstimator,
        prior: Distribution,
        x_o: Optional[Tensor],
        device: str = "cpu",
    ):
        super().__init__(likelihood_estimator, prior, x_o, device)

        warnings.warn(
            "This function is deprecated and will be removed in a future release. Use "
            "`LikelihoodBasedPotential` instead.",
            DeprecationWarning,
            stacklevel=2,
        )

    def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
        with torch.set_grad_enabled(track_gradients):
            prior_log_prob = self.prior.log_prob(theta)  # type: ignore

        # Shape of `x` is (iid_dim, *event_shape)
        theta = reshape_to_batch_event(theta, event_shape=theta.shape[1:])
        x = reshape_to_sample_batch_event(
            self.x_o, event_shape=self.x_o.shape[1:], leading_is_sample=True
        )
        theta_batch_dim = theta.shape[0]
        # Match the number of `x` to the number of conditions (`theta`). This is
        # importantif the potential is simulataneously evaluated at multiple `theta`
        # (e.g. multi-chain MCMC).
        trailing_minus_ones = [-1 for _ in range(x.dim() - 2)]
        x = x.expand(-1, theta_batch_dim, *trailing_minus_ones)

        # Calculate likelihood in one batch.
        with torch.set_grad_enabled(track_gradients):
            # Call the specific log prob method of the mixed likelihood estimator as
            # this optimizes the evaluation of the discrete data part.
            log_likelihood_trial_batch = self.likelihood_estimator.log_prob(
                input=x,
                condition=theta.to(self.device),
            )
            # Reshape to (x-trials x parameters), sum over trial-log likelihoods.
            log_likelihood_trial_sum = log_likelihood_trial_batch.reshape(
                self.x_o.shape[0], -1
            ).sum(0)

        return log_likelihood_trial_sum + prior_log_prob