"""Module for setting model parameters in a PyTorch model."""

from collections import OrderedDict, defaultdict
from logging import WARNING

import torch
from composer.utils import dist
from flwr.common import (
    log,
)
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_model_state_dict,
    set_model_state_dict,
)
from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel

from mt_dao.module_utils import collect_fsdp_modules, is_leaf_fsdp_module
from mt_dao.utils import (
    ParameterShapeMismatchError,
    match_shapes,
)


def set_model_tensors_dict(
    model: torch.nn.Module,
    tensors_dict: OrderedDict[str, torch.Tensor],
) -> None:
    """Set the trainable parameters of a model from an ordered dictionary.

    Note that the `tensors_dict` can also be partial.
    This function handles both Fully Sharded Data Parallel (FSDP) modules and regular
    PyTorch modules. It collects all FSDP modules in the model hierarchy and sets their
    parameters accordingly. If the model is not FSDP, it sets the parameters directly
    from the provided dictionary.

    Parameters
    ----------
    model : torch.nn.Module
        The model whose parameters are to be set.
    tensors_dict : OrderedDict[str, torch.Tensor]
        An ordered dictionary mapping parameter names to their tensor values.
        Parameters that are successfully set will be removed from this dictionary.

    Raises
    ------
    ParameterShapeMismatchError
        If the shapes of the remote and local parameters do not match when setting
        the parameters.

    """
    # Map to store FSDP modules by their name
    fsdp_modules: dict[str, FullyShardedDataParallel] = {}
    non_fsdp_modules: dict[str, torch.nn.Module] = {}

    # Collect all FSDP modules in the model hierarchy
    if hasattr(model, "model") and isinstance(model.model, FullyShardedDataParallel):
        model = model.model
        # First pass: collect all FSDP modules
        (fsdp_modules, non_fsdp_modules) = collect_fsdp_modules(
            "model",
            model,
            (fsdp_modules, non_fsdp_modules),
        )

    if fsdp_modules:
        # Process each FSDP module that's a leaf module or direct parent of leaf modules
        for local_module_name, local_module in fsdp_modules.items():
            if is_leaf_fsdp_module(local_module_name, (fsdp_modules, non_fsdp_modules)):
                set_params_fsdp(local_module_name, local_module, tensors_dict)
    else:
        # If model is not FSDP, handle the regular case
        for local_param_name, local_param in model.named_parameters():
            # Set the parameters only if they require gradients
            if local_param.requires_grad:
                # Remove potential DDP-related prefixes and the "module." prefix
                lookup_name = local_param_name.replace("model.", "").replace(
                    "module.",
                    "",
                )
                if lookup_name not in tensors_dict:
                    log(
                        WARNING,
                        "Parameter %s not found in the list of remote parameters and"
                        " won't be set.",
                        local_param_name,
                    )
                else:
                    success = _set_param_base(
                        tensors_dict[lookup_name],
                        local_param,
                    )
                    # Raise error if the shapes don't match
                    if not success:
                        raise ParameterShapeMismatchError(
                            remote_name=lookup_name,
                            remote_shape=tensors_dict[lookup_name].shape,
                            local_shape=local_param.shape,
                        )
    dist.barrier()


def set_params_fsdp(
    local_module_name: str,
    local_module: FullyShardedDataParallel,
    tensors_dict: OrderedDict[str, torch.Tensor],
) -> None:
    """Set parameters for a FullyShardedDataParallel module.

    This function sets parameters for an FSDP module by handling the complexities of
    sharded parameters. It attempts to directly set parameters and collects any
    parameters that couldn't be set directly into a leftover buffer, which will be
    handled through the state dictionary.

    Parameters
    ----------
    local_module_name : str
        The name of the FSDP module whose parameters are to be set.
    local_module : FullyShardedDataParallel
        The FSDP module instance whose parameters are to be set.
    tensors_dict : OrderedDict[str, torch.Tensor]
        An ordered dictionary mapping parameter names to their tensor values.
        Parameters that are successfully set will be removed from this dictionary.

    """
    # Initialize the leftover buffer for parameters that are sharded across multiple
    # ranks not with full shards
    leftover_buffer: dict[
        str,
        tuple[torch.nn.Module, str, torch.Tensor],
    ] = {}
    # Open the FSDP context manager to summon the parameters
    with FullyShardedDataParallel.summon_full_params(
        local_module,
        writeback=True,
        recurse=False,
    ):
        # Iterate over the named modules of the local module
        for (
            internal_local_module_name,
            internal_local_module,
        ) in local_module.named_modules():
            # Iterate over the named parameters of the current internal local module
            for local_param_name, local_param in internal_local_module.named_parameters(
                recurse=False,
            ):
                # Construct the remote parameter name by removing the "model." and FSDP
                # prefixes and by adding the local module name, the internal local
                # module name, and the local parameter name
                # TODO(anon): Handle modules prefixes globally
                remote_param_name = (
                    (
                        (local_module_name + "." + internal_local_module_name)
                        .replace("model.", "")
                        .replace("_fsdp_wrapped_module.", "")
                        .replace("._fsdp_wrapped_module", "")
                        .replace("._checkpoint_wrapped_module", "")
                    )
                    + "."
                    + local_param_name
                )
                # Set remote parameter name if in the list of remote parameters,
                if remote_param_name in tensors_dict:
                    # Get the remote parameter from the OrderedDict
                    remote_params = tensors_dict.pop(remote_param_name)
                    # Set the parameters
                    success = _set_param_base(
                        remote_params,
                        local_param,
                    )
                    # Append to the leftover buffer if the assignment failed
                    if not success:
                        leftover_buffer[remote_param_name] = (
                            internal_local_module,
                            remote_param_name,
                            remote_params,
                        )
                else:
                    log(
                        WARNING,
                        "set_params_fsdp: Parameter %s not found in the list of remote"
                        " parameters and won't be set",
                        remote_param_name,
                    )
    # Set the leftover parameters through the state dict
    set_params_with_state_dict(leftover_buffer)


