"""Module for all-reduce operations on float across distributed ranks."""

from collections import defaultdict
from logging import WARNING

import torch
from composer import State
from flwr.common.logger import log
from torch import distributed as dist
from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel

from mt_dao.models.ffns.sigma_moe import SigmaMoE
from mt_dao.models.mpt.mpt_mup import MPTMuPModel
from mt_dao.utils import get_device


def _float_to_tensor(value: float) -> torch.Tensor:
    """Convert a float value to a tensor.

    Parameters
    ----------
    value: float
        The value to be converted.

    Returns
    -------
    torch.Tensor
        A tensor containing the value.

    """
    return torch.tensor(
        value,
        dtype=torch.get_default_dtype(),
        device=get_device(),
    )


def all_reduce_sum_across_ranks(value: float) -> float:
    """All-reduce-SUM the given value across all ranks.

    Parameters
    ----------
    value: float
        The value to be all-reduced.

    Returns
    -------
    float
        The all-reduced value.

    """
    tensor_value = _float_to_tensor(value)
    dist.all_reduce(tensor_value, op=dist.ReduceOp.SUM)
    return tensor_value.item()


def all_reduce_avg_across_ranks(value: float) -> float:
    """All-reduce-AVG the given value across all ranks.

    Parameters
    ----------
    value: float
        The value to be all-reduced.

    Returns
    -------
    float
        The all-reduced value.

    """
    tensor_value = _float_to_tensor(value)
    dist.all_reduce(tensor_value, op=dist.ReduceOp.AVG)
    return tensor_value.item()


def all_reduce_max_across_ranks(value: float) -> float:
    """All-reduce-MAX the given value across all ranks.

    Parameters
    ----------
    value: float
        The value to be all-reduced.

    Returns
    -------
    float
        The all-reduced value.

    """
    tensor_value = _float_to_tensor(value)
    dist.all_reduce(tensor_value, op=dist.ReduceOp.MAX)
    return tensor_value.item()


def all_reduce_min_across_ranks(value: float) -> float:
    """All-reduce-MIN the given value across all ranks.

    Parameters
    ----------
    value: float
        The value to be all-reduced.

    Returns
    -------
    float
        The all-reduced value.

    """
    tensor_value = _float_to_tensor(value)
    dist.all_reduce(tensor_value, op=dist.ReduceOp.MIN)
    return tensor_value.item()


def get_module_list_from_state(
    state: State,
) -> torch.nn.ModuleList:
    """Get the ModuleList from the model state.

    Parameters
    ----------
    state: State
        The current training state.

    Returns
    -------
    torch.nn.ModuleList
        The ModuleList from the model state.

    Raises
    ------
    ValueError
        If the model does not contain a ModuleList.

    """
    # Get the list of block from the model recursing into the children modules
    for _child_name, child in state.model.named_modules():
        if isinstance(child, torch.nn.ModuleList):
            return child
    msg = "The model does not contain a torch.nn.ModuleList."
    raise ValueError(msg)


def get_n_bins_from_fd_rule(
    sel_history: torch.Tensor,
    eps: float = 1e-9,
) -> int:
    """Get the number of bins for the histogram using the Freedman-Diaconis rule.

    Notably, this function assumes that the selection history under each rank will
    provide sufficient statistical significance for estimating the quartiles.

    Parameters
    ----------
    sel_history: torch.Tensor
        A tensor containing expert selections for each token, with shape
        (micro_batch_size * seq_length, n_experts).
    eps: float
        A small value to avoid division by zero.

    Returns
    -------
    int
        The number of bins to use for the histogram.

    """
    # NOTE: We are here assuming that this rank selection history will provide
    # sufficient statistical significance for estimating the quartiles.
    # Get the number of datapoints per expert selection history
    n_datapoints = sel_history.shape[0]
    # Get the average interquartile range (IQR) across experts
    q1 = torch.quantile(sel_history.float(), q=0.25, dim=0)
    q3 = torch.quantile(sel_history.float(), q=0.75, dim=0)
    mean_iqr_across_experts = (q3 - q1).mean().item() + eps
    return int(1 / (2 * mean_iqr_across_experts / n_datapoints ** (1 / 3)))


