"""Utility functions for initialising any model state from a saved checkpoint."""

import operator
import re
from pathlib import Path
from tempfile import TemporaryDirectory

import numpy as np
import torch
from composer import Trainer
from composer.utils.checkpoint import _torch_load_with_validation  # noqa: PLC2701
from flwr.common import NDArrays

from mt_dao.conf.base_schema import S3CommConfig
from mt_dao.constants import (
    MODEL_PARAMETERS,
    NUMBER_MODEL_STATES,
    NUMPY_ARRAY_SUFFIX,
    PYTORCH_SUFFIX,
    SAFETENSORS_SUFFIX,
)
from mt_dao.file_utils import load_chunked_model, load_model_states_from_file
from mt_dao.masks_utils import ModelStateNames
from mt_dao.moe_utils.utils import (
    split_experts_layers_tensors,
    transform_torch_checkpoint_to_mt_dao_model_states,
)
from mt_dao.pytorch_utils.model_parameters_getters import get_tensors_dict
from mt_dao.pytorch_utils.utils import get_n_experts_from_model
from mt_dao.server.s3_helper import get_file_from_path, list_objects


def generate_uninitialised_opt_states(model_states: NDArrays) -> NDArrays:
    """Generate unitisalised optimizer states.

    A helper function that is used to generate unitialised optimizer states
    when restoring checkpoints from a centralised run.

    Parameters
    ----------
    model_states : NDArrays
        The parameter model states to provide the shape reference.

    Returns
    -------
    NDArrays
       The unitialised optimizer states.

    """
    optimizer_states = []
    for _ in [
        ModelStateNames.EXP_AVG.value,
        ModelStateNames.EXP_AVG_SQ.value,
    ]:
        optimizer_states.extend(
            [np.zeros_like(parameter) for parameter in model_states],
        )
    return optimizer_states


def restore_check_from_safetensors(
    local_checkpoint_dir: Path,
    _desired_steps: int | None,
    _n_experts: int | None,
    _layer_names: list[str],
    *,
    preserve_optimizer_states: bool,
) -> NDArrays:
    """Restore centralised checkpoint from .npz file.

    A helper method that restores a centralised checkpoint from a .npz file.

    Parameters
    ----------
    local_checkpoint_dir : Path
        The directory where the checkpoint is stored.
    preserve_optimizer_states : bool, optional
        Whether to preserve the saved optimizer states.
    _layer_names : list[str]
        The names of the layers as defined by the Trainer object
        to ensure that we have a correct mapping.

    Returns
    -------
    NDArrays
        The loaded checkpointed model states.

    """
    # NOTE: We are in the case when we are assuming the full structured
    # checkpointing with the chunked updates.
    model_states_dir = local_checkpoint_dir / f"{MODEL_PARAMETERS}"
    model_states, parameter_names = load_chunked_model(
        model_states_dir / f"{ModelStateNames.PARAMETERS.value}",
    )
    model_states = [model_states[parameter_names.index(name)] for name in _layer_names]
    if preserve_optimizer_states:
        for momenta_type in [
            ModelStateNames.EXP_AVG.value,
            ModelStateNames.EXP_AVG_SQ.value,
        ]:
            momenta_path = model_states_dir / f"{momenta_type}"
            momenta_state, momenta_names = load_chunked_model(momenta_path)
            momenta_state = [
                momenta_state[momenta_names.index(name)] for name in _layer_names
            ]
            model_states.extend(momenta_state)
    else:
        optimizer_states = generate_uninitialised_opt_states(
            model_states=model_states,
        )
        model_states.extend(optimizer_states)
    return model_states


