from typing import Callable

import torch
import torch.nn as nn

from .channel_grouping import ChannelGrouping


class Integrand(nn.Module):
    """
    Class that wraps an integrand function and meta-data necessary to use advanced MadNIS
    features like learnable multi-channel weights, grouped channels and channel weight priors.
    """

    def __init__(
        self,
        function: Callable,
        input_dim: int,
        bounds: list[list[float]] | None = None,
        channel_count: int | None = None,
        remapped_dim: int | None = None,
        has_channel_weight_prior: bool = False,
        channel_grouping: ChannelGrouping | None = None,
    ):
        """
        Args:
            function: integrand function.
                The signature depends on the other arguments:

                  - single-channel integration, ``channel_count=None``: ``x -> f``
                  - basic multi-channel integration, ``remapped_dim=None``,
                    ``has_channel_weight_prior=False``: ``(x, c) -> f``
                  - with channel weights, ``remapped_dim=None``, ``has_channel_weight_prior=True``:
                    ``(x, c) -> (f, alpha)`` (no trainable channel weights possible)
                  - with channel-dependent mapping, ``remapped_dim: int``,
                    ``has_channel_weight_prior=False``: ``(x, c) -> (f, y)``
                  - all features, ``remapped_dim: int``, ``has_channel_weight_prior=True``:
                    ``(x, c) -> (f, y, alpha)``

                with the following tensors:

                  - ``x`` is a point generated by the importance sampling, shape (n, input_dim),
                  - ``c`` is the channel index, shape (n, ),
                  - ``f`` is the integrand value, shape (n, ),
                  - ``y`` is the point after applying a channel-dependent mapping, shape
                    (n, remapped_dim)
                  - ``alpha`` is the prior channel weight, shape (n, channel_count).
            input_dim: dimension of the integration space
            bounds: List of pairs ``[lower bound, upper bound]`` of the integration interval for
                all dimensions. The integrand is rescaled so that the MadNIS training can be
                performed on the unit hypercube. If None, the unit hypercube is used as integration
                domain.
            channel_count: None in the single-channel case, specifies the number of channels
                otherwise.
            remapped_dim: If different from None, it gives the dimension of a remapped space,
                with a channel-dependent mapping computed as part of the integrand function.
            has_channel_weight_prior: If True, the integrand returns channel weights
            channel_grouping: ChannelGrouping object or None if all channels are independent
        """
        super().__init__()
        self.input_dim = input_dim
        self.remapped_dim = input_dim if remapped_dim is None else remapped_dim
        self.channel_count = channel_count
        self.has_channel_weight_prior = has_channel_weight_prior
        self.channel_grouping = channel_grouping

        if channel_count is None:
            self.function = lambda x, channels: (function(x), None, None)
        elif remapped_dim is None:
            if has_channel_weight_prior:

                def func(x, channels):
                    w, prior = function(x, channels)
                    return w, None, prior

                self.function = func
            else:
                self.function = lambda x, channels: (function(x, channels), None, None)

        elif has_channel_weight_prior:
            self.function = function
        else:

            def func(x, channels):
                w, y = function(x, channels)
                return w, y, None

            self.function = func

        if bounds is not None:
            bounds = torch.tensor(bounds)
            self.register_buffer("scale", bounds[:, 1] - bounds[:, 0])
            self.register_buffer("offset", bounds[:, 0])
            self.register_buffer("scale_det", self.scale.prod())
            old_func = self.function

            def rescaled_func(x, channels):
                w, y, prior = old_func(self.scale * x + self.offset, channels)
                return self.scale_det * w, y, prior

            self.function = rescaled_func

        self.register_buffer(
            "channel_id_map",
            (
                None
                if self.channel_grouping is None
                else torch.tensor(
                    [
                        channel.group.group_index
                        for channel in self.channel_grouping.channels
                    ]
                )
            ),
        )

    def unique_channel_count(self) -> int:
        """
        Returns the number of channels, or, if some channels are grouped together, the number of
        channel groups
        """
        if self.channel_grouping is None:
            return self.channel_count
        else:
            return len(self.channel_grouping.groups)

    def remap_channels(self, channels: torch.Tensor) -> torch.Tensor:
        """
        Remaps channel indices to the indices of their respective channel groups if a
        ``ChannelGrouping`` object was provided, otherwise returns the indices unchanged.

        Args:
            channels: channel indices, shape (n, )
        Returns:
            remapped channel indices, shape (n, )
        """
        if self.channel_grouping is None:
            return channels
        else:
            return self.channel_id_map[channels]

    def forward(
        self, x: torch.Tensor, channels: torch.Tensor | None
    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
        return self.function(x, channels)