def collect_distributed_selection_history(
    ffn: SigmaMoE,
) -> torch.Tensor:
    """Collect the full selection history from the SigmaMoE across distributed ranks.

    Parameters
    ----------
    ffn: SigmaMoE
        The SigmaMoE from which to collect the full selection history.

    Returns
    -------
    torch.Tensor
        The full selection history from the SigmaMoE.

    """
    # Get the element-wise sigmoided selection scores history whose shape is
    # (batch_size, seq_length, n_experts).
    # NOTE: If we want to use this for some loss we need to make sure the values are
    # consistent by creating a copy here so as to maintain the original data in the
    # variable.
    block_full_selection_history = torch.cat(ffn.sel_hist)
    block_full_selection_history_cpu = (
        block_full_selection_history.detach().clone().cpu().float()
    )
    # Gather across ranks
    if torch.distributed.is_initialized():
        # Get process group
        group = dist.GroupMember.WORLD
        assert group is not None, (
            "The process group is not set. Please ensure that the distributed"
            " environment is initialized correctly."
        )
        # Loop over the ranks and gather the selection history
        for rank in dist.get_process_group_ranks(group=group):
            # Exclude Rank-0 from the gather operation
            if rank == 0:
                continue
            dist.broadcast(
                block_full_selection_history,
                src=rank,
                group=group,
            )
            block_full_selection_history_cpu = torch.cat(
                [
                    block_full_selection_history_cpu,
                    block_full_selection_history.detach().clone().cpu().float(),
                ],
                dim=0,
            )
    return torch.nn.functional.sigmoid(block_full_selection_history_cpu)


def compute_histograms(
    sel_history: torch.Tensor,
    n_bins: int | None = None,
    *,
    scores_histograms: bool = True,
    topk_heatmaps: bool = True,
    log_warning: bool = False,
) -> tuple[
    dict[int, torch.Tensor],
    dict[int, tuple[torch.Tensor, torch.Tensor]],
    int | None,
]:
    """Compute histograms of expert selections for top-k experts.

    Parameters
    ----------
    sel_history: torch.Tensor
        A tensor containing expert selections for each token, with shape
        (micro_batch_size * seq_length, n_experts).
    n_bins: int | None
        The number of bins to use for the histogram. If None, the number of bins will
        be determined using the Freedman-Diaconis rule.
    scores_histograms: bool
        If True, compute histograms of scores for each expert.
    topk_heatmaps: bool
        If True, compute histograms of expert selections for each possible value of k in
        the top-k selection.
    log_warning: bool
        If True, log a warning if the number of bins is determined using the
        Freedman-Diaconis rule.

    Returns
    -------
    tuple[
        dict[int, torch.Tensor],
        dict[int, tuple[torch.Tensor, torch.Tensor]],
        int | None
    ]:
        A tuple containing:
        - A dict of tensors representing the histograms of expert selections for each
            possible value of k in the top-k selection.
        - A dict of tuples, each containing a histogram of scores and the corresponding
            bin edges for each k in the top-k selection.
        - The number of bins used for the histograms or None.
        The histograms are normalized to represent the probability distribution of
        expert selections across the batch and sequence length.

    Raises
    ------
    ValueError
        If both `scores_histograms` and `topk_heatmaps` are False, a ValueError is
        raised indicating that at least one must be True to compute histograms.

    """
    # Return immediately if both boolean flags are False
    if not scores_histograms and not topk_heatmaps:
        msg = (
            "Both scores_histograms and topk_heatmaps are False. At least one must be"
            " True to compute histograms."
        )
        raise ValueError(msg)
    # Initialize the return list
    experts_topk_selection_counts: dict[int, torch.Tensor] = {}
    experts_scores_histograms: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
    # Get the number of experts, sel_history has shape
    # (micro_batch_size * seq_length, n_experts).
    n_experts = sel_history.shape[-1]
    # Get n_bins from `sel_history`
    if n_bins is None:
        n_bins = get_n_bins_from_fd_rule(sel_history)
    if log_warning:
        log(WARNING, "Using %d bins for the histogram", n_bins)
    # Loop over the possible number of experts to select (from 1 to n_experts)
    for e_idx in range(n_experts):
        if topk_heatmaps:
            # Get the value for k of topk
            k = e_idx + 1
            # Get the top-i expert, shape (micro_batch_size * seq_length, i)
            top_k_expert = torch.topk(sel_history, k, dim=-1).indices
            # Compute how many tokens have been assigned to each expert
            expert_usage = torch.bincount(
                # Flatten on all dimensions to obtain the list of assigned experts
                top_k_expert.view(-1),
                minlength=n_experts,
            ).float()
            # Append the histogram to the list
            experts_topk_selection_counts[k] = expert_usage
        if scores_histograms:
            # Compute the histogram of scores for the e_idx-th expert
            (bin_counts, bin_edges) = torch.histogram(
                sel_history[:, e_idx],
                bins=n_bins,
                range=(0, 1),
            )
            experts_scores_histograms[e_idx] = (bin_counts, bin_edges)
    return experts_topk_selection_counts, experts_scores_histograms, n_bins


