from typing import Any, Callable, Optional

import torch
from torchmetrics import Metric

from torchgmm.bayes.core import CovarianceType, covariance_shape


class PriorAggregator(Metric):
    """
    The prior aggregator aggregates component probabilities over batches and process.
    """

    full_state_update = False

    def __init__(
        self,
        num_components: int,
        *,
        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
    ):
        super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore

        self.responsibilities: torch.Tensor
        self.add_state("responsibilities", torch.zeros(num_components), dist_reduce_fx="sum")

    def update(self, responsibilities: torch.Tensor) -> None:
        # Responsibilities have shape [N, K]
        self.responsibilities.add_(responsibilities.sum(0))

    def compute(self) -> torch.Tensor:
        return self.responsibilities / self.responsibilities.sum()


class MeanAggregator(Metric):
    """
    The mean aggregator aggregates component means over batches and processes.
    """

    full_state_update = False

    def __init__(
        self,
        num_components: int,
        num_features: int,
        *,
        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
    ):
        super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore

        self.mean_sum: torch.Tensor
        self.add_state("mean_sum", torch.zeros(num_components, num_features), dist_reduce_fx="sum")

        self.component_weights: torch.Tensor
        self.add_state("component_weights", torch.zeros(num_components), dist_reduce_fx="sum")

    def update(self, data: torch.Tensor, responsibilities: torch.Tensor) -> None:
        # Data has shape [N, D]
        # Responsibilities have shape [N, K]
        self.mean_sum.add_(responsibilities.t().matmul(data))
        self.component_weights.add_(responsibilities.sum(0))

    def compute(self) -> torch.Tensor:
        return self.mean_sum / self.component_weights.unsqueeze(1)


class CovarianceAggregator(Metric):
    """
    The covariance aggregator aggregates component covariances over batches and processes.
    """

    full_state_update = False

    def __init__(
        self,
        num_components: int,
        num_features: int,
        covariance_type: CovarianceType,
        reg: float,
        *,
        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
    ):
        super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore

        self.num_components = num_components
        self.num_features = num_features
        self.covariance_type = covariance_type
        self.reg = reg

        self.covariance_sum: torch.Tensor
        self.add_state(
            "covariance_sum",
            torch.zeros(covariance_shape(num_components, num_features, covariance_type)),
            dist_reduce_fx="sum",
        )

        self.component_weights: torch.Tensor
        self.add_state("component_weights", torch.zeros(num_components), dist_reduce_fx="sum")

    def update(self, data: torch.Tensor, responsibilities: torch.Tensor, means: torch.Tensor) -> None:
        data_component_weights = responsibilities.sum(0)
        self.component_weights.add_(data_component_weights)

        if self.covariance_type in ("spherical", "diag"):
            # For diagonal and spherical covariances we can rely on the identity
            # Var[X] = E[X^2] - (E[X])^2.  The component means ("means" argument)
            # correspond to the running estimate of E[X] for each component, while
            # ``x_prob`` accumulates the weighted sum of X^2 for each component.
            # We therefore accumulate the *unnormalised* numerator
            #   nk * Var[X] = (∑ r_i x_i^2) - nk * μ^2
            # and perform the normalisation (= divide by nk) in ``compute``.
            x_prob = torch.matmul(responsibilities.t(), data.square())
            m_prob = data_component_weights.unsqueeze(-1) * means.square()
            covars = x_prob - m_prob
            if self.covariance_type == "diag":
                self.covariance_sum.add_(covars)
            else:  # covariance_type == "spherical"
                self.covariance_sum.add_(covars.mean(1))
        elif self.covariance_type == "tied":
            # This is taken from https://github.com/scikit-learn/scikit-learn/blob/
            # 844b4be24d20fc42cc13b957374c718956a0db39/sklearn/mixture/_gaussian_mixture.py#L183
            x_sq = data.T.matmul(data)
            mean_sq = (data_component_weights * means.T).matmul(means)
            self.covariance_sum.add_(x_sq - mean_sq)
        else:  # covariance_type == "full":
            # We iterate over each component since this is typically faster...
            for i in range(self.num_components):
                component_diff = data - means[i]
                covars = (responsibilities[:, i].unsqueeze(1) * component_diff).T.matmul(component_diff)
                self.covariance_sum[i].add_(covars)

    def compute(self) -> torch.Tensor:
        if self.covariance_type == "diag":
            return self.covariance_sum / self.component_weights.unsqueeze(-1) + self.reg
        if self.covariance_type == "spherical":
            return self.covariance_sum / self.component_weights + self.reg
        if self.covariance_type == "tied":
            result = self.covariance_sum / self.component_weights.sum()
            shape = result.size()
            result = result.flatten()
            result[:: self.num_features + 1].add_(self.reg)
            return result.view(shape)
        # covariance_type == "full"
        result = self.covariance_sum / self.component_weights.unsqueeze(-1).unsqueeze(-1)
        diag_mask = (
            torch.eye(self.num_features, device=result.device, dtype=result.dtype)
            .bool()
            .unsqueeze(0)
            .expand(self.num_components, -1, -1)
        )
        result[diag_mask] += self.reg
        return result
