"""Utilities for splitting and merging experts module tensors."""

from logging import DEBUG
from typing import Any

import numpy as np
import torch
from flwr.common import NDArrays, log

from mt_dao.masks_utils import ModelStateNames
from mt_dao.models.ffns.sigma_moe import LIST_EXPERTS_PARAMS as FFNS_LIST_EXPERTS_PARAMS
from mt_dao.moe_utils.sigma_moe_expert_utils import (
    recombine_ffn_experts,
    split_ffn_experts,
)
from mt_dao.pytorch_utils.utils import invert_opt_state_dict_keys


def transform_torch_checkpoint_to_mt_dao_model_states(
    model_states_dict: dict[str, Any],
    layer_ordering: list[str],
    n_experts: int | None = None,
    *,
    preserve_optimizer_states: bool = True,
) -> NDArrays:
    """Transform a PyTorch checkpoint to mt_dao model states.

    This function extracts model parameters and optimizer states from a PyTorch
    checkpoint and returns them as a list of NumPy arrays. It handles both the model
    parameters and optimizer states, including special handling for MPT models with
    multiple experts.

    Parameters
    ----------
    model_states_dict : dict[str, Any]
        The model states dictionary containing the model and optimizer states.
    n_experts : int | None, optional
        The number of experts in the model, if applicable. If provided, the function
        will split the model parameters and optimizer states accordingly. Defaults to
        None.
    preserve_optimizer_states : bool
        Whether to preserve the initial optimizer states from the mode, by default True.
    layer_ordering : list[str]
        Defines the ordering in which the parameters should be returned.

    Returns
    -------
    NDArrays
        A list of NumPy arrays representing the model parameters and optimizer states.
        The arrays are detached from the computation graph and moved to the CPU.
        If `n_experts` is specified, the function will split the parameters and
        optimizer states accordingly.

    """
    # Initialize the model states
    model_states: dict[ModelStateNames, NDArrays] = {}
    tensors_dict: dict[str, torch.Tensor] | dict[str, torch.nn.Parameter] = (
        model_states_dict["state"]["model"]
    )
    if n_experts:
        # For MPT models where we need to split the experts
        tensors_dict = split_experts_layers_tensors(
            tensors_dict=tensors_dict,
            n_experts=n_experts,
        )
    #  TODO(anon): Handle prefix manipulation globally.
    tensors_dict = {k.replace("model.", "", 1): v for k, v in tensors_dict.items()}
    log(DEBUG, f"Model parameters found: {list(tensors_dict.keys())}")
    log(DEBUG, f"Layer ordering provided: {layer_ordering}")
    tensors_dict = {k: tensors_dict[k] for k in layer_ordering}
    model_states[ModelStateNames.PARAMETERS] = [
        val.detach().to("cpu").numpy() for val in tensors_dict.values()
    ]
    optimizers_dict: dict[str, Any] = model_states_dict["state"].get(
        "optimizers",
        {},
    )
    if optimizers_dict:
        # If optimizers are present, we assume the first one is the one we want
        optimizer = next(iter(optimizers_dict.values()))
        opt_state: dict[str, dict[str, torch.Tensor]] = optimizer.get(
            "state",
            {},
        )
        if n_experts:
            # For MPT models where we need to split the experts
            opt_state = split_experts_opt_dict_tensors(
                opt_state_dict=opt_state,
                n_experts=n_experts,
            )
        #  TODO(anon): Handle prefix manipulation globally.
        opt_state = {k.replace("model.", "", 1): v for k, v in opt_state.items()}
        opt_state = {k: opt_state[k] for k in layer_ordering}
        for state in opt_state.values():
            for state_type_str, opt_state_tensor in state.items():
                if state_type_str.upper() in [e.value for e in ModelStateNames]:
                    state_type = ModelStateNames[state_type_str.upper()]
                    # Add the optimizer state tensor to the model states
                    if state_type not in model_states:
                        model_states[state_type] = []
                    opt_state_numpy = opt_state_tensor.detach().to("cpu").numpy()
                    if preserve_optimizer_states:
                        model_states[state_type].append(
                            opt_state_numpy,
                        )
                    else:
                        model_states[state_type].append(
                            np.zeros_like(opt_state_numpy),
                        )
    model_states_list: NDArrays = []
    for model_state in model_states.values():
        model_states_list.extend(model_state)
    return model_states_list


def split_experts_opt_dict_tensors(
    opt_state_dict: dict[str, dict[str, torch.Tensor]],
    n_experts: int,
) -> dict[str, dict[str, torch.Tensor]]:
    """Split experts optimizer state tensors into individual expert parameters.

    Parameters
    ----------
    opt_state_dict : dict[str, dict[str, torch.Tensor]]
        Optimizer state dictionary with structure {layer_name: {state_name: tensor}}
    n_experts : int
        Number of experts for FFN

    Returns
    -------
    dict[str, dict[str, torch.Tensor]]
        Optimizer state dictionary with split expert parameters.

    """
    # Separate the optimizer state dict into a state_name and expert tensors.
    inv_opt_state_dict = invert_opt_state_dict_keys(opt_state_dict)

    # Apply split operation to each state group
    split_inv_opt_state_dict = {}
    for state_name, layer_tensors in inv_opt_state_dict.items():
        split_layer_tensors = split_experts_layers_tensors(
            tensors_dict=layer_tensors,
            n_experts=n_experts,
        )
        split_inv_opt_state_dict[state_name] = split_layer_tensors

    # Recompose the original structure with the split tensors.
    return invert_opt_state_dict_keys(split_inv_opt_state_dict)


