"""Monitor expert activations similarity during evaluation.

This callback computes, for each SigmaMoE (expert FFN) layer, the average
pairwise cosine similarity between the activations of all experts on the same
inputs. It does NOT log raw activation values. Metrics are collected only during
evaluation (not during training).

It uses the same SigmaMoE layers referenced by the experts selection callback
(i.e., transformer blocks' FFN layers).
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch
import torch.nn.functional as F  # noqa: N812
from composer.core import Callback, State
from llmfoundry.registry import callbacks

from mt_dao.composer_callbacks.plotting_utils import (
    HeatmapOptions,
    create_heatmap,
)
from mt_dao.models.ffns.sigma_moe import SigmaMoE

if TYPE_CHECKING:
    from collections.abc import Callable

    from composer.loggers import Logger

    import wandb

__all__ = ["ExpertSimilarityMonitor"]


@callbacks.register("expert_similarity_monitor")
class ExpertSimilarityMonitor(Callback):
    """Logs average pairwise similarity between expert activations during eval.

    For each SigmaMoE layer, computes per-expert outputs for the given inputs
    (without routing/weighting), then averages the experts' pairwise cosine
    similarity matrices across all tokens in the batch/sequence.

    Logged keys (eval only):
      - "expert_activation/block_{i}/cos_sim/mean"
      - "expert_activation/block_{i}/cos_sim/std"
      - "expert_activation/block_{i}/cos_sim/min"
      - "expert_activation/block_{i}/cos_sim/max"
      - "expert_activation/block_{i}/cos_sim/off_diag_mean"
      - "expert_activation/block_{i}/cos_sim/off_diag_std"
      - "expert_activation/block_{i}/cos_sim/diag_mean"
      - "expert_activation/block_{i}/cos_sim/heatmap" (image)
    """

    def __init__(
        self,
        *,
        log_key_prefix: str = "expert_activation",
        all_reduce_across_ranks: bool = True,
        accumulate_across_eval: bool = False,
        weighted_by_tokens: bool = True,
        first_n_eval_batches: int | None = None,
    ) -> None:
        """Initialize the activation similarity monitor.

        Parameters
        ----------
        log_key_prefix : str, optional
            Prefix for logged metric keys, by default "expert_activation".
        all_reduce_across_ranks : bool, optional
            If True, average similarity matrices across distributed ranks, by
            default True.
        accumulate_across_eval : bool, optional
            If True, accumulate similarity across all eval batches and log once at
            eval_end. If False, log per batch at eval_batch_end. Default False.
        weighted_by_tokens : bool, optional
            If accumulating across eval, weight batch contributions by number of
            tokens (B*S). Default True.
        first_n_eval_batches : int | None, optional
            If set, only collect/report metrics for the first N eval batches.
            Batches beyond N are ignored (hooks not registered). Default None.

        """
        self.log_key_prefix = log_key_prefix
        self.all_reduce_across_ranks = all_reduce_across_ranks
        self.accumulate_across_eval = accumulate_across_eval
        self.weighted_by_tokens = weighted_by_tokens
        self.first_n_eval_batches = first_n_eval_batches

        # Forward-hook handles and mapping
        self._handles: list[torch.utils.hooks.RemovableHandle] = []
        self._module_to_block_idx: dict[torch.nn.Module, int] = {}

        # Accumulated per-batch results (current batch)
        self._block_cos_mats: dict[int, torch.Tensor] = {}
        self._block_token_counts: dict[int, int] = {}

        # Accumulators across the whole eval
        self._accum_sum_cos_mats: dict[int, torch.Tensor] = {}
        self._accum_denoms: dict[int, float] = {}

        # Eval batch counter for limiting collection
        self._eval_batches_seen: int = 0

    # -------- Internals ---------
    @staticmethod
    def _compute_expert_outputs(module: SigmaMoE, x: torch.Tensor) -> torch.Tensor:
        """Vectorized per-expert outputs without routing/weighting.

        Parameters
        ----------
        module : SigmaMoE
            The SigmaMoE module.
        x : torch.Tensor
            Input tensor of shape [B, S, D].

        Returns
        -------
        torch.Tensor
            Tensor of shape [N, E, V] where N=B*S (tokens), E=n_experts, V=v_dim.

        """
        d_model = module.d_model
        x_flat = x.view(-1, d_model)  # [N, D]

        # Up-projection for all experts: [N, E, expert_size]
        up_outputs = torch.einsum("nd,esd->nes", x_flat, module.up_proj)
        up_outputs = module.activation(up_outputs)

        # Down-projection: [N, E, V]
        return torch.einsum("nes,esv->nev", up_outputs, module.down_proj)

    @staticmethod
    def _average_pairwise_cos(expert_outputs: torch.Tensor) -> torch.Tensor:
        """Compute average pairwise cosine similarity across tokens.

        Parameters
        ----------
        expert_outputs : torch.Tensor
            Tensor with shape [N, E, V].

        Returns
        -------
        torch.Tensor
            Average cosine-similarity matrix with shape [E, E].

        """
        # Normalize per token across feature dim
        exp_norm = F.normalize(expert_outputs, p=2, dim=-1)  # [N, E, V]
        # Batched E x V @ V x E -> E x E per token
        cos_mats = torch.bmm(exp_norm, exp_norm.transpose(1, 2))  # [N, E, E]
        return cos_mats.mean(dim=0)  # [E, E]

    @staticmethod
    def _summarize_cos_matrix(cos: torch.Tensor) -> dict[str, float]:
        metrics: dict[str, float] = {}
        metrics["cos_sim/mean"] = cos.mean().item()
        metrics["cos_sim/std"] = cos.std().item()
        metrics["cos_sim/min"] = cos.min().item()
        metrics["cos_sim/max"] = cos.max().item()

        e = cos.shape[0]
        if e > 1:
            off_diag_mask = ~torch.eye(e, dtype=torch.bool, device=cos.device)
            off_vals = cos[off_diag_mask]
            metrics["cos_sim/off_diag_mean"] = off_vals.mean().item()
            metrics["cos_sim/off_diag_std"] = off_vals.std().item()
        diag_vals = torch.diag(cos)
        metrics["cos_sim/diag_mean"] = diag_vals.mean().item()
        metrics["cos_sim/diag_std"] = diag_vals.std().item()
        return metrics

    @staticmethod
    def _all_reduce_mean_tensor(t: torch.Tensor) -> torch.Tensor:
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.AVG)
        return t

    @staticmethod
    def _all_reduce_sum_tensor(t: torch.Tensor) -> torch.Tensor:
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.SUM)
        return t

    def _register_hooks_for_eval(self, state: State) -> None:
        """Attach forward hooks to SigmaMoE layers for the upcoming eval pass."""
        self._handles.clear()
        self._module_to_block_idx.clear()
        self._block_cos_mats.clear()
        self._block_token_counts.clear()

        # Unwrap DDP if present and get transformer
        original_model = (
            state.model.module
            if isinstance(state.model, torch.nn.parallel.DistributedDataParallel)
            else state.model
        )
        # MPTMuPForCausalLM -> .model.transformer (as used in other callbacks)
        transformer = getattr(
            getattr(original_model, "model", original_model),
            "transformer",
            None,
        )
        if transformer is None:
            return

        for block_idx, block in enumerate(getattr(transformer, "blocks", [])):
            ffn = getattr(block, "ffn", None)
            if isinstance(ffn, SigmaMoE):
                self._module_to_block_idx[ffn] = block_idx
                handle = ffn.register_forward_hook(self._make_forward_hook(block_idx))
                self._handles.append(handle)

    def _make_forward_hook(
        self,
        block_idx: int,
    ) -> Callable[[torch.nn.Module, tuple[torch.Tensor, ...], torch.Tensor], None]:
        @torch.no_grad()  # type: ignore[misc]
        def _hook(
            module: torch.nn.Module,
            inputs: tuple[torch.Tensor, ...],
            output: torch.Tensor,  # noqa: ARG001
        ) -> None:
            if not isinstance(module, SigmaMoE):
                return
            if not inputs:
                return
            x = inputs[0]
            if not isinstance(x, torch.Tensor):
                return
            # Compute per-expert outputs and average cosine similarity matrix
            expert_outputs = self._compute_expert_outputs(module, x)
            avg_cos = self._average_pairwise_cos(expert_outputs).detach()

            # Determine number of tokens in this forward
            min_bs_dims = 3
            if x.dim() >= min_bs_dims:
                n_tokens_local = int(x.shape[0] * x.shape[1])
            else:
                n_tokens_local = int(expert_outputs.shape[0])
            n_tokens_total = n_tokens_local

            # Distributed combine
            if self.all_reduce_across_ranks:
                if self.accumulate_across_eval and self.weighted_by_tokens:
                    # Reduce SUM of per-rank sums and token counts, then back to average
                    sum_cos = avg_cos * float(n_tokens_local)
                    sum_cos = self._all_reduce_sum_tensor(sum_cos)

                    n_tensor = avg_cos.new_tensor([float(n_tokens_local)])
                    if (
                        torch.distributed.is_available()
                        and torch.distributed.is_initialized()
                    ):
                        torch.distributed.all_reduce(
                            n_tensor,
                            op=torch.distributed.ReduceOp.SUM,
                        )
                    n_tokens_total = int(n_tensor.item()) if n_tensor.item() > 0 else 1

                    avg_cos = sum_cos / float(max(n_tokens_total, 1))
                else:
                    avg_cos = self._all_reduce_mean_tensor(avg_cos)

            self._block_cos_mats[block_idx] = avg_cos.cpu().float()
            self._block_token_counts[block_idx] = n_tokens_total

        return _hook

    def _remove_hooks(self) -> None:
        for h in self._handles:
            h.remove()
        self._handles.clear()

    # -------- Composer events (eval only) --------
    def eval_start(self, state: State, logger: Logger) -> None:  # noqa: ARG002
        """Reset accumulators at the beginning of eval."""
        if self.accumulate_across_eval:
            self._accum_sum_cos_mats.clear()
            self._accum_denoms.clear()
        self._eval_batches_seen = 0

    def eval_before_forward(self, state: State, logger: Logger) -> None:  # noqa: ARG002
        """Install hooks before each eval forward pass."""
        # Respect first-N-batches limit by skipping hooks once reached
        if (
            self.first_n_eval_batches is not None
            and self._eval_batches_seen >= self.first_n_eval_batches
        ):
            return
        self._eval_batches_seen += 1
        self._register_hooks_for_eval(state)

    def eval_after_forward(self, state: State, logger: Logger) -> None:  # noqa: ARG002
        """Remove hooks after each eval forward pass."""
        self._remove_hooks()

    def eval_batch_end(self, state: State, logger: Logger) -> None:
        """Accumulate or log cosine-similarity summaries and heatmaps."""
        if not self._block_cos_mats:
            return

        if self.accumulate_across_eval:
            # Accumulate across eval instead of logging per-batch
            for block_idx, cos_mat in self._block_cos_mats.items():
                denom_inc = (
                    float(self._block_token_counts.get(block_idx, 1))
                    if self.weighted_by_tokens
                    else 1.0
                )
                contrib = cos_mat * denom_inc
                if block_idx in self._accum_sum_cos_mats:
                    self._accum_sum_cos_mats[block_idx] += contrib
                    self._accum_denoms[block_idx] += denom_inc
                else:
                    self._accum_sum_cos_mats[block_idx] = contrib.clone()
                    self._accum_denoms[block_idx] = denom_inc

            # Clear for next batch
            self._block_cos_mats.clear()
            self._block_token_counts.clear()
            return

        # ------- Per-batch logging (original behavior) -------
        log_key_prefix = f"{self.log_key_prefix}/eval"
        scalar_metrics: dict[str, float] = {}
        image_metrics: dict[str, wandb.Image] = {}

        for block_idx, cos_mat in self._block_cos_mats.items():
            stats = self._summarize_cos_matrix(cos_mat)
            scalar_metrics.update(
                {
                    f"{log_key_prefix}/block_{block_idx}/{k}": v
                    for k, v in stats.items()
                },
            )
            # Heatmap image for visualization
            image_metrics[f"{log_key_prefix}/block_{block_idx}/cos_sim/heatmap"] = (
                create_heatmap(
                    similarity_matrix=cos_mat,
                    expert_ids=list(range(cos_mat.shape[0])),
                    step=state.eval_timestamp.batch.value,
                    opts=HeatmapOptions(
                        cmap="coolwarm",
                        vmin=-1.0,
                        vmax=1.0,
                        annot=False,
                    ),
                )
            )

        if scalar_metrics:
            logger.log_metrics(scalar_metrics)
        if image_metrics:
            logger.log_metrics(image_metrics)  # type: ignore[arg-type]

        # Clear for next batch
        self._block_cos_mats.clear()
        self._block_token_counts.clear()

    def eval_end(self, state: State, logger: Logger) -> None:
        """On eval end, if accumulating, emit aggregated metrics and heatmaps."""
        if not self.accumulate_across_eval:
            return
        if not self._accum_sum_cos_mats:
            return

        log_key_prefix = f"{self.log_key_prefix}/eval"
        scalar_metrics: dict[str, float] = {}
        image_metrics: dict[str, wandb.Image] = {}

        for block_idx, sum_mat in self._accum_sum_cos_mats.items():
            denom = max(self._accum_denoms.get(block_idx, 0.0), 1.0)
            avg_mat = (sum_mat / denom).cpu().float()

            stats = self._summarize_cos_matrix(avg_mat)
            scalar_metrics.update(
                {
                    f"{log_key_prefix}/block_{block_idx}/{k}": v
                    for k, v in stats.items()
                },
            )
            image_metrics[f"{log_key_prefix}/block_{block_idx}/cos_sim/heatmap"] = (
                create_heatmap(
                    similarity_matrix=avg_mat,
                    expert_ids=list(range(avg_mat.shape[0])),
                    step=state.eval_timestamp.batch.value,
                    opts=HeatmapOptions(
                        cmap="coolwarm",
                        vmin=-1.0,
                        vmax=1.0,
                        annot=False,
                    ),
                )
            )

        if scalar_metrics:
            logger.log_metrics(scalar_metrics)
        if image_metrics:
            logger.log_metrics(image_metrics)  # type: ignore[arg-type]

        # Clear accumulators
        self._accum_sum_cos_mats.clear()
        self._accum_denoms.clear()