def clear_expert_selections(
    transformer: MPTMuPModel,
) -> None:
    """Clear the expert selections in the transformer model.

    Parameters
    ----------
    transformer: MPTMuPModel
        The transformer model from which to clear expert selections.

    """
    # Initialize a boolean that tells whether to make the empty cache call
    cache_needs_empty = False
    # Loop over the transformer blocks
    for block in transformer.blocks:
        # Get the FFN
        ffn = block.ffn
        if not isinstance(ffn, SigmaMoE):
            continue
        if len(ffn.sel_hist) > 0:
            ffn.sel_hist.clear()
            cache_needs_empty = True
    if cache_needs_empty:
        torch.cuda.empty_cache()


def collect_expert_selections(  # noqa: PLR0913
    transformer: MPTMuPModel,
    n_bins: int | None = None,
    experts_per_token_threshold: float = 0.5,
    *,
    full_selection_history: bool = False,
    scores_histograms: bool = True,
    topk_heatmaps: bool = True,
) -> tuple[
    dict[int, torch.Tensor],
    dict[int, dict[int, tuple[torch.Tensor, torch.Tensor]]],
    dict[int, dict[int, torch.Tensor]],
    dict[int, torch.Tensor],
    int | None,
]:
    """Collect expert selections from the transformer model.

    Parameters
    ----------
    transformer: MPTMuPModel
        The transformer model from which to collect expert selections.
    n_bins: int | None
        The number of bins to use for the histogram. If None, the number of bins will
        be determined using the Freedman-Diaconis rule.
    experts_per_token_threshold: float
        The threshold to use when counting how many experts a token should be
        assigned to. Defaults to 0.5.
    full_selection_history: bool
        If True, collect the full selection history from the SigmaMoE.
    scores_histograms: bool
        If True, compute histograms of scores for each expert.
    topk_heatmaps: bool
        If True, compute histograms of expert selections for each possible value of k in
        the top-k selection.

    Returns
    -------
    tuple[
        dict[int, dict[int, torch.Tensor]],
        dict[int, dict[int, tuple[torch.Tensor, torch.Tensor]]],
        dict[int, dict[int, torch.Tensor]],
        dict[int, torch.Tensor],
        int | None,
    ]
        A tuple containing:
        - A dictionary mapping block indices to tensors representing the full
            selection history for each block.
        - A dictionary mapping block indices to dictionaries of expert scores
            histograms, where each dictionary maps expert indices to tuples of
            (histogram, bin edges).
        - A dictionary mapping block indices to dictionaries of expert top-k selection
            counts, where each dictionary maps expert indices to tensors representing
            the counts.
        - A dictionary mapping block indices to tensors representing the token-wise
            thresholded scores histogram counts for each value in [0, n_experts].
        - An int representing the number of bins used for the experts scores histograms.

    """
    # Return immediately if both boolean flags are False
    if not scores_histograms and not topk_heatmaps:
        log(WARNING, "Both scores_histograms and topk_heatmaps are False.")
        return {}, {}, {}, {}, n_bins
    # Initialize the return dictionaries
    blocks_full_selection_history: dict[int, torch.Tensor] = {}
    blocks_scores_histograms: dict[
        int,
        dict[int, tuple[torch.Tensor, torch.Tensor]],
    ] = defaultdict(dict)
    blocks_experts_topk_selection_counts: dict[int, dict[int, torch.Tensor]] = (
        defaultdict(
            dict,
        )
    )
    blocks_token_wise_threshold_histogram: dict[int, torch.Tensor] = defaultdict(
        torch.Tensor,
    )
    # Loop over the transformer blocks
    for i, block in enumerate(transformer.blocks):
        # Get the FFN
        ffn = block.ffn
        if not isinstance(ffn, SigmaMoE):
            continue
        assert len(ffn.sel_hist) == 1, (
            f"Block {i} is expected to have one tensor in the sel_hist, but it has"
            f" {len(ffn.sel_hist)}."
        )
        # Collect the full selection history
        selection_history = collect_distributed_selection_history(ffn)
        if full_selection_history:
            blocks_full_selection_history[i] = selection_history
        # Round the scores given a threshold of `experts_per_token_threshold`.
        # Values below the threshold become zero; values above become one
        # (strictly greater than threshold -> 1).
        # Keep dtype to be compatible with downstream ops.
        rounded_selection_history = (
            selection_history > experts_per_token_threshold
        ).to(selection_history.dtype)
        # Count token-wise the number of ones in the rounded selection history
        token_wise_threshold_counts = rounded_selection_history.sum(dim=-1)
        # Bin the token wise thresholded counts into a histogram binned by the number of
        # total experts
        blocks_token_wise_threshold_histogram[i] = torch.histc(
            token_wise_threshold_counts.float(),
            # NOTE: We include the bin for those tokens requiring zero experts.
            bins=ffn.n_experts + 1,
            min=0,
            max=ffn.n_experts,
        )
        # Get the list of histograms for the top-k experts
        with torch.no_grad():
            topk_histograms, experts_scores_histograms, n_bins = compute_histograms(
                # NOTE: The following flattening is likely unecessary
                # Flatten on the token dimensions (batch_size * seq_length, n_experts)
                sel_history=selection_history.view(-1, selection_history.shape[-1]),
                n_bins=n_bins,
                scores_histograms=scores_histograms,
                topk_heatmaps=topk_heatmaps,
            )
        # Append the histogram to the list
        blocks_experts_topk_selection_counts[i] = topk_histograms
        blocks_scores_histograms[i] = experts_scores_histograms
    return (
        blocks_full_selection_history,
        blocks_scores_histograms,
        blocks_experts_topk_selection_counts,
        blocks_token_wise_threshold_histogram,
        n_bins,
    )