def restore_check_from_numpy(
    local_checkpoint_dir: Path,
    _desired_steps: int | None,
    _n_experts: int | None,
    _layer_names: list[str],
    *,
    preserve_optimizer_states: bool,
) -> NDArrays:
    """Restore centralised checkpoint from .npz file.

    A helper method that restores a centralised checkpoint from a .npz file.

    Parameters
    ----------
    local_checkpoint_dir : Path
        The directory where the checkpoint is stored.
    preserve_optimizer_states : bool, optional
        Whether to preserve the saved optimizer states.

    Returns
    -------
    NDArrays
        The loaded checkpointed model states.

    """
    # NOTE: We are in a case where are a restoring from a single checkpoint:
    # We are assuming that there is only a single .npz checkpoint contained
    # in the directory!.
    local_file_name = next(
        iter(local_checkpoint_dir.glob(f"*{NUMPY_ARRAY_SUFFIX}")),
    )
    _, model_states = load_model_states_from_file(
        local_file_name,
    )
    old_len = len(model_states)
    assert len(model_states) % NUMBER_MODEL_STATES == 0
    if not preserve_optimizer_states:
        new_model_states = model_states[: len(model_states) // NUMBER_MODEL_STATES]
        optimizer_states = generate_uninitialised_opt_states(
            model_states=new_model_states,
        )
        new_model_states.extend(optimizer_states)
        model_states = new_model_states
    assert old_len == len(model_states)
    return model_states


def restore_check_from_torch(
    local_checkpoint_dir: Path,
    _desired_steps: int | None,
    _n_experts: int | None,
    _layer_names: list[str],
    *,
    preserve_optimizer_states: bool,
) -> NDArrays:
    """Restore centralised checkpoint from .pt file.

    A helper method that restores a centralised checkpoint from a .pt
    file.

    Parameters
    ----------
    local_checkpoint_dir : Path
        The directory where the checkpoint is stored.
    preserve_optimizer_states : bool, optional
        Whether to preserve the saved optimizer states.
    _desired_steps: int | None
        The desired number of steps for loaded .pt checkpoint.
    _n_experts: int | None
        The number of experts in the MoE model should it be of
        this type.
    _layer_names : list[str]
        The names of the layers as defined by the Trainer object
        to ensure that we have a correct mapping.

    Returns
    -------
    NDArrays
        The loaded checkpointed model states.

    """
    checkpoints_to_download = list(
        map(
            str,
            local_checkpoint_dir.glob(f"*{PYTORCH_SUFFIX}"),
        ),
    )
    file_path = validate_torch_check(
        checkpoints_to_download=checkpoints_to_download,
        desired_steps=_desired_steps,
    )
    return transform_torch_checkpoint_to_mt_dao_model_states(
        model_states_dict=_torch_load_with_validation(
            checkpoint_filepath=str(local_checkpoint_dir) + file_path,
            map_location="cpu",
        ),
        n_experts=_n_experts,
        preserve_optimizer_states=preserve_optimizer_states,
        layer_ordering=_layer_names,
    )


def determine_checkpoints_and_type(
    centralised_checkpoint_dir: str,
) -> tuple[bool, list[str], set[str]]:
    """Determine checkpoints to download.

    This is a helper method that determine the checkpoint states to download.
    In addition to the actual files to download, this method determines the
    location of the checkpoints and the type of checkpoint we are expecting.

    Parameters
    ----------
    centralised_checkpoint_dir : str
        The directory where the checkpoints are located.

    Returns
    -------
    tuple[bool, list[str], set[str]]
        A tuple containing:
            - A bool determining the location of the checkpoints,
              where True indicates that the checkpoint is remote.
            - A list of strings containing in the checkpoint paths
            - A set of suffixes that determine the checkpoint types.

    Raises
    ------
    ValueError
        If no checkpoints are found at the checkpoint directory.

    """
    remote_or_local, checkpoints_to_download = list_objects(centralised_checkpoint_dir)
    if not checkpoints_to_download:
        msg = f"No centralised checkpoints found at {centralised_checkpoint_dir}"
        raise ValueError(msg)
    # Determine type of checkpoints we are dealing with.
    checkpoint_suffixes = {
        Path(checkpoint).suffix for checkpoint in checkpoints_to_download
    }
    return remote_or_local, checkpoints_to_download, checkpoint_suffixes


def validate_torch_check(
    checkpoints_to_download: list[str],
    desired_steps: int | None,
) -> str:
    """Validate a .pt checkpoint and return the path.

    A helper function that is used to validate the the requested .pt
    checkpoint exists, and if it does, returns the path of the checkpoint.

    Parameters
    ----------
    checkpoints_to_download : list[str]
        The possible checkpoints to download.
    desired_steps : int | None
        The number of steps that should be taken.
    centralised_checkpoint_dir : str
        The location of the centralised checkpoint.

    Returns
    -------
    str
        The path of the centralised checkpoint, correctly formatted.

    Raises
    ------
    ValueError
        If no checkpoint with the desired number of batches can be found locally/inS3.

    """
    sorted_pairs = sorted(
        [
            (
                int(reg.group(1)),  # epoch number
                int(reg.group(2)),  # number of batches
            )
            for path in checkpoints_to_download
            if (reg := re.search(r"/ep(\d+)-ba(\d+)", path)) is not None
        ],
        key=operator.itemgetter(1),
    )
    if desired_steps == -1:
        # NOTE: We always choose the longest training duration.
        path_to_check = next(reversed(sorted_pairs), None)
    else:
        path_to_check = next(
            (
                (epoch, batches)
                for epoch, batches in sorted_pairs
                if batches == desired_steps
            ),
            None,
        )
    if path_to_check is None:
        msg = f"Could not find a checkpoint with {desired_steps} batches"
        raise ValueError(msg)
    epoch, batches = path_to_check
    # NOTE: This is the format that centralised checkpoints are typcally written in.
    # We are always assuming that it is rank0.
    return f"/ep{epoch}-ba{batches}-" + "rank0.pt"


def extract_layer_names_and_experts(
    trainer: Trainer,
) -> tuple[list[str], int | None]:
    """Extract the layer names and experts from a Trainer object.

    A helper method that constructs the layer names and the number of
    experts for a model defined by the Trainer object.

    Parameters
    ----------
    trainer : Trainer
        The Trainer object.

    Returns
    -------
    tuple[list[str], int | None]
        A tuple containing:
        - A list of Strings containing the layer names.
        - The number of experts in the model configuration,
        if any.

    """
    # Obtain the list of trainable parameters names
    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,
    )
    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,
    )
    if n_experts:
        # For DEPToE models we need to split the experts
        tensors_dict = split_experts_layers_tensors(
            tensors_dict=tensors_dict,
            n_experts=n_experts,
        )
    param_names = list(
        tensors_dict.keys(),
    )
    param_names = [param_name.replace("model.", "", 1) for param_name in param_names]
    return param_names, n_experts


