"""Callback to log expert selection statistics for MPTMuP models."""

from collections import defaultdict
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 (
    create_expert_heatmap,
    create_expert_scores_histogram,
    create_expert_tokens_scores_histogram,
    create_thresholded_expert_scores_histogram,
    create_tokens_kl_histograms,
)
from mt_dao.composer_callbacks.utils import (
    clear_expert_selections,
    collect_expert_selections,
)

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


@callbacks.register("experts_selection_callback")
class ExpertSelectionCallback(Callback):
    """Callback to log expert selection statistics for MPTMuP models."""

    def __init__(  # noqa: PLR0913
        self,
        log_interval: str = "16ba",
        log_key_prefix: str = "expert_selection",
        *,
        scores_histograms: bool = True,
        topk_heatmaps: bool = True,
        token_wise_scores: bool = False,
        token_wise_max_entries: int = 10_000,
        token_wise_samples: int = 8,
        experts_per_token_threshold: float = 0.5,
    ) -> None:
        """Initialize the ExpertSelectionCallback.

        Parameters
        ----------
        log_interval: str
            The interval at which to log expert selection statistics.
            Defaults to "16ba".
        log_key_prefix: str
            The prefix to use for logging keys.
            Defaults to "expert_selection".
        scores_histograms: bool
            Whether to log histograms of expert scores.
            Defaults to True.
        topk_heatmaps: bool
            Whether to log heatmaps of top-k expert selections.
            Defaults to True.
        token_wise_scores: bool
            Whether to log token-wise scores for expert selection.
            Defaults to False.
        token_wise_max_entries: int
            The maximum number of entries to include in the token-wise scores logging.
            Defaults to 10,000.
        token_wise_samples: int
            The number of samples to use for token-wise scores logging.
            Defaults to 8.
        experts_per_token_threshold: float
            The threshold to use when counting how many experts a token should be
            assigned to. Defaults to 0.5.

        """
        self.log_interval = (
            Time.from_timestring(log_interval)
            if isinstance(log_interval, str)
            else Time(log_interval, TimeUnit.BATCH)
        )
        self.log_key_prefix = log_key_prefix
        self.scores_histograms = scores_histograms
        self.topk_heatmaps = topk_heatmaps
        self.token_wise_scores = token_wise_scores
        self.token_wise_max_entries = token_wise_max_entries
        self.token_wise_samples = token_wise_samples
        self.experts_per_token_threshold = experts_per_token_threshold
        self.last_batch_logged = -1
        self.blocks_experts_topk_selection_counts: dict[
            int,
            dict[int, torch.Tensor],
        ] = defaultdict(lambda: defaultdict(torch.Tensor))
        self.blocks_scores_histograms: dict[
            int,
            dict[int, tuple[torch.Tensor, torch.Tensor]],
        ] = defaultdict(lambda: defaultdict(lambda: (torch.Tensor(), torch.Tensor())))
        self.blocks_token_wise_threshold_histogram: dict[int, torch.Tensor] = (
            defaultdict(torch.Tensor)
        )
        # NOTE: The n_bins attribute is shared across all blocks but computed using the
        # Freedman-Diaconis rule from a single forward (microbatch) block's data. This
        # could lead to suboptimal binning for blocks with different statistical
        # properties, suboptimality which we are happy to pay in favor of efficiency.
        self.n_bins: int | None = None
        self.full_selection_history: dict[int, torch.Tensor] = {}
        # 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_collect(self, state: State) -> bool:
        if isinstance(self.log_interval, Time):
            # Get the correct timestamp depending on the state of the model
            reference_timestamp = (
                state.timestamp if state.model.training else state.eval_timestamp
            )
            # NOTE: Always log first step
            if reference_timestamp.batch == 0:
                return True
            # NOTE: In case of non batch-based logging, it is impossible here to know
            # whether we should not log as either of the objects available are aware of
            # wheter this batch will increase other time units or not. As such we return
            # immediately without any check.
            if self.log_interval.unit != TimeUnit.BATCH:
                return reference_timestamp.batch + 1 != self.last_batch_logged
            # NOTE: At this stage, the batch number is not yet updated, so we use it
            # increased by one
            increased_batch = reference_timestamp.batch + 1
            increased_timestamp = reference_timestamp.get(self.log_interval.unit) + 1
            return (
                increased_batch != self.last_batch_logged
                and increased_timestamp % self.log_interval.value == 0
            )
        return False

    def _should_log(self, state: State) -> bool:
        if isinstance(self.log_interval, Time):
            # Get the correct timestamp depending on the state of the model
            reference_timestamp = (
                state.timestamp if state.model.training else state.eval_timestamp
            )
            # NOTE: Always log first step
            if reference_timestamp.batch == 1:
                return True
            # Get the correct timestamp depending on the state of the model
            return (
                reference_timestamp.batch != self.last_batch_logged
                and reference_timestamp.get(self.log_interval.unit)
                % self.log_interval.value
                == 0
            )
        return False

    def _store_blocks_full_selection_history(
        self,
        blocks_full_selection_history: dict[int, torch.Tensor],
    ) -> None:
        # Loop over the blocks full selection history
        for (
            b_idx,
            block_full_selection_history,
        ) in blocks_full_selection_history.items():
            self.full_selection_history[b_idx] = (
                torch.cat(
                    [
                        block_full_selection_history,
                        self.full_selection_history[b_idx],
                    ],
                )
                if b_idx in self.full_selection_history
                else block_full_selection_history
            )

    def _store_blocks_scores_histograms(
        self,
        blocks_scores_histograms: dict[
            int,
            dict[int, tuple[torch.Tensor, torch.Tensor]],
        ],
    ) -> None:
        # Loop over the blocks histogram
        for b_idx, block_scores_histograms in blocks_scores_histograms.items():
            # Loop over the expert scores histograms for each block
            for e_idx, (
                block_expert_scores_counts,
                block_expert_scores_bins,
            ) in block_scores_histograms.items():
                block_new_scores_counts = (
                    torch.add(
                        block_expert_scores_counts,
                        self.blocks_scores_histograms[b_idx][e_idx][0],
                    )
                    if b_idx in self.blocks_scores_histograms
                    and e_idx in self.blocks_scores_histograms[b_idx]
                    else block_expert_scores_counts
                )
                self.blocks_scores_histograms[b_idx][e_idx] = (
                    block_new_scores_counts,
                    block_expert_scores_bins,
                )

    def _store_blocks_experts_topk_selection_counts(
        self,
        blocks_experts_topk_selection_counts: dict[int, dict[int, torch.Tensor]],
    ) -> None:
        # Loop over the blocks histogram
        for (
            b_idx,
            block_experts_topk_selection_counts,
        ) in blocks_experts_topk_selection_counts.items():
            # Loop over the expert selection counts for each block
            for (
                k_idx,
                block_k_topk_selection_counts,
            ) in block_experts_topk_selection_counts.items():
                self.blocks_experts_topk_selection_counts[b_idx][k_idx] = (
                    torch.add(
                        block_k_topk_selection_counts,
                        self.blocks_experts_topk_selection_counts[b_idx][k_idx],
                    )
                    if b_idx in self.blocks_experts_topk_selection_counts
                    and k_idx in self.blocks_experts_topk_selection_counts[b_idx]
                    else block_k_topk_selection_counts
                )

    def _store_blocks_token_wise_threshold_histogram(
        self,
        blocks_token_wise_threshold_histogram: dict[int, torch.Tensor],
    ) -> None:
        # Loop over the blocks full selection history
        for (
            b_idx,
            block_token_wise_threshold_histogram,
        ) in blocks_token_wise_threshold_histogram.items():
            self.blocks_token_wise_threshold_histogram[b_idx] = (
                torch.add(
                    block_token_wise_threshold_histogram,
                    self.blocks_token_wise_threshold_histogram[b_idx],
                )
                if b_idx in self.blocks_token_wise_threshold_histogram
                else block_token_wise_threshold_histogram
            )

    def before_train_batch(  # noqa: PLR6301
        self,
        state: State,
        logger: Logger,  # noqa: ARG002
    ) -> None:
        """Clear the expert selections cache before the training batch.

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

        """
        # 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
        # Clear the expert selections cache if not empty
        clear_expert_selections(transformer=transformer)

    def after_forward(
        self,
        state: State,
        logger: Logger,  # noqa: ARG002
    ) -> None:
        """Accumulate expert selection statistics after the forward pass.

        This method collects expert selection data from the model and updates the
        blocks histogram. It is called after the forward pass of the training step.

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

        """
        # Check if we should collect
        if not self._should_collect(state):
            return
        # 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
        # Collect expert selection data (dict[int, list[torch.Tensor]], containing one
        # histogram for each possible `k` to topk and for each block)
        (
            blocks_full_selection_history,
            blocks_scores_histograms,
            blocks_experts_topk_selection_counts,
            blocks_token_wise_threshold_histogram,
            self.n_bins,
        ) = collect_expert_selections(
            transformer=transformer,
            n_bins=self.n_bins,
            full_selection_history=self.token_wise_scores,
            scores_histograms=self.scores_histograms,
            topk_heatmaps=self.topk_heatmaps,
            experts_per_token_threshold=self.experts_per_token_threshold,
        )
        # Handle the returned objects
        if blocks_full_selection_history:
            self._store_blocks_full_selection_history(blocks_full_selection_history)
        if blocks_scores_histograms:
            self._store_blocks_scores_histograms(blocks_scores_histograms)
        if blocks_experts_topk_selection_counts:
            self._store_blocks_experts_topk_selection_counts(
                blocks_experts_topk_selection_counts,
            )
        if blocks_token_wise_threshold_histogram:
            self._store_blocks_token_wise_threshold_histogram(
                blocks_token_wise_threshold_histogram,
            )

    def batch_end(self, state: State, logger: Logger) -> None:
        """Log expert selection statistics at the end of the batch.

        This method logs the expert selection statistics after each batch, including
        the expert selection heatmap and load balance metrics for each block.
        It is called at the end of each batch during training.

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

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

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

        # 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"
        )

        # Create token-wise scores histograms and figures
        if self.token_wise_scores:
            for (
                block_idx,
                block_full_selection_history,
            ) in self.full_selection_history.items():
                # Create tokens KL divergence histograms
                figures_dict = create_tokens_kl_histograms(
                    block_full_selection_history=block_full_selection_history,
                    expert_ids=self.expert_ids,
                    step=(
                        state.timestamp.batch.value
                        if state.model.training
                        else state.eval_timestamp.batch.value
                    ),
                    max_entries=self.token_wise_max_entries,
                )
                log_dict_images.update(
                    {
                        (
                            f"{log_key_prefix}/token_histogram/"
                            f"block_{block_idx}/{fig_name}"
                        ): fig
                        for fig_name, fig in figures_dict.items()
                    },
                )
                # Randomly slice 16 samples from the full selection history
                indices = torch.randperm(
                    block_full_selection_history.shape[0],
                )[: self.token_wise_samples]
                # Create a heatmap for the expert selection
                log_dict_images[
                    f"{log_key_prefix}/token_histogram/block_{block_idx}"
                ] = create_expert_tokens_scores_histogram(
                    block_full_selection_history=block_full_selection_history[
                        indices,
                        :,
                        :,
                    ],
                    expert_ids=self.expert_ids,
                    step=(
                        state.timestamp.batch.value
                        if state.model.training
                        else state.eval_timestamp.batch.value
                    ),
                )
        # Create heatmaps for FFN selections
        for (
            block_idx,
            block_histograms,
        ) in self.blocks_experts_topk_selection_counts.items():
            # For each value of k in the topk selection, compute the load balance as the
            # KL divergence from a uniform distribution
            for k, block_histogram in block_histograms.items():
                # Computing the Kullback-Leibler divergence against a uniform
                # distribution for the block histogram after normalization
                log_dict_scalar_metrics[
                    (f"{log_key_prefix}/block_{block_idx}/top-{k}/kl_vs_uniform")
                ] = torch.nn.functional.kl_div(
                    # NOTE: The block_histogram countains the binned counts
                    ((block_histogram / block_histogram.sum()) + 1e-8).log(),
                    # Uniform distribution over block_histogram.shape[0] values
                    torch.ones(block_histogram.shape) / block_histogram.shape[0],
                    reduction="sum",
                ).item()
            # Plot the heatmap for the top-k expert selections
            log_dict_images[f"{log_key_prefix}/heatmap/block_{block_idx}"] = (
                create_expert_heatmap(
                    block_histograms=block_histograms,
                    expert_ids=self.expert_ids,
                    step=(
                        state.timestamp.batch.value
                        if state.model.training
                        else state.eval_timestamp.batch.value
                    ),
                )
            )
        # Create histograms for FFN selection scores
        for block_idx, block_scores_histograms in self.blocks_scores_histograms.items():
            # TODO(anon): Compare each experts distribution with a beta distribution
            # with alpha=beta < 1.0, which we believe being the optimal routers should
            # follow for balanced and informative gradient flow
            # Create a histogram for the expert selection scores and collect statistics
            fig, statistics_dict = create_expert_scores_histogram(
                block_scores_histograms=block_scores_histograms,
                expert_ids=self.expert_ids,
                step=(
                    state.timestamp.batch.value
                    if state.model.training
                    else state.eval_timestamp.batch.value
                ),
            )
            # Add the statistics to the metrics dictionary
            log_dict_scalar_metrics.update(
                {
                    f"{log_key_prefix}/scores_histogram/block_{block_idx}/"
                    f"expert_{expert_id}/{stat_name}": stat_value
                    for expert_id, expert_stat_dict in statistics_dict.items()
                    for stat_name, stat_value in expert_stat_dict.items()
                },
            )
            # Add the histogram figure to the metrics dictionary
            log_dict_images[f"{log_key_prefix}/scores_histogram/block_{block_idx}"] = (
                fig
            )
        # Create histograms for the thresholded selection scores
        for (
            block_idx,
            block_token_wise_threshold_histogram,
        ) in self.blocks_token_wise_threshold_histogram.items():
            # Create a histogram for the expert selection scores and collect statistics
            fig = create_thresholded_expert_scores_histogram(
                block_token_wise_threshold_histogram=block_token_wise_threshold_histogram,
                step=(
                    state.timestamp.batch.value
                    if state.model.training
                    else state.eval_timestamp.batch.value
                ),
            )
            # Add the histogram figure to the metrics dictionary
            log_dict_images[
                f"{log_key_prefix}/thresholded_scores_histogram/block_{block_idx}"
            ] = fig

        # Log scalar metrics
        logger.log_metrics(log_dict_scalar_metrics)

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

        # 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
        )
        self.blocks_experts_topk_selection_counts.clear()
        self.blocks_scores_histograms.clear()
        self.full_selection_history.clear()
        self.n_bins = None

    def eval_after_forward(self, state: State, logger: Logger) -> None:
        """Execute after the forward pass during evaluation.

        This method is called after the forward pass during evaluation and
        collects expert selection statistics if the model is in evaluation mode.

        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_forward(state, logger)

    def eval_batch_end(self, state: State, logger: Logger) -> None:
        """Execute at the end of the batch during evaluation.

        This method is called at the end of each batch during evaluation and logs
        expert selection statistics if the model is in evaluation mode.

        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.batch_end(state, logger)