def collect_router_weights_statistics(
    router_weights: torch.Tensor,
) -> tuple[dict[str, float], torch.Tensor]:
    """Collect statistics on router weights.

    Parameters
    ----------
    router_weights: torch.Tensor
        The router weights tensor, expected to have shape (n_experts, n_features).

    Returns
    -------
    tuple[dict[str, float], torch.Tensor]
        A tuple containing:
        - A dictionary of statistics related to cosine similarity.
        - A cosine similarity matrix of the router weights.

    """
    # Initialize the metrics dictionary
    metrics: dict[str, float] = {}
    # Row-wise L2 normalization
    router_weights_normalized = torch.nn.functional.normalize(
        router_weights,
        p=2,
        dim=1,
    )

    # TODO(anon): Assess whether we can come up with a metric that takes out the
    # "contribution" of the random initialization, i.e., the distance between randomly
    # initialized vectors.
    # TODO(anon): Add a metric that assesses linear independence of the vectors
    # (eigenvectors).
    # TODO(anon): Add eigenvalues.
    # Compute cosine similarity matrix: W @ W.T
    cos_sim_matrix = torch.mm(
        router_weights_normalized,
        router_weights_normalized.t(),
    )

    # Overall cosine similarity statistics
    metrics["cos_sim/mean"] = cos_sim_matrix.mean().item()
    metrics["cos_sim/std"] = cos_sim_matrix.std().item()
    metrics["cos_sim/min"] = cos_sim_matrix.min().item()
    metrics["cos_sim/max"] = cos_sim_matrix.max().item()

    # Off-diagonal elements (similarities between different experts)
    n_experts = cos_sim_matrix.shape[0]
    if n_experts > 1:
        # Create mask for off-diagonal elements
        off_diag_mask = ~torch.eye(
            n_experts,
            dtype=torch.bool,
            device=router_weights_normalized.device,
        )
        off_diag_similarities = cos_sim_matrix[off_diag_mask]

        metrics["cos_sim/off_diag_mean"] = off_diag_similarities.mean().item()
        metrics["cos_sim/off_diag_std"] = off_diag_similarities.std().item()
        metrics["cos_sim/off_diag_min"] = off_diag_similarities.min().item()
        metrics["cos_sim/off_diag_max"] = off_diag_similarities.max().item()

    # Diagonal elements (should be close to 1.0 for normalized vectors)
    diagonal_similarities = torch.diag(cos_sim_matrix)
    metrics["cos_sim/diag_mean"] = diagonal_similarities.mean().item()
    metrics["cos_sim/diag_std"] = diagonal_similarities.std().item()
    return metrics, cos_sim_matrix