def retrieve_model_states_and_types(  # noqa: PLR0913
    local_checkpoint_directory: Path,
    checkpoint_suffixes: set[str],
    desired_steps: int | None,
    param_names: list[str],
    n_experts: int | None,
    *,
    preserve_optimizer_states: bool,
) -> tuple[NDArrays, list[tuple[str, ModelStateNames]]]:
    """Retrieve model states and layer types.

    A helper method that retrieves model states and names for
    a given Trainer context and local_checkpoint_directory, based
    on the filetypes present in the checkpoint suffixes. If the checkpoint
    suffix contains a .pt file, this will be loaded directly from the Trainer
    object. If the checkpoints suffixes contain either .npz or .safetensor
    files, these will be loaded directly from the local directory where
    they are located. Additionally, in all cases, the model state can be
    loaded with or without preserving the optimizer states.

    Parameters
    ----------
    local_checkpoint_directory : Path | None
        The directory where the checkpoints are stored.
    checkpoint_suffixes : set[str]
        A set containing the checkpoint suffix types,
        used to infer which checkpoints we are loading.
    preserve_optimizer_states : bool
        Whether to preserve the optimizer states in the
        checkpoint or not.
    desired_steps : int | None
        The desired number of steps that we should use to
        determine the correct checkpoint for .pt files.
    n_experts : int | None
        The number of experts in the model configuration, if any.
    param_names : list[str]
        A list of strings containing the names of the layer names.

    Returns
    -------
    tuple[NDArrays, list[tuple[str, ModelStateNames]]]
        A tuple containing:
        - model_states: The model parameters and optimizer states as NDArrays
        - layer_names_and_types: A list mapping parameter names to their state types

    Raises
    ------
    ValueError
        If the checkpoint in the checkpoint suffix is not
        supported for loading.

    """
    if SAFETENSORS_SUFFIX in checkpoint_suffixes:
        restore_fn = restore_check_from_safetensors
    elif NUMPY_ARRAY_SUFFIX in checkpoint_suffixes:
        restore_fn = restore_check_from_numpy
    elif PYTORCH_SUFFIX in checkpoint_suffixes:
        restore_fn = restore_check_from_torch
    else:
        msg = "The suffix type is not supported."
        raise ValueError(msg)
    model_states = restore_fn(
        local_checkpoint_dir=local_checkpoint_directory,
        preserve_optimizer_states=preserve_optimizer_states,
        _layer_names=param_names,
        _desired_steps=desired_steps,
        _n_experts=n_experts,
    )
    # Construct the layer names and types data structure
    layer_names_and_types: list[tuple[str, ModelStateNames]] = []
    for i in [
        ModelStateNames.PARAMETERS,
        ModelStateNames.EXP_AVG,
        ModelStateNames.EXP_AVG_SQ,
    ]:
        layer_names_and_types.extend([(param_name, i) for param_name in param_names])
    return model_states, layer_names_and_types


