"""Module for setting the trainer's model and optimizer states."""

import gc
import time
from collections import OrderedDict, defaultdict
from logging import DEBUG
from typing import cast

import numpy as np
import torch
from composer import Trainer
from composer.trainer.trainer import _clear_incomplete_train_states  # noqa: PLC2701
from composer.utils import dist
from composer.utils.iter_helpers import ensure_tuple
from flwr.common import (
    NDArrays,
    Scalar,
    log,
)
from omegaconf import DictConfig

from mt_dao.clients.configs import EvaluateConfig, FitConfig
from mt_dao.clients.utils import (
    randomize_layers,
)
from mt_dao.composer_trainer_utils.shared_utils import (
    filter_expert_parameters,
    get_ndarrays_and_names_from_payload,
    pass_expert_ids_to_callbacks,
)
from mt_dao.masks_utils import ModelStateNames
from mt_dao.moe_utils.utils import (
    merge_experts_layers_tensors,
    merge_experts_opt_dict_tensors,
    split_experts_layers_tensors,
)
from mt_dao.optimizers import ADOPT, QHADOPT, DecoupledAdamW
from mt_dao.pytorch_utils.model_parameters_getters import (
    get_tensors_dict,
)
from mt_dao.pytorch_utils.model_parameters_setters import (
    set_model_tensors_dict,
)
from mt_dao.pytorch_utils.optimizer_states_setters import set_optimizers_states
from mt_dao.pytorch_utils.utils import (
    get_n_experts_from_model,
    invert_opt_state_dict_keys,
)
from mt_dao.shm.utils import compress_with_strict
from mt_dao.utils import (
    ClientState,
    create_tensors_dict,
    set_frozen_layers_tensors_dict,
)


def set_trainer_tensors_dict(
    trainer: Trainer,
    tensors_dict: OrderedDict[str, torch.Tensor],
) -> None:
    """Set the tensors dictionary for the trainer's model.

    This function sets the tensors dictionary for the trainer's model. It handles both
    FSDP (Fully Sharded Data Parallel) and non-FSDP model architectures. For FSDP
    modules, it creates helper state dictionaries and sets the tensors for each leaf
    FSDP module. For non-FSDP modules, it loops over each parameter and sets the tensors
    directly.

    Parameters
    ----------
    trainer : Trainer
        The trainer object containing the model and optimizer.
    tensors_dict : OrderedDict[str, torch.Tensor]
        An ordered dictionary containing the tensors to be set in the model. The keys
        are the names of the tensors, and the values are the corresponding tensor
        values. It can be incomplete, meaning it may not contain all the tensors
        present in the model.

    """
    if n_experts := get_n_experts_from_model(
        (
            trainer._original_model  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
            if isinstance(
                trainer.state.model,
                torch.nn.parallel.DistributedDataParallel,
            )
            else trainer.state.model
        ),
    ):
        # For DEPToE models we need to split the experts
        tensors_dict = OrderedDict(
            merge_experts_layers_tensors(
                tensors_dict=dict(tensors_dict),
                n_experts=n_experts,
            ),
        )
    # TODO(anon): Potential checks on the trainer model
    set_model_tensors_dict(
        model=trainer.state.model,
        tensors_dict=tensors_dict,
    )


