"""Utility functions shared across getters and setters for ComposerTrainer."""

import copy
from concurrent.futures import ThreadPoolExecutor
from logging import DEBUG

import numpy as np
from composer import Trainer
from flwr.common import NDArrays, Scalar, log

from mt_dao.clients.configs import FitConfig
from mt_dao.conf.base_schema import DeltaModes
from mt_dao.masks_utils import ModelStateNames
from mt_dao.shm.utils import compress_with_strict
from mt_dao.utils import sum_of_squares, sum_of_squares_per_layer


def report_l2_norms(  # noqa: PLR0913, PLR0917
    layer_names: list[str],
    layer_deltas: dict[str, float],
    transmission_state_sum_of_squares: dict[str, float],
    aggregation_state_sum_of_squares: dict[str, float],
    train_metrics: dict[str, Scalar],
    metric_keys: tuple[str, str, str],
) -> None:
    """Report the L2 norms for the given layers, deltas and states.

    Parameters
    ----------
    layer_names : list[str]
        The layer names to report
    layer_deltas : dict[str, float]
        The delta per layer
    transmission_state_sum_of_squares : dict[str, float]
       The sum of squares per layer for the transmission state.
    aggregation_state_sum_of_squares : dict[str, float]
        The sum of squares per layer for the aggregation state.
    train_metrics : dict[str, Scalar]
        A dictionary to store training metrics.
    metric_keys : tuple[str, str, str]
        A tuple containing:
        - pre_key: The key for the tranmission state metrics
        - post_key: The key for the aggregation state metrics
        - delta_key: The key for the delta metrics

    """
    pre_key, post_key, delta_key = metric_keys
    total_delta = 0.0
    total_baseline_sum_of_squares = 0.0
    total_returned_sum_of_squares = 0.0
    for layer_name in layer_names:
        total_delta += layer_deltas[layer_name]
        total_baseline_sum_of_squares += transmission_state_sum_of_squares[layer_name]
        total_returned_sum_of_squares += aggregation_state_sum_of_squares[layer_name]

    train_metrics[pre_key] = float(
        (total_baseline_sum_of_squares),
    )
    train_metrics[post_key] = float(
        (total_returned_sum_of_squares),
    )
    train_metrics[delta_key] = float(
        np.sqrt(total_delta),
    )


def extract_l2_norm_for_model_state(
    transmission_state: tuple[NDArrays, list[str]],
    aggregation_state: tuple[NDArrays, list[str]],
    layers_for_calculation: set[str],
    train_metrics: dict[str, Scalar],
    pseudo_gradient: str,
) -> dict[str, float]:
    """Extract the L2 norm for a given subset layers.

     This function retrieves the L2 norm against the model parameters after training
     and the baseline server representation and determines the L2 norm delta metrics
     based on the specific layers chosen for calculation.

    Parameters
    ----------
    transmission_state: tuple[NDArrays, list[str]],
         A tuple containing the layers and the corresponding
         layer names of the server.
    aggregation_state: tuple[NDArrays, list[str]],
         A tuple containing the layers and the corresponding
         layer names of the clients.
    layers_for_calculation: set[str],
         A set of all the layers that should be used for the delta calculation
    train_metrics: dict[str, Scalar],
         A dictionary to store training metrics.
    pseudo_gradient: str,
         The the key used to define the metric per layer.

    Returns
    -------
    dict[str, float]
         A dictionary containing the L2 norm per layer.

    """
    if not transmission_state or not aggregation_state:
        log(DEBUG, "No model state to compute metrics")
        return None

    # NOTE: We do not need to have the situation where the transmission and
    # the aggregation are of the same length as the layers for calculation
    # will determien the specific layers we need only.
    def calculate_layer_delta(
        args: tuple[str, NDArrays, list[str], NDArrays, list[str]],
    ) -> tuple[str, float, float]:
        """Calculate the layer delta for a single layer.

        Parameters
        ----------
        args : tuple[str, NDArrays, list[str], NDArrays, list[str]]
            A tuple containing:
            - The name of the delta layer as a str.
            - The transmission layers
            - The tranmission layer names
            - The aggregation layers
            - The aggregation layer names

        Returns
        -------
        tuple[str, float, str, float]
            A tuple containing:
            - The name of the delta layer as a str.
            - The sum of squares for the delta layer.
            - The square root of the sum of squares.

        """
        (
            delta_layer,
            transmission_ndarrays,
            transmission_layer_names,
            aggregation_ndarrays,
            aggregation_layer_names,
        ) = args

        base = transmission_ndarrays[transmission_layer_names.index(delta_layer)]
        curr = aggregation_ndarrays[aggregation_layer_names.index(delta_layer)]
        diff = base - curr
        layer_delta_sum_of_squares = sum_of_squares([diff])

        metric_value = float(np.sqrt(layer_delta_sum_of_squares))

        return delta_layer, layer_delta_sum_of_squares, metric_value

    transmission_ndarrays, transmission_layer_names = transmission_state
    aggregation_ndarrays, aggregation_layer_names = aggregation_state
    layer_deltas: dict[str, float] = {}

    # Prepare arguments for each layer calculation
    args_list = [
        (
            delta_layer,
            transmission_ndarrays,
            transmission_layer_names,
            aggregation_ndarrays,
            aggregation_layer_names,
        )
        for delta_layer in layers_for_calculation
    ]

    # Use ThreadPoolExecutor to parallelize the calculations
    with ThreadPoolExecutor() as executor:
        results = list(executor.map(calculate_layer_delta, args_list))

    # Process results
    for delta_layer, layer_delta_sum_of_squares, metric_value in results:
        layer_deltas[delta_layer] = layer_delta_sum_of_squares
        train_metrics[f"{pseudo_gradient}/{delta_layer}"] = metric_value

    return layer_deltas


