"""Utility functions of general use for PyTorch."""

from collections import defaultdict
from typing import TYPE_CHECKING, cast

import torch

if TYPE_CHECKING:
    from llmfoundry import MPTConfig


def get_n_experts_from_model(
    model: torch.nn.Module,
) -> int | None:
    """Retrieve the number of experts from the model configuration.

    Parameters
    ----------
    model : torch.nn.Module
        The model from which to retrieve the number of experts.

    Returns
    -------
    int | None
        The number of experts if defined in the model configuration, otherwise None.

    """
    n_experts: int | None = None
    if cast("MPTConfig", model.config).ffn_config["ffn_type"] == "sigma_moe":
        n_experts = cast("MPTConfig", model.config).ffn_config.get(
            "ff_n_experts",
        )
        assert n_experts is not None, (
            "The number of experts must be defined in the MPTConfig "
            "ffn_config dictionary under the key 'ff_n_experts'. ffn_config dict ="
            f" {cast('MPTConfig', model.config).ffn_config}"
        )
        n_experts = int(n_experts)
    return n_experts


def invert_opt_state_dict_keys(
    opt_state_dict: dict[str, dict[str, torch.Tensor]],
) -> dict[str, dict[str, torch.Tensor]]:
    """Invert the keys of the optimizer state dict.

    This function takes an optimizer state dictionary where the outer keys are layer
    names and the inner keys are optimizer state names. It returns a new dictionary with
    the keys inverted, such that the outer keys become the inner keys and vice versa.

    Notably, the function will also work in the opposite direction, allowing you to
    convert a dictionary with inverted keys back to its original form.

    Parameters
    ----------
    opt_state_dict : dict[str, dict[str, torch.Tensor]]
        The optimizer state dictionary. The keys of the outer dictionary are usually the
        layer names, while the keys of the inner dictionaries are usually the optimizer
        state names (e.g., 'exp_avg', 'exp_avg_sq'). It can also be inverted, where the
        outer keys are optimizer state names and the inner keys are layer names.

    Returns
    -------
    dict[str, dict[str, torch.Tensor]]
        A new dictionary with inverted keys.

    """
    inverted_dict: dict[str, dict[str, torch.Tensor]] = defaultdict(dict)

    # Get all unique inner keys by examining all outer keys
    inner_keys: set[str] = set()
    for inner_dict in opt_state_dict.values():
        inner_keys.update(inner_dict.keys())

    # Create the inverted structure
    for inner_key in inner_keys:
        inverted_dict[inner_key] = {}
        for outer_key, inner_dict in opt_state_dict.items():
            if inner_key in inner_dict:
                inverted_dict[inner_key][outer_key] = inner_dict[inner_key]

    return inverted_dict
