"""Module for getting model parameters from a PyTorch model."""

from functools import reduce

import torch
from composer.utils import dist
from flwr.common import (
    NDArray,
)
from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel

# TODO(anon): Handle modules prefixes globally
MODULE_PREFIXES_TO_REMOVE = [
    "model.",
    "module.",
    "_fsdp_wrapped_",
    "_checkpoint_wrapped_",
]

WEIGHT_TOKEN_EMBEDDING_NAME = "wte"  # noqa: S105


def get_wte_parameters_from_model(model: torch.nn.Module) -> tuple[str, NDArray]:
    """Get the WTE parameters from the model.

    This function retrieves the WTE parameters from the model's state dictionary.
    It assumes that the WTE parameters are uniquely identified by their names containing
    "wte". If the WTE parameters are not found or are not unique, it raises a
    ValueError.

    Parameters
    ----------
    model : torch.nn.Module
        The model from which to extract the WTE parameters.

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

    Raises
    ------
    ValueError
        If the model is None, if there are no WTE parameters, or if the WTE parameters
        are not unique.

    """
    # Get the WTE parameters
    wte_parameters_dict = {
        name: param.detach().to("cpu").numpy()
        for name, param in get_tensors_dict(model).items()
        if WEIGHT_TOKEN_EMBEDDING_NAME in name
    }
    # Return the WTE parameters
    wte_parameters = list(wte_parameters_dict.values())
    wte_parameters_name = list(wte_parameters_dict.keys())
    if len(wte_parameters) <= 0:
        msg = "There are no WTE parameters"
        raise ValueError(msg)
    if len(wte_parameters) != 1:
        msg = "WTE parameters are not unique"
        raise ValueError(msg)
    return wte_parameters_name[0], wte_parameters[0]


def get_tensors_dict(
    model: torch.nn.Module,
    *,
    only_requires_grad: bool = False,
    sort_dict: bool = True,
    no_detach_and_clone: bool = False,
) -> dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]:
    """Get the parameters of a model as an OrderedDict.

    Parameters
    ----------
    model : torch.nn.Module
        The model.
    only_requires_grad : bool, optional
        If True, only returns states referring to parameters that require gradients
        (default is False).
    sort_dict : bool, optional
        Whether to sort the dictionary, by default True.
    no_detach_and_clone : bool, optional
        Whether to detach and clone the parameters, by default False.

    Returns
    -------
    dict[str, torch.nn.Parameter] | dict[str, torch.Tensor]
        The trainable parameters.

    Raises
    ------
    ValueError
        If the model is None.

    """
    params_dict: dict[str, torch.nn.Parameter] | dict[str, torch.Tensor] = {}
    # NOTE: This function is weird because the encapsulation done to support FSDP and
    # DDP is weird. Since they are both likely to change, we MUST maintain this very
    # well and implement as many checkers as we can.
    if hasattr(model, "model") and type(model.model) is FullyShardedDataParallel:
        if model.model is None:  # type: ignore[reportUnnecessaryComparison]
            error_message = "Model is None"
            raise ValueError(error_message)
        inner_model = model.model
        # NOTE: This doesn't work in the case in use_orig_params is True if the FSDP
        # configuration as the tensors returned are flattened breaking some assumptions
        # of the rest of the codebase
        with FullyShardedDataParallel.summon_full_params(
            inner_model,
            recurse=True,
            writeback=False,
            rank0_only=True,
            offload_to_cpu=True,
            with_grads=False,
        ):
            # NOTE: This parameter dict using the above parameters, i.e., (recurse=True,
            # writeback=False, rank0_only=True, offload_to_cpu=True, with_grads=False,),
            # will be complete only on rank 0. The other ranks will have zero-shaped
            # tensors for those layers that are not "living" in there.
            # NOTE: If the FSDP configuration use the original parameters
            # (use_orig_params=true), then the tensors in rank 0 have the correct
            # original shape. In the other ranks they are flattened anyway.
            # NOTE: On ranks > 0 the dictionary won't be empty. It will contain the
            # parameters that are "living" in that rank and will have zero-shaped
            # tensors for the others.
            params_dict = {
                reduce(
                    lambda n, prefix: n.replace(prefix, ""),
                    MODULE_PREFIXES_TO_REMOVE,
                    name,
                ): (param.detach().clone() if not no_detach_and_clone else param)
                for name, param in inner_model.named_parameters()
                if not only_requires_grad or param.requires_grad
            }
    else:
        params_dict = {
            name: param.detach().clone() if not no_detach_and_clone else param
            for name, param in model.named_parameters()
            if not only_requires_grad or param.requires_grad
        }
    if sort_dict:
        params_dict = dict(sorted(params_dict.items()))
    dist.barrier()
    return params_dict
