"""Utilities for splitting and recombining SwitchHead expert parameters.

This module provides functions to split SwitchHead expert-related weights into
individual experts per attention head as if they were in a ModuleList, and to recombine
them back to the original format.
"""

from collections import defaultdict

import torch


def split_switch_head_experts(
    switch_head_params_dict: dict[str, torch.Tensor],
    n_experts: int,
    n_heads: int,
) -> dict[str, torch.Tensor]:
    """Split SwitchHead expert parameters into individual experts per attention head.

    SwitchHead parameters have the expert dimension encapsulated into the first
    dimension and multiplied by the number of heads. This function splits them into
    individual expert parameters with ModuleList-style naming.

    Parameters
    ----------
    switch_head_params_dict: dict[str, torch.Tensor]
        dictionary of SwitchHead parameters where the first dimension
        contains both expert and head information (n_experts * n_heads)
    n_experts: int
        Number of experts
    n_heads: int
        Number of attention heads

    Returns
    -------
    dict[str, torch.Tensor]
        A new dictionary with individual expert parameters named as if they were in a
        ModuleList.
        Parameter names are transformed from:
            "prefix.switch_head.param_name"
        to:
            "prefix.switch_head.experts.{expert_idx}.heads.{head_idx}.param_name"

    """
    if n_experts == 1:
        # If there's only one expert, return the parameters as is
        return switch_head_params_dict
    split_params = {}

    for param_name, param_tensor in switch_head_params_dict.items():
        # Extract the base parameter name (remove the full path prefix)
        base_name = param_name.split("switch_head.")[-1]

        for expert_idx in range(n_experts):
            for head_idx in range(n_heads):
                if param_tensor.ndim == 0:
                    # Scalar parameters, no expert or head dimension
                    # NOTE: We assume that such parameters are shared across all
                    # experts and heads, so we keep them as is.
                    expert_head_param = param_tensor
                else:
                    expert_head_param = param_tensor[expert_idx, head_idx]
                expert_param_name = param_name.replace(
                    f"switch_head.{base_name}",
                    f"switch_head.experts.{expert_idx}.heads.{head_idx}.{base_name}",
                )
                split_params[expert_param_name] = expert_head_param

    return split_params


def recombine_switch_head_experts(
    split_params_dict: dict[str, torch.Tensor],
    n_experts: int,
    n_heads: int,
) -> dict[str, torch.Tensor]:
    """Recombine individual expert parameters back into the original SwitchHead format.

    This function takes individual expert parameters with ModuleList-style naming
    and recombines them back into the original SwitchHead format.

    Parameters
    ----------
    split_params_dict: dict[str, torch.Tensor]
        Dictionary with individual expert parameters. Expected naming pattern:
        "prefix.switch_head.experts.{expert_idx}.heads.{head_idx}.param_name"
        where:
            - prefix is the model prefix (e.g., "model.")
            - expert_idx is the index of the expert (0 to n_experts-1)
            - head_idx is the index of the attention head (0 to n_heads-1)
            - param_name is the base parameter name (e.g., "v", "o", "sel_v", "sel_o")
    n_experts: int
        Number of experts
    n_heads: int
        Number of attention heads

    Returns
    -------
    dict[str, torch.Tensor]
        A new dictionary with recombined parameters in original SwitchHead format.
        Parameter names are transformed from:
            "prefix.switch_head.experts.{expert_idx}.heads.{head_idx}.param_name"
        back to:
            "prefix.switch_head.param_name"

    Raises
    ------
    ValueError: If any expert/head combination is missing for a parameter.

    """
    if n_experts == 1:
        # If there's only one expert, return the parameters as is
        return split_params_dict
    recombined_params: dict[str, torch.Tensor] = defaultdict(torch.Tensor)

    # Group parameters by their base name
    param_groups: dict[str, dict[tuple[int, int], torch.Tensor]] = defaultdict(dict)
    for param_name, param_tensor in split_params_dict.items():
        # Extract base parameter name
        if "switch_head.experts." in param_name:
            parts = param_name.split("switch_head.experts.")
            prefix = parts[0]
            expert_head_and_param = parts[1]  # e.g., "0.heads.1.v"

            # Parse expert_idx, head_idx, and base_param
            expert_parts = expert_head_and_param.split(".")
            expert_idx = int(expert_parts[0])
            # Skip "heads" keyword
            head_idx = int(expert_parts[2])
            base_param = ".".join(expert_parts[3:])  # Handle nested parameter names

            # Reconstruct original parameter name
            original_param_name = f"{prefix}switch_head.{base_param}"

            if original_param_name not in param_groups:
                param_groups[original_param_name] = defaultdict(torch.Tensor)
            param_groups[original_param_name][expert_idx, head_idx] = param_tensor

    # Recombine parameters
    for original_param_name, expert_head_tensors in param_groups.items():
        flattened_experts_tensors = []
        for expert_idx in range(n_experts):
            flattened_attn_heads_tensors = []
            for head_idx in range(n_heads):
                key = (expert_idx, head_idx)
                if key in expert_head_tensors:
                    flattened_attn_heads_tensors.append(expert_head_tensors[key])
                else:
                    msg = (
                        f"Missing expert {expert_idx}, head {head_idx} for"
                        f" parameter {original_param_name}"
                    )
                    raise ValueError(msg)
            flattened_experts_tensors.append(
                (
                    flattened_attn_heads_tensors[0]
                    if len(flattened_attn_heads_tensors) == 1
                    else torch.stack(flattened_attn_heads_tensors, dim=0)
                ),
            )

        # Stack along first dimension (flattened expert * head dimension)
        recombined_params[original_param_name] = (
            flattened_experts_tensors[0]
            if len(flattened_experts_tensors) == 1
            else torch.stack(
                flattened_experts_tensors,
                dim=0,
            )
        )

    return recombined_params