def merge_experts_opt_dict_tensors(
    opt_state_dict: dict[str, dict[str, torch.Tensor]],
    n_experts: int,
) -> dict[str, dict[str, torch.Tensor]]:
    """Merge split experts optimizer state tensors back into original format.

    Parameters
    ----------
    opt_state_dict : dict[str, dict[str, torch.Tensor]]
        Optimizer state dictionary with split expert parameters
    n_experts : int
        Number of experts for FFN

    Returns
    -------
    dict[str, dict[str, torch.Tensor]]
        Optimizer state dictionary with merged parameters in original format.

    """
    # Separate the optimizer state dict into a state_name and expert tensors.
    inv_opt_state_dict = invert_opt_state_dict_keys(opt_state_dict)

    # Apply merge operation to each state group
    merged_inv_opt_state_dict = {}
    for state_name, layer_tensors in inv_opt_state_dict.items():
        merged_layer_tensors = merge_experts_layers_tensors(
            layer_tensors,
            n_experts,
        )
        merged_inv_opt_state_dict[state_name] = merged_layer_tensors

    # Recompose the original structure with the recombined tensors.
    return invert_opt_state_dict_keys(merged_inv_opt_state_dict)


def split_experts_layers_tensors(
    tensors_dict: dict[str, torch.nn.Parameter] | dict[str, torch.Tensor],
    n_experts: int,
) -> dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]:
    """Split expert modules parameters into individual expert parameters.

    Parameters
    ----------
    tensors_dict : dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]
        Dictionary of model parameters
    n_experts : int
        Number of experts for FFN

    Returns
    -------
    dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]
        Dictionary with split expert parameters.

    """
    # Make a copy to avoid modifying the original dict
    result_dict = dict(tensors_dict)

    # Take out the SigmaMoE layers from the tensors dict
    sigma_moe_params = {
        k: v
        for k, v in tensors_dict.items()
        # NOTE: The boolean condition was `param in k` before. However, that was also
        # capturing non-expert FFNs since their signature was
        # `<pattern_in_param>.weight` which contains `<pattern_in_param>` contained in
        # `FFNS_LIST_EXPERTS_PARAMS`. Using
        # `any(k.endswith(param) for param in FFNS_LIST_EXPERTS_PARAMS)` to filter fixes
        # this.
        if any(k.endswith(param) for param in FFNS_LIST_EXPERTS_PARAMS)
    }

    # Remove original SigmaMoE and SwitchHead params from result
    for key in sigma_moe_params:
        result_dict.pop(key, None)

    # Split SigmaMoE layers
    if sigma_moe_params:
        split_sigma_moe_params = split_ffn_experts(sigma_moe_params, n_experts)
        result_dict |= split_sigma_moe_params

    return result_dict


def merge_experts_layers_tensors(
    tensors_dict: dict[str, torch.nn.Parameter] | dict[str, torch.Tensor],
    n_experts: int,
) -> dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]:
    """Merge split expert modules parameters back into original format.

    Parameters
    ----------
    tensors_dict : dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]
        Dictionary with split expert parameters
    n_experts : int
        Number of experts for FFN

    Returns
    -------
    dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]
        Dictionary with merged parameters in original format.

    """
    # Make a copy to avoid modifying the original dict
    result_dict = dict(tensors_dict)

    # Take out the split FFN expert layers from the tensors dict
    split_ffn_params = {k: v for k, v in tensors_dict.items() if "ffn.experts." in k}

    # Remove split expert params from result
    for key in split_ffn_params:
        result_dict.pop(key, None)

    # Merge FFN layers
    if split_ffn_params:
        merged_ffn_params = recombine_ffn_experts(
            split_ffn_params,
            n_experts,
        )
        result_dict |= merged_ffn_params

    return result_dict


def verify_experts_layers_split_merge_consistency(
    original_tensors: dict[str, torch.nn.Parameter] | dict[str, torch.Tensor],
    split_tensors: dict[str, torch.nn.Parameter] | dict[str, torch.Tensor],
    n_experts: int,
    *,
    atol: float = 1e-6,
) -> bool:
    """Verify that split and merge operations are consistent for experts layers.

    Parameters
    ----------
    original_tensors : dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]
        Dictionary of original experts parameters
    split_tensors : dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]
        Dictionary of split expert parameters
    n_experts : int
        Number of experts for FFN
    atol : float
        Absolute tolerance for comparing tensors

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

    """
    try:
        merged = merge_experts_layers_tensors(
            split_tensors,
            n_experts,
        )

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

    except (ValueError, KeyError):
        return False
    else:
        return True


def verify_experts_opt_dict_split_merge_consistency(
    original_opt_state_dict: dict[str, dict[str, torch.Tensor]],
    split_opt_state_dict: dict[str, dict[str, torch.Tensor]],
    n_experts: int,
    *,
    atol: float = 1e-6,
) -> bool:
    """Verify that split and merge operations are correct for experts optimizer state.

    Parameters
    ----------
    original_opt_state_dict : dict[str, dict[str, torch.Tensor]]
        Dictionary of original optimizer state with structure
        {layer_name: {state_name: tensor}}
    split_opt_state_dict : dict[str, dict[str, torch.Tensor]]
        Dictionary of split optimizer state with expert parameters
    n_experts : int
        Number of experts for FFN
    atol : float
        Absolute tolerance for comparing tensors

    Returns
    -------
    bool
        True if merged optimizer state matches original, False otherwise

    """
    try:
        merged = merge_experts_opt_dict_tensors(
            split_opt_state_dict,
            n_experts,
        )

        for layer_name, layer_states in original_opt_state_dict.items():
            if layer_name not in merged:
                return False
            for state_name, state_tensor in layer_states.items():
                if state_name not in merged[layer_name]:
                    return False
                if not torch.allclose(
                    state_tensor,
                    merged[layer_name][state_name],
                    atol=atol,
                ):
                    return False

    except (ValueError, KeyError):
        return False
    else:
        return True