def set_params_with_state_dict(
    modules_buffer: dict[str, tuple[torch.nn.Module, str, torch.Tensor]],
) -> None:
    """Set parameters for modules using state dictionaries.

    This function efficiently updates parameters across multiple modules by leveraging
    PyTorch's state dictionary functionality. It first reorganizes the input buffer to
    process each module only once, then retrieves each module's state dict, updates
    the parameters, and applies the changes back to the module.

    Parameters
    ----------
    modules_buffer : dict[str, tuple[torch.nn.Module, str, torch.Tensor]]
        A dictionary mapping keys to tuples containing:
        - The module instance to update
        - The parameter name as it appears in the module's state dict
        - The tensor containing the new parameter values

    Notes
    -----
    This function is particularly useful for updating parameters in complex module
    hierarchies, as it handles the state dict management required for proper parameter
    updates. It's designed to work with both regular PyTorch modules and modules with
    special parameter handling like FSDP.

    """
    # Construct the internal mapping that allows us to iterate over each module in the
    # modules_buffer only once
    internal_buffer: dict[torch.nn.Module, dict[str, torch.Tensor]] = defaultdict(dict)
    for (
        internal_local_module,
        remote_param_name,
        remote_param,
    ) in modules_buffer.values():
        internal_buffer[internal_local_module] |= {remote_param_name: remote_param}
    # Iterate over the internal buffer and set the parameters
    for (
        internal_local_module,
        remote_param_dict,
    ) in internal_buffer.items():
        # Get the state dict of the local module
        state = get_model_state_dict(
            internal_local_module,
            options=StateDictOptions(
                full_state_dict=True,
                cpu_offload=True,
                broadcast_from_rank0=True,
            ),
        )
        # Set the parameters
        for local_param_name, local_param in state.items():
            assert isinstance(local_param, torch.Tensor)
            _set_param_base(
                remote_param_dict[local_param_name],
                local_param,
            )
        # Set the state dict of the local module
        set_model_state_dict(
            internal_local_module,
            state,
            options=StateDictOptions(
                full_state_dict=True,
                cpu_offload=True,
                broadcast_from_rank0=True,
            ),
        )


def _set_param_base(
    remote_param: torch.Tensor,
    local_param: torch.nn.Parameter | torch.Tensor,
) -> bool:
    """Set a local parameter's data with a remote parameter's value.

    This function handles the base case of setting a parameter's value, taking care of
    shape matching and device placement. It's used by higher-level parameter setting
    functions to handle individual parameter assignments. It also sets the
    `requires_grad` flag of the local parameter to match the remote parameter.

    Parameters
    ----------
    remote_param : torch.Tensor
        The remote parameter tensor containing the values to assign.
    local_param : torch.nn.Parameter | torch.Tensor
        The local parameter or tensor whose data will be updated.

    Returns
    -------
    bool
        True if the assignment was successful or not required (empty local parameter),
        False if the assignment failed due to shape mismatch.

    """
    # Skip the setting if the remote parameter is None or the local parameter is empty.
    # The first case happens if the parameter is not found in the remote model, the
    # second case happens if the parameter is not found in the local model, e.g., the
    # current shard doesn't have the parameter.
    if local_param.shape != torch.Size([0]):
        # Try to match the shapes
        remote_param = match_shapes(remote_param, local_param)
        # Set the parameters only if the shape matches
        if local_param.shape == remote_param.shape:
            local_param.data.copy_(
                remote_param.to(device=local_param.device, dtype=local_param.dtype),
            )
            # Set the requires_grad flag to match the remote parameter
            local_param.requires_grad = remote_param.requires_grad
        # Return False if the assignment failed
        return local_param.shape == remote_param.shape
    # Return True if the assignment was not required (empty local parameter)
    return True