def get_expert_sel_weights(
    block: torch.nn.Module,
) -> torch.Tensor | None:
    """Retrieve the expert selection weights from a SigmaMoE block.

    Parameters
    ----------
    block: torch.nn.Module
        The block from which to retrieve the expert selection weights.

    Returns
    -------
    torch.Tensor | None
        The expert selection weights tensor if found, otherwise None.
        If the block is a FullyShardedDataParallel, it will summon the full parameters
        to access the expert selection weights.

    Raises
    ------
    ValueError
        If no expert selection weights are found in the block.

    """
    return_tensor_dict: dict[str, torch.Tensor] = {}
    with FullyShardedDataParallel.summon_full_params(
        block,
        recurse=True,
        writeback=False,
        rank0_only=True,
        offload_to_cpu=True,
        with_grads=False,
    ):
        for name, param in block.named_parameters():
            if "expert_sel" in name:
                return_tensor_dict[name] = param.data.detach().clone().cpu().float()
    # NOTE: Here, we raise only on Rank-0 which can be misleading in distributed
    # environments where error handling may be problematic.
    # TODO(anon): Propagate error to other ranks too.
    if not return_tensor_dict and dist.get_rank() == 0:
        msg = f"Expert selection weights not found in {block}."
        raise ValueError(msg)
    if return_tensor_dict:
        assert len(return_tensor_dict) == 1, (
            "Expected exactly one expert selection weight,"
            f" but found {len(return_tensor_dict)}."
        )
        return next(iter(return_tensor_dict.values()))
    return None


def collect_routers_weights_statistics(
    transformer: MPTMuPModel,
) -> tuple[dict[str, float], dict[int, torch.Tensor]]:
    """Collect statistics on router weights from the transformer model.

    Parameters
    ----------
    transformer: MPTMuPModel
        The transformer model from which to collect router weights statistics.

    Returns
    -------
    tuple[dict[str, float], dict[int, torch.Tensor]]
        A tuple containing:
        - A dictionary of metrics related to router weights.
        - A dictionary mapping block indices to cosine similarity matrices of router
          weights.

    """
    # Initialize the metrics dictionary and the dict of cosine similarity matrices
    expert_sel_metrics: dict[str, float] = {}
    cos_sim_matrices: dict[int, torch.Tensor] = {}
    # Loop over the transformer blocks
    for block_idx, block in enumerate(transformer.blocks):
        # Get the FFN
        ffn = block.ffn
        if not isinstance(ffn, SigmaMoE):
            continue
        router_weights = (
            get_expert_sel_weights(block)
            if isinstance(
                block,
                FullyShardedDataParallel,
            )
            else ffn.expert_sel.detach().clone().cpu().float()
        )
        if router_weights is not None:
            # Warn if router weights are flat
            if router_weights.ndim == 1:
                log(
                    WARNING,
                    "Router weights in block %d are flat, they will be reshaped"
                    " depending on the number of experts in the model"
                    " (`router_weights.view(ffn.n_experts, -1)`) but it might not"
                    " reflect what PyTorch does under the hood for FSDP.",
                    block_idx,
                )
            # Collect router weights statistics
            block_metrics, block_cos_sim_matrix = collect_router_weights_statistics(
                router_weights.view(ffn.n_experts, -1),
            )
            # Store the cosine similarity matrix for this block
            cos_sim_matrices[block_idx] = block_cos_sim_matrix
            # Update the metrics dictionary with block-specific metrics
            expert_sel_metrics.update(
                {f"block_{block_idx}/{k}": v for k, v in block_metrics.items()},
            )
        dist.barrier()
    return expert_sel_metrics, cos_sim_matrices


def get_full_selection_history(
    transformer: MPTMuPModel,
) -> dict[int, torch.Tensor]:
    """Get the full selection history from the transformer model.

    Parameters
    ----------
    transformer: MPTMuPModel
        The transformer model from which to collect the full selection history.

    Returns
    -------
    dict[int, torch.Tensor]
        A dictionary mapping block indices to tensors representing the full
        selection history for each block.

    """
    full_selection_history: dict[int, torch.Tensor] = {}
    # Loop over the transformer blocks
    for i, block in enumerate(transformer.blocks):
        # Get the FFN
        ffn = block.ffn
        if not isinstance(ffn, SigmaMoE):
            continue
        assert len(ffn.sel_hist) == 1, (
            f"Block {i} is expected to have one tensor in the sel_hist, but it has"
            f" {len(ffn.sel_hist)}."
        )
        # Collect the full selection history
        full_selection_history[i] = collect_distributed_selection_history(ffn)
    return full_selection_history