def determine_layers_for_statistics(
    delta_mode: DeltaModes,
    aggregation_layer_names: list[str],
    transmission_layer_names: list[str],
) -> list[str]:
    """Determine which layers will be sent for L2 norm delta metrics.

     This function retrieves the model parameters after training and the baseline server
     representation and determines which model layers will be used for L2 norm delta
     metrics based on the selection of a certain mode.

    Parameters
    ----------
    delta_mode: DeltaModes
         The mode for calculating the delta statistic.
    aggregation_layer_names: tuple[NDArrays, list[str]],
         A tuple containing the layers and the corresponding
         layer names of the clients.
    transmission_layer_names: tuple[NDArrays, list[str]],
         A tuple containing the layers and the corresponding
         layer names of the server.


    Returns
    -------
    list[str]
         The list of reference layers that were used to determine which layers
         should be passed to the L2 norm calculation.

    """
    if delta_mode == DeltaModes.ALL_PARAMS:
        log(DEBUG, "Calculating full delta.")
        reference_layers = transmission_layer_names
    elif delta_mode == DeltaModes.TRANSMISSION:
        log(DEBUG, "Transmission determined delta mode. Filtering received.")
        reference_layers = transmission_layer_names
    elif delta_mode == DeltaModes.AGGREGATION:
        log(DEBUG, "Aggregated determined delta mode. Filtering transitted.")
        reference_layers = aggregation_layer_names
    elif delta_mode == DeltaModes.INTERSECTION:
        log(DEBUG, "Intersection delta mode. Calculating.")
        reference_layers = list(
            set(transmission_layer_names) & set(aggregation_layer_names),
        )
    return reference_layers


