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

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

from collections import defaultdict

import torch


def split_ffn_experts(
    sigma_moe_params_dict: dict[str, torch.Tensor],
    n_experts: int,
) -> dict[str, torch.Tensor]:
    """Split SigmaMoE expert parameters into individual expert parameters.

    This function takes SigmaMoE parameters where the expert dimension is the first
    dimension and splits them into individual expert parameters with ModuleList-style
    naming.

    Parameters
    ----------
    sigma_moe_params_dict: dict[str, torch.Tensor]
        Dictionary of SigmaMoE parameters where expert dimension is first.
        Expected parameters are: up_projection, down_projection, expert_selection
    n_experts: int
        Number of experts in the SigmaMoE model.

    Returns
    -------
    dict[str, torch.Tensor]
        A new dictionary with individual expert parameters.
        Parameter names are transformed from:
            "prefix.ffn.param_name"
        to:
            "prefix.ffn.experts.{expert_idx}.param_name"

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

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

        # Split along the expert dimension (first dimension)
        for expert_idx in range(n_experts):
            if param_tensor.ndim == 0:
                # Scalar parameters, no expert dimension
                # NOTE: We assume that such parameters are shared across all experts,
                # so we keep them as is.
                expert_param = param_tensor
            else:
                expert_param = param_tensor[expert_idx]
            # Create ModuleList-style naming: e.g., "ffn.experts.0.up_projection"
            expert_param_name = param_name.replace(
                f"ffn.{base_name}",
                f"ffn.experts.{expert_idx}.{base_name}",
            )
            split_params[expert_param_name] = expert_param

    return split_params


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

    This function takes individual expert parameters with ModuleList-style naming
    and recombines them back into the original SigmaMoE format where the expert
    dimension is the first dimension.

    Parameters
    ----------
    split_params_dict: dict[str, torch.Tensor]
        Dictionary of individual expert parameters.
            Expected naming pattern: "prefix.ffn.experts.{expert_idx}.param_name"
            Example: {"ffn.experts.0.up_projection": torch.Tensor([...]),
                    "ffn.experts.1.up_projection": torch.Tensor([...])}
    n_experts: int
        Number of experts in the SigmaMoE model.

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

    Raises
    ------
    ValueError
        If any expert index 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 = {}

    # Group parameters by their base name
    param_groups: dict[str, dict[int, torch.Tensor]] = defaultdict(dict)
    for param_name, param_tensor in split_params_dict.items():
        # Extract base parameter name,
        # e.g., "up_projection" from "ffn.experts.0.up_projection")
        if "ffn.experts." in param_name:
            parts = param_name.split("ffn.experts.")
            prefix = parts[0]
            expert_and_param = parts[1]  # e.g., "0.up_projection"
            expert_idx_str, base_param = expert_and_param.split(".", 1)
            expert_idx = int(expert_idx_str)

            # Reconstruct original parameter name
            original_param_name = f"{prefix}ffn.{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] = param_tensor

    # Recombine parameters by stacking along expert dimension
    for original_param_name, expert_tensors in param_groups.items():
        # Stack tensors in expert order
        stacked_tensors = []
        for expert_idx in range(n_experts):
            if expert_idx in expert_tensors:
                stacked_tensors.append(expert_tensors[expert_idx])
            else:
                msg = f"Missing expert {expert_idx} for parameter {original_param_name}"
                raise ValueError(msg)

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

    return recombined_params


def get_expert_parameters(
    split_params_dict: dict[str, torch.Tensor],
    expert_idx: int,
) -> dict[str, torch.Tensor]:
    """Extract parameters for a specific expert from split parameters.

    Parameters
    ----------
    split_params_dict: dict[str, torch.Tensor]
        Dictionary of split expert parameters.
        Expected naming pattern: "prefix.ffn.experts.{expert_idx}.param_name"
        Example: {"ffn.experts.0.up_projection": torch.Tensor([...]),
                  "ffn.experts.1.up_projection": torch.Tensor([...])}
    expert_idx: int
        Index of the expert to extract parameters for.

    Returns
    -------
    dict[str, torch.Tensor]
        A new dictionary containing only the parameters for the specified expert.
        Parameter names are transformed from:
            "prefix.ffn.experts.{expert_idx}.param_name"
        to:
            "prefix.ffn.param_name"

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


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

    Parameters
    ----------
    original_params: dict[str, torch.Tensor]
        Dictionary of original SigmaMoE parameters.
        Expected parameters are: up_projection, down_projection, expert_selection
        These should have the expert dimension as the first dimension.
        Example: {"ffn.up_projection": torch.Tensor([n_experts, ...])
        where n_experts is the number of experts.
    split_params: dict[str, torch.Tensor]
        Dictionary of split expert parameters.
        Expected naming pattern: "prefix.ffn.experts.{expert_idx}.param_name"
        Example: {"ffn.experts.0.up_projection": torch.Tensor([...]),
                  "ffn.experts.1.up_projection": torch.Tensor([...])}
    n_experts: int
        Number of experts in the SigmaMoE model.
    atol: float
        Absolute tolerance for comparing tensors.

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

    """
    try:
        recombined = recombine_ffn_experts(split_params, n_experts)

        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:
        # If no exceptions, the operations are consistent
        return True
