"""Monitor expert selection weights in SigmaMoE layers during training."""

from typing import TYPE_CHECKING, cast

import torch
from composer.core import Callback, State, Time, TimeUnit
from composer.loggers import Logger
from llmfoundry.registry import callbacks

from mt_dao.composer_callbacks.plotting_utils import HeatmapOptions, create_heatmap
from mt_dao.composer_callbacks.utils import (
    collect_routers_weights_statistics,
)

if TYPE_CHECKING:
    import wandb
    from mt_dao.models.mpt.mpt_mup import ComposerMPTMuPCausalLM, MPTMuPForCausalLM

__all__ = ["RouterSimilarityCallback"]


@callbacks.register("router_similarity_callback")
class RouterSimilarityCallback(Callback):
    """Monitor metrics related to the expert_sel parameter in SigmaMoE layers.

    This callback tracks cosine similarity metrics for the expert selection weights
    in SigmaMoE layers. It computes the cosine similarity between rows of the
    expert_sel matrix by normalizing row-wise and computing W @ W.T.

    The callback logs metrics for each transformer block's SigmaMoE expert_sel
    parameter:
    - Cosine similarity matrix statistics (mean, std, min, max)
    - Off-diagonal cosine similarity statistics
    - Diagonal cosine similarity statistics (should be close to 1.0)
    - Heatmap images of the cosine similarity matrix

    The callback is designed to be used with SigmaMoE models that have expert_sel
    parameters in their layers.

    Example:
        .. doctest::

            >>> from composer import Trainer
            >>> from mt_dao.composer_callbacks import RouterSimilarityCallback
            >>> # constructing trainer object with this callback
            >>> trainer = Trainer(
            ...     model=model,
            ...     train_dataloader=train_dataloader,
            ...     eval_dataloader=eval_dataloader,
            ...     optimizers=optimizer,
            ...     max_duration="1ep",
            ...     callbacks=[RouterSimilarityCallback()],
            ... )

    The metrics are logged with the following keys:
    - `expert_sel/cosine_sim/mean/{layer_name}` - Mean cosine similarity
    - `expert_sel/cosine_sim/std/{layer_name}` - Std deviation of cosine similarities
    - `expert_sel/cosine_sim/min/{layer_name}` - Minimum cosine similarity
    - `expert_sel/cosine_sim/max/{layer_name}` - Maximum cosine similarity
    - `expert_sel/cosine_sim/off_diag_mean/{layer_name}` - Mean off-diagonal similarity
    - `expert_sel/cosine_sim/off_diag_std/{layer_name}` - Std off-diagonal similarity
    - `expert_sel/cosine_sim/diag_mean/{layer_name}` - Mean diagonal similarity
    - `expert_sel/cosine_sim/heatmap/{layer_name}` - Full cosine similarity matrix

    """

    def __init__(
        self,
        *,
        log_interval: str = "8ba",
        log_key_prefix: str = "expert_router",
    ) -> None:
        """Initialize the RouterSimilarityCallback.

        Parameters
        ----------
        log_interval : str, optional
            How often to log the metrics, by default "8ba"
        log_key_prefix: str = "expert_router",
            Prefix for the logged keys, by default "expert_router"

        """
        self.log_interval = (
            Time.from_timestring(log_interval)
            if isinstance(log_interval, str)
            else Time(log_interval, TimeUnit.BATCH)
        )
        self.last_batch_logged = -1
        self.log_key_prefix = log_key_prefix
        # Initialize the list of experts IDs in case we need a form of remapping when
        # logging (e.g., when we are logging client-side information and it got a subset
        # of the global model experts only)
        self.expert_ids: list[int] = []

    def _should_log(self, state: State) -> bool:
        if not state.model.training:
            # If the model is not training, we log every time (once per evaluation)
            return True
        if isinstance(self.log_interval, Time):
            # NOTE: Always log first step
            if state.timestamp.batch == 0:
                return True
            # Get the correct timestamp depending on the state of the model
            return (
                state.timestamp.batch != self.last_batch_logged
                and state.timestamp.get(self.log_interval.unit)
                % self.log_interval.value
                == 0
            )
        return False

    def after_train_batch(self, state: State, logger: Logger) -> None:
        """Monitor expert_sel parameters at the end of each batch.

        This method is executed before the call to the optimizer step.

        Parameters
        ----------
        state : State
            The current state of training
        logger : Logger
            The logger to log metrics to

        """
        # Check if we should log
        if not self._should_log(state):
            return

        # Get the training/evaluation prefix for the metrics
        log_key_prefix = (
            f"{self.log_key_prefix}/train"
            if state.model.training
            else f"{self.log_key_prefix}/eval"
        )

        # Separate scalar metrics and prepare heatmap images
        log_dict_images: dict[str, wandb.Image] = {}

        # Get the transformer part of the model
        original_model = cast(
            "ComposerMPTMuPCausalLM",
            (
                state.model.module
                if isinstance(
                    state.model,
                    torch.nn.parallel.DistributedDataParallel,
                )
                else state.model
            ),
        )
        transformer = cast("MPTMuPForCausalLM", original_model.model).transformer

        expert_sel_metrics, cos_sim_matrices = collect_routers_weights_statistics(
            transformer,
        )

        # Create heatmap image
        if cos_sim_matrices:
            for block_idx, cosine_sim_matrix in cos_sim_matrices.items():
                log_dict_images.update(
                    {
                        f"{log_key_prefix}/block_{block_idx}/cosine_sim/heatmap": (
                            create_heatmap(
                                similarity_matrix=cosine_sim_matrix,
                                expert_ids=self.expert_ids,
                                step=(
                                    state.timestamp.batch.value
                                    if state.model.training
                                    else state.eval_timestamp.batch.value
                                ),
                                opts=HeatmapOptions(
                                    cmap="coolwarm",
                                    vmin=-1.0,
                                    vmax=1.0,
                                    annot=False,
                                ),
                            )
                        ),
                    },
                )
            # The logger cannot interpret the image input from its typing even though it
            # handles it correctly, so ignoring mypy here.
            logger.log_metrics(log_dict_images)  # type: ignore[arg-type]

        # Log scalar metrics
        if expert_sel_metrics:
            logger.log_metrics(
                {f"{log_key_prefix}/{k}": v for k, v in expert_sel_metrics.items()},
            )

        # Update last logged batch and clear the blocks histogram
        self.last_batch_logged = (
            state.timestamp.batch.value
            if state.model.training
            else state.eval_timestamp.batch.value
        )

    def before_eval_all(self, state: State, logger: Logger) -> None:
        """Collect and log expert selection parameters stats at final evaluation.

        Parameters
        ----------
        state: State
            The current state of the evaluation process, which includes the model and
            other relevant information.
        logger: Logger
            The logger used to log metrics and other information during evaluation.

        """
        return self.after_train_batch(state, logger)

    def eval_before_all(self, state: State, logger: Logger) -> None:
        """Collect and log expert selection parameters stats at intra fit evaluation.

        Parameters
        ----------
        state: State
            The current state of the evaluation process, which includes the model and
            other relevant information.
        logger: Logger
            The logger used to log metrics and other information during evaluation.

        """
        return self.after_train_batch(state, logger)