def retrieve_checkpoint_states(
    checkpoint_dir: str,
    s3_comm_config: S3CommConfig,
    run_uuid: str,
    desired_steps: int | None = -1,
) -> tuple[Path, TemporaryDirectory | None, bool, set[str]]:
    """Retrieve the checkpoint states for given checkpoints.

    This is a helper method that retrieves the checkpoints that
    are specified in a given directory. This occurs for both local
    and remote s3 directories. In the case of the local directory,
    the presence of the local files is verified. In the case of remote
    s3 checkpoints, the checkpoints are downloaded to a temporary
    directory, which should be deleted later.

    Parameters
    ----------
    checkpoint_dir : str
        The directory that contains the checkpoints.
    s3_comm_config : S3CommConfig
        The config for s3_communication.
    run_uuid : str
        The uuid of the centralised checkpoints.
    desired_steps : int | None
        The desired steps from the checkpoint to be downloaded
        in case of a .pt.

    Returns
    -------
    tuple[Path, TemporaryDirectory | None]
        A tuple containing:
        - The local checkpoint path.
        - The temporary directory that was used to store the checkpoints,
        locally.
        - A bool indicating whether the checkpoitn is local or remote,
        where True indicates remote.
        - A set of Strings that determines which suffixes are contained in
        the checkpoints, which determines the file type.

    Raises
    ------
    ValueError
        If the requested checkpoints are not in the supported
        file types.

    """
    checkpoint_data = determine_checkpoints_and_type(
        centralised_checkpoint_dir=checkpoint_dir,
    )
    remote_or_local, checkpoints_to_download, checkpoint_suffixes = checkpoint_data
    if (
        NUMPY_ARRAY_SUFFIX in checkpoint_suffixes
        or SAFETENSORS_SUFFIX in checkpoint_suffixes
        or PYTORCH_SUFFIX in checkpoint_suffixes
    ):
        if remote_or_local:
            tmp_dir = TemporaryDirectory(delete=False)
            local_checkpoint_directory = Path(tmp_dir.name)
        else:
            tmp_dir = None
            local_checkpoint_directory = Path(checkpoint_dir)
        # NOTE: Filter out .pt checkpoints to be only the batch we have requested
        # so that we do not download unecessarily.
        if PYTORCH_SUFFIX in checkpoint_suffixes:
            file_str = validate_torch_check(
                checkpoints_to_download=checkpoints_to_download,
                desired_steps=desired_steps,
            )
            for file_path in checkpoints_to_download:
                if file_str in file_path:
                    checkpoints_to_download = [file_path]
                    break
        # Downloading the files to a temp local directory if necessary.
        for file_name in checkpoints_to_download:
            if remote_or_local:
                # Format s3 string correctly.
                bucket = s3_comm_config.bucket_name
                file_to_download = f"s3://{bucket}/" + file_name
                relative_header = str(
                    Path(
                        *Path(checkpoint_dir).parts[2:],
                    ),
                )
            else:
                file_to_download = file_name
                relative_header = ""
            _ = get_file_from_path(
                output_parent_path=str(local_checkpoint_directory),
                input_file_path=file_to_download,
                run_uuid=run_uuid,
                s3_comm_config=s3_comm_config,
                relative_header=relative_header,
            )
    else:
        msg = f"We cannot currently support these checkpoints: {checkpoint_suffixes}"
        raise ValueError(msg)
    return local_checkpoint_directory, tmp_dir, remote_or_local, checkpoint_suffixes