def filter_expert_parameters(
    mask_transmission: tuple[bool, ...],
    layer_names_transmission: list[str],
    layer_types_transmission: list[ModelStateNames],
    experts_ids: list[int],
    *,
    change_names: bool = True,
) -> tuple[tuple[bool, ...], list[str], list[ModelStateNames]]:
    """Filter expert parameters based on allowed expert IDs.

    Parameters
    ----------
    mask_transmission : tuple[bool, ...]
        Boolean mask for parameter transmission
    layer_names_transmission : list[str]
        List of layer names
    layer_types_transmission : list[ModelStateNames]
        List of layer types
    experts_ids : list[int]
        List of allowed expert IDs
    change_names : bool
        Whether to change the names of the expert parameters to reflect the new indices.
        If False, the names will remain unchanged, but the mask will still be filtered.

    Returns
    -------
    tuple[tuple[bool, ...], list[str], list[ModelStateNames]]
        Filtered mask, layer names, and layer types

    """
    # Create a mapping from old expert indices to new expert indices
    expert_id_mapping = {
        old_idx: new_idx for new_idx, old_idx in enumerate(experts_ids)
    }

    # Filter mask, layer names, and layer types for allowed experts
    filtered_mask = list(copy.deepcopy(mask_transmission))
    filtered_layer_names = []
    filtered_layer_types = []

    # Get the true indices of the original mask
    true_indices = [i for i, val in enumerate(mask_transmission) if val]

    # Loop over the tuple elements
    for i, (layer_name, layer_type) in enumerate(
        zip(
            layer_names_transmission,
            layer_types_transmission,
            strict=True,
        ),
    ):
        should_include = True
        updated_layer_name = layer_name

        # Check if this is a FFN expert parameter
        if "ffn" in layer_name and "experts" in layer_name:
            should_include, updated_layer_name = process_sigma_moe_expert(
                layer_name=layer_name,
                expert_id_mapping=expert_id_mapping,
            )

        # NOTE: We skip SwitchHead expert parameters for the time being, but the same
        # logic can be easily implemented for them in the future.

        # Include the parameter if it should be included
        if should_include:
            filtered_layer_names.append(
                updated_layer_name if change_names else layer_name,
            )
            filtered_layer_types.append(layer_type)
            # NOTE: We don't touch the mask here as it should be correct for this index
        else:
            # NOTE: Since the boolean mask is in full unlike the layer names and types,
            # we need to find the i-th True element of the original mask and make it
            # False in the filtered one
            filtered_mask[true_indices[i]] = False

    return tuple(filtered_mask), filtered_layer_names, filtered_layer_types


def process_sigma_moe_expert(
    layer_name: str,
    expert_id_mapping: dict[int, int],
) -> tuple[bool, str]:
    """Process SigmaMoE expert parameter name and check if it should be included.

    Parameters
    ----------
    layer_name : str
        Original layer name
    expert_id_mapping : dict[int, int]
        Mapping from old expert IDs to new expert IDs
    min_parts_length : int
        Minimum required parts length for processing

    Returns
    -------
    tuple[bool, str]
        (should_include, updated_layer_name)

    """
    # Check the initial condition is satisfied
    # e.g., w/ expert: "model.layers.0.mlp.ffn.experts.5.up_projection"
    # e.g., w/o expert: "model.layers.0.mlp.ffn.up_projection"
    assert (
        "ffn" in layer_name
    ), f"This function is for FFN only, not {layer_name}"
    assert (
        "experts" in layer_name
    ), f"This function is for experts only, not {layer_name}"
    # Split the layer name to extract expert information
    parts = layer_name.split("ffn.experts.")
    expert_info = parts[1].split(".", 1)
    assert (
        len(expert_info) > 0
    ), f"Expected at least one part after 'ffn.experts.' in {layer_name}"

    # Get the expert index from the name
    expert_idx = int(expert_info[0])
    # Only include if this expert is in the allowed list
    if expert_idx in expert_id_mapping:
        # Remap the expert index to the new range [0, len(experts_ids)-1]
        new_expert_idx = expert_id_mapping[expert_idx]
        if len(expert_info) > 1:
            updated_layer_name = (
                f"{parts[0]}ffn.experts.{new_expert_idx}.{expert_info[1]}"
            )
        else:
            updated_layer_name = f"{parts[0]}ffn.experts.{new_expert_idx}"
        return True, updated_layer_name
    # If we successfully parsed the expert index, but we have not returned yet it
    # means that we have to exclude it. So, we return the name as it is and it won't
    # be appended.
    return False, layer_name