def get_expert_head_parameters(
    split_params_dict: dict[str, torch.Tensor],
    expert_idx: int,
    head_idx: int,
) -> dict[str, torch.Tensor]:
    """Extract parameters given a expert and head combination from split parameters.

    Parameters
    ----------
    split_params_dict: dict[str, torch.Tensor]
        A dictionary with individual expert parameters
        Expected naming pattern:
            "prefix.switch_head.experts.{expert_idx}.heads.{head_idx}.param_name"
        where:
            - prefix is the model prefix (e.g., "model.")
            - expert_idx is the index of the expert (0 to n_experts-1)
            - head_idx is the index of the attention head (0 to n_heads-1)
            - param_name is the base parameter name (e.g., "v", "o", "sel_v", "sel_o")
    expert_idx: int
        Index of the expert to extract parameters for
    head_idx: int
        Index of the attention head to extract parameters for

    Returns
    -------
    dict[str, torch.Tensor]
        A new dictionary with parameters for the specified expert and head.
        The keys are the parameter names, and the values are the corresponding tensors.

    """
    expert_head_pattern = f".experts.{expert_idx}.heads.{head_idx}."
    return {
        param_name: param_tensor
        for param_name, param_tensor in split_params_dict.items()
        if expert_head_pattern in param_name
    }


def verify_split_recombine_consistency(
    original_params: dict[str, torch.Tensor],
    split_params: dict[str, torch.Tensor],
    n_experts: int,
    n_heads: int,
    atol: float = 1e-6,
) -> bool:
    """Verify that split+recombine operations are consistent for SwitchHead parameters.

    Parameters
    ----------
    original_params: dict[str, torch.Tensor]
        Dictionary of original SwitchHead parameters.
        The keys are the parameter names, and the values are the corresponding tensors.
    split_params: dict[str, torch.Tensor]
        Dictionary of split SwitchHead parameters.
        The keys are the parameter names, and the values are the corresponding tensors.
    n_experts: int
        Number of experts.
    n_heads: int
        Number of attention heads.
    atol: float
        Absolute tolerance for comparison.

    Returns
    -------
    bool
        True if recombined parameters match original parameters, False otherwise

    """
    try:
        recombined = recombine_switch_head_experts(split_params, n_experts, n_heads)

        for param_name, param in original_params.items():
            if param_name not in recombined:
                return False
            if not torch.allclose(
                param,
                recombined[param_name],
                atol=atol,
            ):
                return False

    except (ValueError, KeyError):
        return False
    else:
        # This else block is not necessary, but included for clarity
        return True