def set_trainer_states_for_fit(  # noqa: PLR0914
    payload: NDArrays,
    trainer: Trainer,
    configs: tuple[FitConfig, DictConfig],
    client_state_struct: ClientState,
    train_metrics: dict[str, Scalar],
) -> None:
    """Set the trainer's model and optimizer states for training.

    This function sets the model and optimizer states for the trainer using the provided
    payload and configurations. It handles the extraction and setting of parameters,
    momenta, and other relevant states.
    It also manages the randomization of layers based on the fit configuration and
    client state.

    Parameters
    ----------
    payload : NDArrays
        The payload containing the model parameters and optimizer states.
    trainer : Trainer
        The trainer object containing the model and optimizer.
    configs : tuple[FitConfig, DictConfig]
        A tuple containing the fit configuration and a dummy configuration.
    client_state_struct : ClientState
        The client state structure containing information about the client's local
        steps.
    train_metrics : dict[str, Scalar]
        A dictionary to store training metrics, such as time taken to set optimizer
        state.

    """
    fit_config, dummy_config = configs

    # Get parameter names from trainer
    assert fit_config.transmission_mask is not None, "Transmission mask is None"

    # Retrieve the transmission mask
    # TODO(anon): Exclude from the transmission masks the personalized layers,
    # which are contained in `fit_config.personalized_layers`, if there's any
    (mask_transmission, layer_names_transmission, layer_types_transmission) = (
        fit_config.transmission_mask
    )

    # Filter 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:
        log(DEBUG, "Setting experts_ids for the client: %s", fit_config.experts_ids)
        # Pass the experts_ids to the loggers of experts metrics
        pass_expert_ids_to_callbacks(
            trainer=trainer,
            experts_ids=fit_config.experts_ids,
        )
        # NOTE: We are replacing the references of these variables and NOT changing the
        # `fit_config.transmission_mask` in place.
        (
            mask_transmission,
            layer_names_transmission,
            layer_types_transmission,
        ) = filter_expert_parameters(
            mask_transmission,
            layer_names_transmission,
            layer_types_transmission,
            fit_config.experts_ids,
        )

    transmission_payload = list(
        compress_with_strict(
            data=payload,
            selectors=mask_transmission,
            strict=True,
        ),
    )

    new_parameters, new_parameters_names = get_ndarrays_and_names_from_payload(
        payload=transmission_payload,
        layer_names=layer_names_transmission,
        layer_types=layer_types_transmission,
        model_state_type=ModelStateNames.PARAMETERS,
    )

    # Get list of random layers from config
    random_layers: list[str] = (
        fit_config.random_layers if fit_config.random_layers is not None else []
    )
    # Randomize layers
    # TODO(anon,anon): Use the underlying re-init methods of the single modules for
    # this one and make sure the logic is clearly reseting modules
    if (
        random_layers
        and fit_config.random_init_freq > 0
        and client_state_struct.local_steps_cumulative % fit_config.random_init_freq
        == 0
    ):
        tensors_dict = get_tensors_dict(
            trainer.state.model,
            # NOTE: The random model has all trainable parameters that we should
            # care about with the correct flag set
            only_requires_grad=True,
            # NOTE: We need to sort the parameters to ensure that the order is
            # consistent with the names later
            sort_dict=True,
        )

        if n_experts := get_n_experts_from_model(
            (
                trainer._original_model  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
                if isinstance(
                    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,
            )
        randomize_layers(
            parameters=new_parameters,
            dummy_config=dummy_config,
            names=list(
                tensors_dict.keys(),
            ),
            random_layers=random_layers,
            truly_random_init=fit_config.truly_random_init,
            cid=int(fit_config.cid),
            server_round=fit_config.server_round,
        )

    if fit_config.reset_optimizer:
        # Reset the optimizer state
        log(DEBUG, "Resetting optimizer states to zero, step included.")
        fit_config.server_steps_cumulative = 0
        # Get the NDArrays from the full mask
        (_, full_layer_names, full_layer_types) = fit_config.full_mask
        first_momentum, first_momentum_names = get_ndarrays_and_names_from_payload(
            payload=payload,
            layer_names=full_layer_names,
            layer_types=full_layer_types,
            model_state_type=ModelStateNames.EXP_AVG,
        )
        second_momentum, second_momentum_names = get_ndarrays_and_names_from_payload(
            payload=payload,
            layer_names=full_layer_names,
            layer_types=full_layer_types,
            model_state_type=ModelStateNames.EXP_AVG_SQ,
        )
        # Zero out the momentum vectors
        first_momentum = [np.zeros_like(m) for m in first_momentum]
        second_momentum = [np.zeros_like(m) for m in second_momentum]
    else:
        first_momentum, first_momentum_names = get_ndarrays_and_names_from_payload(
            payload=transmission_payload,
            layer_names=layer_names_transmission,
            layer_types=layer_types_transmission,
            model_state_type=ModelStateNames.EXP_AVG,
        )

        second_momentum, second_momentum_names = get_ndarrays_and_names_from_payload(
            payload=transmission_payload,
            layer_names=layer_names_transmission,
            layer_types=layer_types_transmission,
            model_state_type=ModelStateNames.EXP_AVG_SQ,
        )

    assert fit_config.server_steps_cumulative is not None
    if len(first_momentum + second_momentum) > 0:
        time_to_set = set_trainer_optimizer_states(
            trainer=trainer,
            momentum_vectors_set=(
                first_momentum + second_momentum,
                first_momentum_names + second_momentum_names,
                [ModelStateNames.EXP_AVG.value] * len(first_momentum_names)
                + [ModelStateNames.EXP_AVG_SQ.value] * len(second_momentum_names),
            ),
            fit_config=fit_config,
        )
        log(DEBUG, "Optimizer state set in %s seconds", time_to_set)
        train_metrics |= {"client/local_adopt/set_optimizer_state_time": time_to_set}
    dist.barrier()

    # Set the model tensors dictionary
    set_trainer_tensors_dict(
        trainer=trainer,
        tensors_dict=set_frozen_layers_tensors_dict(
            create_tensors_dict(
                parameters=new_parameters,
                parameters_names=new_parameters_names,
            ),
            frozen_layers=fit_config.frozen_layers,
        ),
    )
    gc.collect()
    torch.cuda.empty_cache()


def set_trainer_states_for_eval(
    payload: NDArrays,
    trainer: Trainer,
    eval_config: EvaluateConfig,
) -> None:
    """Set the trainer's model states for evaluation.

    This function sets the model states for the trainer using the provided payload and
    evaluation configuration.

    Parameters
    ----------
    payload : NDArrays
        The payload containing the model parameters.
    trainer : Trainer
        The trainer object containing the model.
    eval_config : EvaluateConfig
        The evaluation configuration containing the transmission mask and other
        parameters.

    """
    # Get parameter names from trainer
    assert eval_config.transmission_mask is not None, "Transmission mask is None"

    # Retrieve the transmission mask
    # TODO(anon): Exclude from the transmission masks the personalized layers,
    # which are contained in `eval_config.personalized_layers`, if there's any
    (mask_transmission, layer_names_transmission, layer_types_transmission) = (
        eval_config.transmission_mask
    )
    transmission_payload = list(
        compress_with_strict(
            data=payload,
            selectors=mask_transmission,
            strict=True,
        ),
    )

    # Filter experts that this client is not supposed to evaluate
    if eval_config.experts_ids is not None:
        # TODO(anon): Find a way to pass the experts_ids to the loggers of experts
        # metrics
        (
            mask_transmission,
            layer_names_transmission,
            layer_types_transmission,
        ) = filter_expert_parameters(
            mask_transmission,
            layer_names_transmission,
            layer_types_transmission,
            eval_config.experts_ids,
        )

    new_parameters, new_parameters_names = get_ndarrays_and_names_from_payload(
        payload=transmission_payload,
        layer_names=layer_names_transmission,
        layer_types=layer_types_transmission,
        model_state_type=ModelStateNames.PARAMETERS,
    )

    new_tensors_dict = create_tensors_dict(
        parameters=new_parameters,
        parameters_names=new_parameters_names,
    )

    if n_experts := get_n_experts_from_model(
        (
            trainer._original_model  # pyright: ignore[reportAttributeAccessIssue]  # noqa: SLF001
            if isinstance(
                trainer.state.model,
                torch.nn.parallel.DistributedDataParallel,
            )
            else trainer.state.model
        ),
    ):
        # For DEPToE models we need to split the experts
        new_tensors_dict = OrderedDict(
            merge_experts_layers_tensors(
                tensors_dict=new_tensors_dict,
                n_experts=n_experts,
            ),
        )

    set_trainer_tensors_dict(
        trainer=trainer,
        tensors_dict=new_tensors_dict,
    )
    gc.collect()
    torch.cuda.empty_cache()


def set_trainer_optimizer_states(
    momentum_vectors_set: tuple[NDArrays, list[str], list[str]],
    trainer: Trainer,
    fit_config: FitConfig,
) -> float:
    """Set optimizer states with momentum vectors for model training.

    This function sets the optimizer states (momentum vectors) for a trainer's model.
    It handles both FSDP (Fully Sharded Data Parallel) and non-FSDP model architectures.
    For FSDP modules, it creates helper state dictionaries and sets the optimizer state
    for each leaf FSDP module. For non-FSDP modules, it loops over each parameter and
    sets the optimizer state directly.

    Parameters
    ----------
    momentum_vectors_set : tuple[NDArrays, list[str], list[str]]
        A tuple containing:
        - NDArrays of momentum values
        - List of momentum parameter names
        - List of momentum types (e.g., 'exp_avg', 'exp_avg_sq')
    trainer : Trainer
        The trainer object containing the model and optimizer.
    fit_config : FitConfig
        The configuration containing server steps information.

    Returns
    -------
    float
        The time in seconds taken to set the optimizer state.

    Raises
    ------
        AssertionError
            If server_steps_cumulative is None or the optimizer type is not supported.

    """
    log(DEBUG, "Setting optimizer state")
    assert fit_config.server_steps_cumulative is not None, (
        "Server steps cumulative is None"
    )
    dist.barrier()
    astart_time = time.perf_counter()

    # Get the optimizer from the trainer object
    optimizer: torch.optim.Optimizer = ensure_tuple(trainer.state.optimizers)[0]
    assert isinstance(
        optimizer,
        ADOPT | DecoupledAdamW | QHADOPT,
    ), f"The optimizer {type(optimizer)} is not supported for setting the states."
    # NOTE: It seems necessary to leave this call. The optimizer results incomplete with
    # partial or missing states otherwise. It does not seem to represent a big time
    # overhead as far as we have experimented. The log is left as monitor.
    _clear_incomplete_train_states(trainer.state)

    # Create a helper optimizer state dict for NDArrays
    opt_state_dict: dict[str, dict[str, torch.Tensor]] = cast(
        "dict[str, dict[str, torch.Tensor]]",
        defaultdict(lambda: defaultdict(dict)),
    )
    for momentum, momentum_name, momentum_type in zip(
        *momentum_vectors_set,
        strict=True,
    ):
        opt_state_dict[momentum_name][momentum_type] = torch.as_tensor(momentum)

    # If there are experts to merge, we should do it here
    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 MPT models where we need to split the experts
        opt_state_dict = merge_experts_opt_dict_tensors(
            opt_state_dict=opt_state_dict,
            n_experts=n_experts,
        )

    set_optimizers_states(
        model=trainer.state.model,
        optimizer=optimizer,
        step=fit_config.server_steps_cumulative,
        opt_state_dict=invert_opt_state_dict_keys(opt_state_dict),
    )

    dist.barrier()
    return time.perf_counter() - astart_time
