"""Module for extracting model and optimizer states from a Composer Trainer."""

import time
from collections import OrderedDict
from collections.abc import Iterable
from logging import DEBUG

import torch
from composer import Trainer
from composer.utils import dist
from composer.utils.iter_helpers import ensure_tuple
from flwr.common import (
    NDArray,
    NDArrays,
    Scalar,
    log,
)

from mt_dao.clients.configs import FitConfig
from mt_dao.composer_trainer_utils.shared_utils import (
    calculate_l2_norms_per_state_and_mode,
    filter_expert_parameters,
    get_ndarrays_and_names_from_payload,
    modify_aggregation_mask_with_frozen_layers,
)
from mt_dao.masks_utils import ModelStateNames
from mt_dao.moe_utils.utils import (
    split_experts_layers_tensors,
    split_experts_opt_dict_tensors,
)
from mt_dao.pytorch_utils.model_parameters_getters import (
    get_tensors_dict,
    get_wte_parameters_from_model,
)
from mt_dao.pytorch_utils.optimizer_states_getters import get_optimizer_states_base
from mt_dao.pytorch_utils.utils import get_n_experts_from_model
from mt_dao.shm.utils import compress_with_strict
from mt_dao.utils import ClientState, is_literal_for_ast


def get_wte_parameters_from_trainer(trainer: Trainer) -> tuple[str, NDArray]:
    """Extract WTE parameters from the trainer.

    This function retrieves the WTE parameters from the trainer's model. It assumes that
    the model has a specific structure where the WTE parameters can be identified by
    their names containing "wte". If the WTE parameters are not found or are not unique,
    it raises a ValueError.

    Parameters
    ----------
    trainer : Trainer
        The trainer object containing the model from which to extract WTE parameters.

    Returns
    -------
    tuple[str, NDArray]
        A tuple containing the name of the WTE parameters and the parameters themselves
        as a NumPy array.

    """
    # Get the model from the Trainer
    trainer_model = trainer.state.model
    # TODO(anon): Potential checks on the trainer model
    return get_wte_parameters_from_model(trainer_model)


def get_parameters_from_trainer(
    trainer: Trainer,
    parameter_names: Iterable[str] | None = None,
) -> NDArrays:
    """Extract model parameters from the trainer.

    This function retrieves the model parameters from the trainer's state dictionary.
    It can filter the parameters by names if `parameter_names` is provided.

    Parameters
    ----------
    trainer : Trainer
        The trainer object containing the model parameters.
    parameter_names : Iterable[str] | None, optional
        A list of parameter names to filter the parameters. If None, all parameters are
        returned.

    Returns
    -------
    NDArrays
        The extracted model parameters as NumPy arrays.

    """
    tensors_dict = get_tensors_dict(trainer.state.model)

    if n_experts := get_n_experts_from_model(
        trainer._original_model
        if isinstance(  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
            trainer.state.model,
            torch.nn.parallel.DistributedDataParallel,
        )
        else trainer.state.model,
    ):
        # For DEPToE models we need to split the experts
        tensors_dict = split_experts_layers_tensors(
            tensors_dict=tensors_dict,
            n_experts=n_experts,
        )
    return (
        [
            tensors_dict[next((rn for rn in tensors_dict if name in rn), "None")]
            .detach()
            .to("cpu")
            .numpy()
            for name in parameter_names
        ]
        if parameter_names is not None
        else [val.detach().to("cpu").numpy() for val in tensors_dict.values()]
    )


def extract_optimizer_states(
    opt_state_dict: dict[str, dict[str, torch.Tensor]],
    layer_names: list[str],
    layer_types: list[ModelStateNames],
    accumulator_per_type: dict[ModelStateNames, NDArrays],
) -> None:
    """Extract optimizer state tensors for specified layers and types.

    This function retrieves the optimizer state tensors (like momentum values) for
    specified layers and types from the optimizer state dictionary. It converts the
    optimizer state tensors to NumPy arrays for each layer and organizes them by state
    type.

    Parameters
    ----------
    opt_state_dict : dict[str, dict[str, torch.Tensor]]
        The optimizer state dictionary containing the optimizer states for each layer.
    layer_names : list[str]
        The names of the layers for which to extract optimizer states.
    layer_types : list[ModelStateNames]
        The types of optimizer states to extract (e.g., EXP_AVG, EXP_AVG_SQ).
    accumulator_per_type : dict[ModelStateNames, NDArrays]
        A dictionary to accumulate the extracted optimizer states, where keys are state
        types and values are lists of NumPy arrays containing the optimizer state
        tensors for that layer type.

    Raises
    ------
    ValueError
        If the optimizer state dictionary does not contain the expected keys for the
        specified layer names and types, or if the layer names do not match the expected
        format.

    """
    for layer_name, layer_type in zip(layer_names, layer_types, strict=True):
        local_layer_name = next(
            (
                local_layer_name_in_state
                for local_layer_name_in_state in opt_state_dict
                if layer_name in local_layer_name_in_state
            ),
            None,
        )

        if local_layer_name is None:
            msg = f"Parameter {layer_name} not found in optimizer state"
            raise ValueError(msg)

        if (lower_case_layer_type := layer_type.value.lower()) in opt_state_dict[
            local_layer_name
        ]:
            if layer_type not in accumulator_per_type:
                accumulator_per_type[ModelStateNames[layer_type.upper()]] = []

            accumulator_per_type[ModelStateNames[layer_type.upper()]].append(
                opt_state_dict[local_layer_name][lower_case_layer_type]
                .detach()
                .to("cpu")
                .numpy(),
            )