def remap_expert_parameters_to_original_indices(
    layer_names: list[str],
    experts_ids: list[int],
) -> list[str]:
    """Remap expert parameter names back to their original indices.

    This function is the inverse of the remapping done in filter_expert_parameters.
    It takes layer names that have been remapped to the range [0, len(experts_ids)-1]
    and maps them back to their original expert indices.

    Parameters
    ----------
    layer_names : list[str]
        List of layer names with remapped expert indices
    experts_ids : list[int]
        List of original expert IDs used for the remapping

    Returns
    -------
    list[str]
        List of layer names with original expert indices restored

    """
    # Create a mapping from new expert indices back to original expert indices
    new_to_original_mapping = dict(enumerate(experts_ids))

    remapped_layer_names = []
    for layer_name in layer_names:
        # Check if this is a FFN expert parameter
        if "ffn" in layer_name and "experts" in layer_name:
            remapped_name = remap_sigma_moe_expert_to_original(
                layer_name=layer_name,
                new_to_original_mapping=new_to_original_mapping,
            )
            remapped_layer_names.append(remapped_name)
        else:
            # For non-expert parameters, keep the name as is
            remapped_layer_names.append(layer_name)

    return remapped_layer_names


def remap_sigma_moe_expert_to_original(
    layer_name: str,
    new_to_original_mapping: dict[int, int],
) -> str:
    """Remap a SigmaMoE expert parameter name back to its original index.

    Parameters
    ----------
    layer_name : str
        Layer name with remapped expert index
    new_to_original_mapping : dict[int, int]
        Mapping from new expert indices to original expert indices

    Returns
    -------
    str
        Layer name with original expert index restored

    """
    # Check the initial condition is satisfied
    assert (
        "ffn" in layer_name
    ), f"This function is for FFN only, not {layer_name}"
    assert (
        "experts" in layer_name
    ), f"This function is for experts only, not {layer_name}"

    # Split the layer name to extract expert information
    parts = layer_name.split("ffn.experts.")
    expert_info = parts[1].split(".", 1)
    assert (
        len(expert_info) > 0
    ), f"Expected at least one part after 'ffn.experts.' in {layer_name}"

    try:
        new_expert_idx = int(expert_info[0])
    except (ValueError, IndexError):
        # If we can't parse the expert index, return the original name
        return layer_name
    else:
        # Map back to the original expert index
        if new_expert_idx in new_to_original_mapping:
            original_expert_idx = new_to_original_mapping[new_expert_idx]
            if len(expert_info) > 1:
                return (
                    f"{parts[0]}ffn.experts."
                    f"{original_expert_idx}.{expert_info[1]}"
                )
            return f"{parts[0]}ffn.experts.{original_expert_idx}"
        # If the mapping is not found, return the original name
        return layer_name


def calculate_l2_norms_per_state_and_mode(
    aggregation_state: tuple[NDArrays, list[str]],
    transmission_state: tuple[NDArrays, list[str]],
    delta_modes: list[DeltaModes],
    train_metrics: dict[str, Scalar],
    metric_keys: tuple[str, str, str],
) -> None:
    """Calculate the L2 norms according to the given states and for the given modes.

    Parameters
    ----------
    aggregation_state : tuple[NDArrays, list[str]]
        The aggregated state received.
    transmission_state : tuple[NDArrays, list[str]]
       The transmitted ste received.
    delta_modes : list[DeltaModes]
        The modes for delta calculations.
    train_metrics : dict[str, Scalar]
        A dictionary to store training metrics.
    metric_keys : tuple[str, str, str]
        A tuple containing:
        - pre_key: The key for the transmission state metrics
        - post_key: The key for the aggregation state metrics
        - delta_key: The key for the delta metrics

    """
    global_model, local_model, pseudo_gradient = metric_keys
    aggregation_state_sum_of_squares = sum_of_squares_per_layer(aggregation_state)
    transmission_state_sum_of_squares = sum_of_squares_per_layer(transmission_state)
    layers_per_mode_type = {
        mode: determine_layers_for_statistics(
            delta_mode=mode,
            transmission_layer_names=transmission_state[1],
            aggregation_layer_names=aggregation_state[1],
        )
        for mode in delta_modes
    }
    all_layers_for_l2_norm_calculation = set().union(
        *[layers for _, layers in layers_per_mode_type.items()],
    )
    layer_deltas = extract_l2_norm_for_model_state(
        transmission_state=transmission_state,
        aggregation_state=aggregation_state,
        layers_for_calculation=all_layers_for_l2_norm_calculation,
        pseudo_gradient=pseudo_gradient,
        train_metrics=train_metrics,
    )
    for mode, layer_names in layers_per_mode_type.items():
        report_l2_norms(
            layer_names=layer_names,
            layer_deltas=layer_deltas,
            transmission_state_sum_of_squares=transmission_state_sum_of_squares,
            aggregation_state_sum_of_squares=aggregation_state_sum_of_squares,
            train_metrics=train_metrics,
            metric_keys=(
                f"{mode.value}/{global_model}",
                f"{mode.value}/{local_model}",
                f"{mode.value}/{pseudo_gradient}",
            ),
        )