def get_optimizer_states_from_trainer(
    trainer: Trainer,
    layer_names: list[str],
    layer_types: list[ModelStateNames],
) -> dict[ModelStateNames, NDArrays]:
    """Extract optimizer state tensors for specified layers from the trainer.

    This function retrieves the optimizer state for specified layers and types from the
    trainer's state dictionary. It converts the optimizer state tensors (like momentum
    values) to NumPy arrays for each layer and organizes them by state type.

    Parameters
    ----------
    trainer : Trainer
        The trainer object containing the optimizer state.
    layer_names : list[str]
        The names of the layers for which to extract optimizer states.
    layer_types : list[ModelStateNames]
        The types of optimizer states to extract (e.g., EXP_AVG, EXP_AVG_SQ).

    Returns
    -------
    dict[ModelStateNames, NDArrays]
        A dictionary mapping each layer type to a list of NumPy arrays containing the
        optimizer state tensors for that layer type.

    """
    # NOTE: for some data parallelism implementations
    # the optimizer state dict may be empty
    # for all workers except rank 0
    # thus it only makes sense to add optimizer
    # states for rank 0
    accumulator_per_type: dict[ModelStateNames, NDArrays] = OrderedDict()

    dist.barrier()
    start_time = time.perf_counter()

    state = get_optimizer_states_base(
        model=trainer.state.model,
        optimizer=ensure_tuple(trainer.state.optimizers)[0],
        fsdp_state_dict_type=trainer.state.fsdp_state_dict_type,
        fsdp_enabled=trainer.state.fsdp_enabled,
    )

    if dist.get_global_rank() == 0:
        assert state is not None, (
            "Optimizer state is None. This may happen if the optimizer state is not "
            "initialized or if the model is not wrapped in FSDP."
        )
        log(DEBUG, "Checking if we have to split optimizer states by experts")

        if n_experts := get_n_experts_from_model(
            trainer._original_model
            if isinstance(  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
                trainer.state.model,
                torch.nn.parallel.DistributedDataParallel,
            )
            else trainer.state.model,
        ):
            log(DEBUG, "Splitting optimizer state for %s experts", n_experts)
            # For MPT models where we need to split the experts
            state = split_experts_opt_dict_tensors(
                opt_state_dict=state,
                n_experts=n_experts,
            )

        extract_optimizer_states(
            opt_state_dict=state,
            layer_names=layer_names,
            layer_types=layer_types,
            accumulator_per_type=accumulator_per_type,
        )
    dist.barrier()
    log(
        DEBUG,
        "Optimizer state extraction time: %s",
        time.perf_counter() - start_time,
    )
    return accumulator_per_type


def update_client_state(
    client_state: ClientState,
    trainer: Trainer,
    fit_config: FitConfig,
) -> None:
    """Update the client state with the current training progress.

    This function updates the client state with the number of local steps done,
    cumulative steps done, and the current timestamp. It is useful for tracking the
    training progress of the client in a distributed training setup.

    Parameters
    ----------
    client_state : ClientState
        The client state object to be updated.
    trainer : Trainer
        The trainer object containing the current training state.
    fit_config : FitConfig
        The fit configuration containing the training parameters.

    """
    # Get the number of local steps done by the current client
    num_batches_trained = fit_config.n_local_steps
    # Update the client state with the number of batches trained
    client_state.steps_done = num_batches_trained
    # Update the client state with the cumulative steps done
    client_state.local_steps_cumulative += num_batches_trained
    # Update the client state with the current timestamp
    client_state.local_timestamp = {
        k: v
        for k, v in trainer.state.timestamp.copy().state_dict().items()
        if is_literal_for_ast(repr(v))
    }