def get_ndarrays_and_names_from_payload(
    payload: NDArrays,
    layer_names: list[str],
    layer_types: list[ModelStateNames],
    model_state_type: ModelStateNames,
) -> tuple[NDArrays, list[str]]:
    """Get the NDArrays and names from the payload.

    NDArrays are the model parameters/optimizer states, and names are the layer names.
    The layer types uniquely identify what kind of state a give ndarray is

    Based on the model state type, this function filters the payload and returns

    Parameters
    ----------
    payload : NDArrays
        The payload containing the model parameters.
    layer_names : list[str]
        The names of the layers in the model.
    layer_types : list[ModelStateNames]
        The types of the layers in the model.
    model_state_type : ModelStateNames
        The model state type to filter the payload.

    Returns
    -------
    tuple[NDArrays, list[str]]
        A tuple containing:
        - The filtered NDArrays from the payload.
        - The corresponding layer names.

    """
    selector = [layer_type == model_state_type for layer_type in layer_types]

    return list(compress_with_strict(payload, selector, strict=True)), list(
        compress_with_strict(layer_names, selector, strict=True),
    )


def modify_aggregation_mask_with_frozen_layers(
    client_config: FitConfig,
) -> None:
    """Modify the aggregation mask to include frozen layers.

    This function modifies the aggregation mask in place to include frozen layers
    specified in the client configuration. It updates the mask to ensure that the
    frozen layers are included in the aggregation process.

    Parameters
    ----------
    client_config : FitConfig
        The client configuration containing the aggregation mask and frozen layers.

    """
    if client_config.frozen_layers is not None:
        # NOTE: We need to modify the aggregation mask in place
        # to exclude frozen layers, so they won't be sent over
        tuple_of_bools, layer_names, layer_types = client_config.aggregation_mask
        assert tuple_of_bools is not None, "Bools in the aggregation mask are None"
        list_of_bools = list(tuple_of_bools)
        # List of indices of the True elements in the list_of_bools
        true_indices = [i for i, x in enumerate(list_of_bools) if x]
        trimmed_layer_names: list[str] = []
        trimmed_layer_types: list[ModelStateNames] = []
        for i, (layer_name, layer_type) in enumerate(
            zip(layer_names, layer_types, strict=True),
        ):
            if layer_name in client_config.frozen_layers:
                # Set the corresponding index in the aggregation mask to False
                list_of_bools[true_indices[i]] = False
            else:
                trimmed_layer_names.append(layer_name)
                trimmed_layer_types.append(layer_type)
        client_config.aggregation_mask = (
            tuple(list_of_bools),
            trimmed_layer_names,
            trimmed_layer_types,
        )


def pass_expert_ids_to_callbacks(
    trainer: Trainer,
    experts_ids: list[int],
) -> None:
    """Pass the list of experts IDs to the callbacks.

    This function iterates over the callbacks in the Trainer and sets the
    `expert_ids` attribute for the `ExpertsScoresStatisticsMonitor` and
    `ExpertsAssignmentsMonitor` callbacks to the provided list of experts IDs.

    Parameters
    ----------
    trainer : Trainer
        The Composer Trainer instance containing the callbacks.
    experts_ids : list[int]
        The list of expert IDs to be set in the callbacks.

    """
    # Collect the callbacks from the Trainer
    callbacks = trainer.state.callbacks
    # Iterate over the callbacks and set the experts_ids only for
    # ExpertsScoresStatisticsMonitor and ExpertsAssignmentsMonitor callbacks
    for callback in callbacks:
        if hasattr(callback, "expert_ids"):
            callback.expert_ids = experts_ids  # type: ignore[reportAttributeAccessIssue]