def calculate_trainer_post_training_l2_norms(
    model_states: dict[ModelStateNames, NDArrays],
    train_metrics: dict[str, Scalar],
    payload: NDArrays,
    fit_config: FitConfig,
) -> None:
    """Calculate L2 norms for model states after training.

    This function computes the L2 norms of the model states after training and updates
    the training metrics with the calculated norms. It uses the transmission mask to
    filter the relevant model states and calculates the norms for each state type.

    Parameters
    ----------
    model_states : dict[ModelStateNames, NDArrays]
        A dictionary containing the model states, where keys are state names and values
        are lists of NumPy arrays representing the state tensors.
    train_metrics : dict[str, Scalar]
        A dictionary to store the training metrics, where keys are metric names and
        values are the corresponding scalar values.
    payload : NDArrays
        The payload containing the model states to be processed.
    fit_config : FitConfig
        The fit configuration containing the transmission and aggregation masks, as well
        as the delta modes for calculating the L2 norms.

    """
    start_time = time.perf_counter()
    # Get the transmission and aggregation masks from the fit configuration
    (mask_transmission, layer_names_transmission, layer_types_transmission) = (
        fit_config.transmission_mask
    )
    (mask_aggregation, layer_names_aggregation, layer_types_aggregation) = (
        fit_config.aggregation_mask
    )
    # Filter from the mask experts that this client is not supposed to train
    # NOTE: The filtering must always be local and NOT modifying in place the mask
    if fit_config.experts_ids is not None:
        # Get the new boolean mask for the aggregation, but keep the global names
        (
            mask_aggregation,
            layer_names_aggregation,
            layer_types_aggregation,
        ) = filter_expert_parameters(
            mask_aggregation,
            layer_names_aggregation,
            layer_types_aggregation,
            fit_config.experts_ids,
            change_names=False,
        )
        # Get the new boolean mask for the transmission, but keep the global names
        (
            mask_transmission,
            layer_names_transmission,
            layer_types_transmission,
        ) = filter_expert_parameters(
            mask_transmission,
            layer_names_transmission,
            layer_types_transmission,
            fit_config.experts_ids,
            change_names=False,
        )
    for state_name, states in model_states.items():
        calculate_l2_norms_per_state_and_mode(
            aggregation_state=(states, layer_names_aggregation),
            transmission_state=get_ndarrays_and_names_from_payload(
                payload=list(
                    compress_with_strict(
                        data=payload,
                        selectors=mask_transmission,
                        strict=True,
                    ),
                ),
                layer_names=layer_names_transmission,
                layer_types=layer_types_transmission,
                model_state_type=state_name,
            ),
            delta_modes=fit_config.delta_modes,
            train_metrics=train_metrics,
            metric_keys=(
                f"client/l2_norm/{state_name.value}/global_model",
                f"client/l2_norm/{state_name.value}/local_model",
                f"client/l2_norm/{state_name.value}/pseudo_gradient",
            ),
        )
    train_metrics |= {
        "client/fit_metrics_collection_time": time.perf_counter() - start_time,
    }


def get_trainer_model_states(
    trainer: Trainer,
    fit_config: FitConfig,
    train_metrics: dict[str, Scalar],
) -> dict[ModelStateNames, NDArrays]:
    """Get model states from the trainer.

    This function retrieves the model states from the trainer, including parameters and
    optimizer states. It uses the aggregation mask from the fit configuration to filter
    the relevant model states. The function also measures the time taken to retrieve
    the model parameters and optimizer states, updating the training metrics
    accordingly.

    Parameters
    ----------
    trainer : Trainer
        The trainer object containing the model and optimizer states.
    fit_config : FitConfig
        The fit configuration containing the aggregation mask and other parameters.
    train_metrics : dict[str, Scalar]
        A dictionary to store the training metrics, where keys are metric names and
        values are the corresponding scalar values.

    Returns
    -------
    dict[ModelStateNames, NDArrays]
        A dictionary containing the model states, including parameters and optimizer
        states.

    """
    # Initialize the model states
    model_states: dict[ModelStateNames, NDArrays] = {}
    # Modify the aggregation mask in place in case there are any frozen layers
    modify_aggregation_mask_with_frozen_layers(fit_config)
    # Get the aggregation mask
    (mask_aggregation, layer_names_aggregation, layer_types_aggregation) = (
        fit_config.aggregation_mask
    )
    # Filter from the mask experts that this client is not supposed to train
    # NOTE: The filtering must always be local and NOT modifying in place the mask
    if fit_config.experts_ids is not None:
        # NOTE: We are replacing the references of these variables and NOT changing the
        # `fit_config.aggregation_mask` in place.
        (
            mask_aggregation,
            layer_names_aggregation,
            layer_types_aggregation,
        ) = filter_expert_parameters(
            mask_aggregation,
            layer_names_aggregation,
            layer_types_aggregation,
            fit_config.experts_ids,
        )
    start_get_model_parameters_time = time.perf_counter()
    # Get the aggregation model parameter names
    aggregation_parameter_names = [
        name
        for name, layer_type in zip(
            layer_names_aggregation,
            layer_types_aggregation,
            strict=True,
        )
        if layer_type == ModelStateNames.PARAMETERS
    ]
    model_states[ModelStateNames.PARAMETERS] = get_parameters_from_trainer(
        trainer,
        parameter_names=aggregation_parameter_names,
    )
    train_metrics |= {
        "client/fit_get_parameters_time": (
            time.perf_counter() - start_get_model_parameters_time
        ),
    }

    # NOTE: We removed here the parameter checkers for efficiency reasons.
    # Re-add them if needed

    # NOTE: We need to add the optimizer state to the model parameters
    start_get_optimizer_states_time = time.perf_counter()
    model_states.update(
        get_optimizer_states_from_trainer(
            trainer,
            layer_names=layer_names_aggregation,
            layer_types=layer_types_aggregation,
        ),
    )
    train_metrics |= {
        "client/fit_get_optimizer_states_time": (
            time.perf_counter() - start_get_optimizer_states_time
        ),
    }

    return model_states
