"""Utility functions for running S3-related tasks on main server loop in flwr next."""

import operator
import pickle  # noqa: S403
import random
import re
import shutil
import time
from collections.abc import Callable
from itertools import groupby
from logging import DEBUG, INFO, WARNING
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import cast

import numpy as np
import torch
from composer import Trainer
from composer.loggers import RemoteUploaderDownloader
from composer.utils import (
    S3ObjectStore,
    parse_uri,
)
from composer.utils.file_helpers import (
    list_remote_objects,
    validate_given_remote_path,
)
from flwr.common import (
    NDArrays,
    log,
)

from mt_dao.composer_trainer_utils.composer_trainer_setters import (
    set_trainer_tensors_dict,
)
from mt_dao.conf.base_schema import BaseConfig, S3CommConfig, StrategyName
from mt_dao.constants import (
    BIN_SUFFIX,
    CURRENT_SERVER_STATE_FILENAME,
    FIRST_MOMENTUM,
    MODEL_PARAMETERS,
    NUMPY_ARRAY_SUFFIX,
    SECOND_MOMENTUM,
)
from mt_dao.file_utils import (
    custom_pickle_load,
    download_file_from_s3,
    load_chunked_model,
    load_model_states_from_file,
    upload_file_to_s3,
)
from mt_dao.moe_utils.utils import split_experts_layers_tensors
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.client_sampler import ClientSampler
from mt_dao.server.param_scheduler_dispatcher import dispatch_model_state_scheduler
from mt_dao.server.s3_helper import (
    delete_clients_checkpoints,
    delete_object,
    delete_rounds,
    generic_model_state_upload_to_s3,
    get_file_from_path,
    key_exists,
    list_objects,
    obtain_sorted_runs,
)
from mt_dao.server.server_state import (
    ServerState,
    check_server_state_lightweight_dict,
    create_server_state_lightweight_dict,
    read_server_state_lightweight_dict,
)
from mt_dao.strategy_functional.dispatcher import dispatch_strategy_state_keys
from mt_dao.utils import (
    ClientState,
    ModelStateNames,
    create_tensors_dict,
    get_random_model_states_from_config,
)
from mt_dao.wandb_history import WandbHistory


class NoCheckpointsFoundError(Exception):
    """Exception raised when there are no checkpoints in the path looked up."""


def interpret_resume_round(
    resume_round: int | None,
    run_uuid_path: str,
    state_keys: tuple[str, ...],
    *,
    raise_error: bool = True,
) -> int | None:
    """Interpret the resume round parameter for server checkpoint resumption.

    This function interprets the `resume_round` parameter, which specifies the round
    to resume server operations from. If `resume_round` is negative, it is treated as
    an index into the list of sorted rounds obtained from the server's path, allowing
    for reverse indexing. If `resume_round` is None, the function returns None,
    indicating no specific round to resume from. An error is raised if no checkpoints
    are found when `raise_error` is True and `resume_round` is negative but no rounds
    are available.

    Parameters
    ----------
    resume_round : int | None
        The round number to resume from. If negative, treated as a reverse index. If
        None, indicates no resumption is required.
    run_uuid_path : str
        The path to the run uuid root.
    state_keys : tuple[str, ...]
        A tuple of state keys used to identify the federated rounds.
    raise_error : bool, optional
        Whether to raise an error if no checkpoints are found and `resume_round` is
        negative. Default is True.

    Returns
    -------
    int | None
        The interpreted round number to resume from, or None if no resumption.

    Raises
    ------
    NoCheckpointsFoundError
        If `raise_error` is True, no checkpoints are found, and `resume_round` < 0.

    """
    log(
        DEBUG,
        "The parameter `resume_round=%s` will be interpret as an index "
        "for the list of rounds for the run_uuid_path=%s",
        resume_round,
        run_uuid_path,
    )
    if resume_round is None:
        return None
    if resume_round < 0:
        server_round_indices = obtain_sorted_runs(run_uuid_path, state_keys)
        log(DEBUG, "Found server round indices %s", server_round_indices)
        if not server_round_indices and raise_error:
            raise NoCheckpointsFoundError
        if server_round_indices:
            resume_round = server_round_indices[resume_round]
    return resume_round


def upload_server_state(
    server_state: ServerState,
    remote_up_down: RemoteUploaderDownloader | None,
    local_checkpoint_path: str,
) -> None:
    """Upload the server state to the S3 Object Store.

    This function saves the server state to a local file and then uploads it to the S3
    Object Store. The server state includes the history (metrics, etc.), current round,
    elapsed time, client states, and server step counts. It creates a local file path
    based on the checkpoint path and current round, dumps the state to this file in
    binary format using pickle, and then uploads the file to S3.

    Parameters
    ----------
    server_state : ServerState
        The state of the server, including history, current round, elapsed time,
        client states, and server step counts.
    remote_up_down : RemoteUploaderDownloader | None
        The uploader/downloader object for interacting with the S3 Object Store
    local_checkpoint_path : str
        The local path where checkpoint files are temporarily stored before upload

    """
    # Create the server state dictionary containing light stuff
    current_server_state = create_server_state_lightweight_dict(server_state)
    check_server_state_lightweight_dict(
        current_server_state,
        raise_if_not_lightweight=True,
    )
    # Dump the server state to disk and upload to the S3 Object Store
    local_path = (
        Path(local_checkpoint_path)
        / "server"
        / f"{server_state.current_round}"
        / CURRENT_SERVER_STATE_FILENAME
    )
    local_path.parent.mkdir(parents=True, exist_ok=True)
    if not local_path.exists():
        with (local_path).open("wb") as f:
            pickle.dump(current_server_state, f)
    if remote_up_down:
        log(DEBUG, "Push server state to S3")
        upload_file_to_s3(
            remote_up_down,
            f"{server_state.current_round}/{CURRENT_SERVER_STATE_FILENAME}",
            local_path,
        )
    else:
        log(DEBUG, "Only local checkpoint is maintained.")


def upload_server_checkpoint_server_state(
    server_state: ServerState,
) -> None:
    """Upload the server state checkpoint to the S3 Object Store.

    This function uploads only the server state (including history and client states)
    to the S3 Object Store for the current round. It retrieves the remote
    uploader/downloader object from the server state and calls `upload_server_state` to
    perform the upload. The function asserts that the remote uploader/downloader is
    present and only uploads the server state if both history and client states are
    available.

    Parameters
    ----------
    server_state : ServerState
        The server state object containing history, client states, and the remote
        uploader/downloader for S3 communication.

    Raises
    ------
        AssertionError
            If the remote uploader/downloader is not present in the server state.

    Notes
    -----
    This function uploads only the server state and does not upload model parameters
    or momentum vectors. Use `upload_server_checkpoint` to upload the full checkpoint.

    """
    # Retrieve remote_up_down object
    remote_up_down = server_state.remote_up_down
    # Uploading the server state
    if server_state.history and server_state.client_states:
        upload_server_state(
            server_state=server_state,
            remote_up_down=remote_up_down,
            local_checkpoint_path=server_state.local_checkpoint_path,
        )


def upload_server_checkpoint_model_states(
    server_state: ServerState,
) -> None:
    """Upload server model states and momentum vectors to the S3 Object Store.

    This function uploads the server's model parameters and momentum vectors (if
    present) to the S3 Object Store for the current round. It retrieves the remote
    uploader/downloader object from the server state and calls the appropriate helper
    functions to perform the upload. The function asserts that the remote
    uploader/downloader is present and uploads the first and second momentum vectors (if
    available), followed by the model parameters.

    Parameters
    ----------
    server_state : ServerState
        The server state object containing model parameters, momentum vectors, and the
        remote uploader/downloader for S3 communication.

    Raises
    ------
        AssertionError
            If the remote uploader/downloader is not present in the server state.

    Notes
    -----
    This function uploads only the model parameters and momentum vectors. It does not
    upload the server state (history, client states, etc.). Use
    `upload_server_checkpoint` to upload the full checkpoint including server state.

    """
    # Retrieve remote_up_down object
    remote_up_down = server_state.remote_up_down

    if server_state.momentum_vector is not None:
        # Have flexibility for handling momenta of momenta.
        generic_model_state_upload_to_s3(
            model_states=server_state.momentum_vector,
            layer_names_and_types=server_state.layer_names_and_types,
            current_round=server_state.current_round,
            remote_up_down=remote_up_down,
            local_checkpoint_path=server_state.local_checkpoint_path,
            parameter_upload=False,
            is_second_momentum=False,
            max_shard_size=server_state.max_shard_size,
        )
    if server_state.second_momentum_vector is not None:
        generic_model_state_upload_to_s3(
            model_states=server_state.second_momentum_vector,
            layer_names_and_types=server_state.layer_names_and_types,
            current_round=server_state.current_round,
            remote_up_down=remote_up_down,
            local_checkpoint_path=server_state.local_checkpoint_path,
            parameter_upload=False,
            is_second_momentum=True,
            max_shard_size=server_state.max_shard_size,
        )
    # Dump and upload model parameters
    generic_model_state_upload_to_s3(
        model_states=server_state.model_states,
        layer_names_and_types=server_state.layer_names_and_types,
        current_round=server_state.current_round,
        remote_up_down=remote_up_down,
        local_checkpoint_path=server_state.local_checkpoint_path,
        parameter_upload=True,
        max_shard_size=server_state.max_shard_size,
    )


def _check_existence_paths_for_parameters(
    base_path_to_check: str,
    check_fn: Callable[[str], bool],
) -> bool:
    """Check if model parameter files exist at a given path using a check function.

    This function verifies the existence of model parameter files (in .bin, .npz, or
    safetensors format) at the specified base path. It uses the provided check function
    to determine file existence. For safetensors, it checks for all required model state
    files.

    Parameters
    ----------
    base_path_to_check : str
        The base path to check for model parameter files.
    check_fn : Callable[[str], bool]
        A function that takes a file path and returns True if the file exists, False
        otherwise.

    Returns
    -------
    bool
        True if any .bin or .npz file exists, or if all required safetensors files
        exist; False otherwise.

    """
    file_name_no_ext = base_path_to_check + f"/{MODEL_PARAMETERS}"
    for suffix in [BIN_SUFFIX, NUMPY_ARRAY_SUFFIX]:
        log(
            DEBUG,
            "Checking file %s",
            file_name_no_ext + suffix,
        )
        if check_fn(
            file_name_no_ext + suffix,
        ):
            return True
    safetensors_files_are_there = True
    # We need to load through all of the states to build up
    # the model parameters and the exponential averages.
    # NOTE: If any of these files doesn't exist, we will return False
    for model_state in [
        ModelStateNames.PARAMETERS.value,
        ModelStateNames.EXP_AVG.value,
        ModelStateNames.EXP_AVG_SQ.value,
    ]:
        state_path = file_name_no_ext + f"/{model_state}"
        log(
            DEBUG,
            "Checking file %s",
            state_path,
        )
        safetensors_files_are_there = safetensors_files_are_there and check_fn(
            state_path,
        )
    return safetensors_files_are_there


def _prepare_model_parameters_paths(
    local_checkpoint_path: str,
    server_checkpoint_path: str,
    remote_up_down: RemoteUploaderDownloader | None,
    timeout: float = 0.5,
) -> None:
    """Ensure model parameter files are available locally or downloading from remote.

    This function checks for the existence of model parameter files at the local
    checkpoint path. If not found, and a remote uploader/downloader is provided, it
    checks the remote path and downloads the files to the local directory. It waits for
    the remote files to become available if necessary.

    Parameters
    ----------
    local_checkpoint_path : str
        The local path to check for model parameter files.
    server_checkpoint_path : str
        The remote server path to check/download model parameter files from.
    remote_up_down : RemoteUploaderDownloader | None
        The uploader/downloader object for interacting with the S3 Object Store.
    timeout : float, optional
        Time to wait between attempts to validate remote paths, by default 0.5.

    Raises
    ------
    ValueError
        If local files are not found and no remote uploader/downloader is provided.

    """
    # If the local paths are not available checks for remote ones
    if not _check_existence_paths_for_parameters(
        local_checkpoint_path,
        lambda x: Path(x).exists(),
    ):
        log(
            DEBUG,
            "No local checkpoint found, looking at remote paths: %s",
            server_checkpoint_path,
        )
        if remote_up_down is None:
            msg = "Local files not found and no uploader object to check remote store."
            raise ValueError(msg)
        if _check_existence_paths_for_parameters(
            server_checkpoint_path,
            validate_given_remote_path,
        ):
            log(
                DEBUG,
                "Found remote checkpoints, downloading...",
            )
            # Check whether the server model states exist
            file_found = False
            remote_file_name_no_ext = server_checkpoint_path + f"/{MODEL_PARAMETERS}"
            # Look for remote paths
            while not file_found:
                file_found = validate_given_remote_path(remote_file_name_no_ext)
                time.sleep(timeout)
            files_to_download = list_remote_objects(
                remote_path=str(remote_file_name_no_ext),
            )
            local_checkpoint = Path(local_checkpoint_path) / f"{MODEL_PARAMETERS}"
            for file_name in files_to_download:
                parsed_file_name = file_name.split(MODEL_PARAMETERS)[-1]
                local_file_name = str(local_checkpoint) + parsed_file_name
                file_to_download = (
                    f"{server_checkpoint_path.rsplit('/', maxsplit=-1)[-1]}/"
                    f"{MODEL_PARAMETERS}" + parsed_file_name
                )
                log(
                    DEBUG,
                    "Downloading remote checkpoint %s to %s",
                    file_to_download,
                    local_file_name,
                )
                download_file_from_s3(remote_up_down, file_to_download, local_file_name)


def get_model_parameters_checkpoint(
    cfg: BaseConfig,
    remote_up_down: RemoteUploaderDownloader | None,
    server_path: str,
    local_checkpoint_path: Path,
    timeout: float = 0.5,
) -> tuple[NDArrays, tuple[tuple[str, ModelStateNames], ...]]:
    """Retrieve model parameters from local storage or S3 Object Store.

    This function loads model parameters from either local storage or the S3 Object
    Store. It first attempts to find parameter files locally, and if not found,
    downloads them from S3. The function supports both .bin and .npz file formats for
    model parameters. It also ensures that the loaded parameters have the same structure
    as the initial parameters expected by the model.

    Parameters
    ----------
    cfg : BaseConfig
        Configuration parameters containing resume round information.
    remote_up_down : RemoteUploaderDownloader
        The uploader/downloader object for interacting with the S3 Object Store.
    server_path : str
        The S3 path where server checkpoint files are stored.
    local_checkpoint_path : Path
        The local path where checkpoint files are stored or downloaded to.
    timeout : float, optional
        Time to wait between attempts to validate remote paths, by default 0.5.

    Returns
    -------
    tuple[NDArrays, tuple[tuple[str, ModelStateNames], ...]]
        A tuple containing:
        - The model parameters loaded from the checkpoint
        - A tuple of (layer_name, layer_type) pairs defining the model state structure

    Notes
    -----
    The function first checks for local parameter files (.bin or .npz) and if not found,
    attempts to download them from S3. It will poll for the availability of remote files
    with a specified timeout interval until they are found.

    """
    # Initialize the model parameters and layer names mapping from a random model
    initial_model_states, layer_names_and_types = get_random_model_states_from_config(
        cfg,
    )
    # Initialize checkpoint model states and layer names
    server_checkpoint_path = server_path + f"{cfg.mt_dao.resume_round}"
    checkpoint_model_states: NDArrays = []
    checkpoint_layer_names_and_types: list[tuple[str, ModelStateNames]] = []

    # Prepare model parameter paths
    _prepare_model_parameters_paths(
        str(local_checkpoint_path),
        server_checkpoint_path,
        remote_up_down,
        timeout,
    )

    # TODO(anon): Discontinue this clause. Check and convert checkpoints first
    # For backward compatibility, check for parameters in the old format
    for suffixes in [BIN_SUFFIX, NUMPY_ARRAY_SUFFIX]:
        path_to_check = Path(local_checkpoint_path) / f"{MODEL_PARAMETERS}{suffixes}"
        if (path_to_check).exists():
            _, checkpoint_model_states = load_model_states_from_file(path_to_check)
            # NOTE: Since these formats don't have the layer names, we assume they are
            # ordered correctly and use those from the randomly initialized model.
            checkpoint_layer_names_and_types = list(layer_names_and_types)
            log(
                DEBUG,
                "Loaded model parameters from path: %s",
                path_to_check,
            )

    # If the model checkpoint states are still empty, check for the latest version
    if not checkpoint_model_states:
        # We need to load through all of the states to build up
        # the model parameters and the exponential averages.
        for model_state in [
            ModelStateNames.PARAMETERS.value,
            ModelStateNames.EXP_AVG.value,
            ModelStateNames.EXP_AVG_SQ.value,
        ]:
            local_path_name = (
                Path(local_checkpoint_path) / f"{MODEL_PARAMETERS}" / f"{model_state}"
            )
            if local_path_name.exists():
                model_states, layer_names = load_chunked_model(local_path_name)
                checkpoint_model_states.extend(model_states)
                checkpoint_layer_names_and_types.extend(
                    (layer_name, ModelStateNames[model_state])
                    for layer_name in layer_names
                )
                log(
                    DEBUG,
                    "Loaded model %s from path: %s",
                    model_state,
                    local_path_name,
                )

    # Final checks before returning the checkpoint
    assert len(initial_model_states) == len(checkpoint_model_states), (
        f"Randomly initialized parameters ({len(initial_model_states)})"
        f" and checkpoints ({len(checkpoint_model_states)}) must have the same length."
    )
    checkpoint_model_states_dict = dict(
        zip(checkpoint_layer_names_and_types, checkpoint_model_states, strict=True),
    )
    initial_model_states_dict = dict(
        zip(layer_names_and_types, initial_model_states, strict=True),
    )
    assert all(
        initial_model_state.shape
        == checkpoint_model_states_dict[layer_name_and_type].shape
        for (
            layer_name_and_type
        ), initial_model_state in initial_model_states_dict.items()
    ), "Randomly initialized parameters and checkpoints must have the same shapes."
    return [
        checkpoint_model_states_dict[layer_name_and_type]
        for layer_name_and_type in layer_names_and_types
    ], layer_names_and_types


def get_server_state_checkpoint(
    cfg: BaseConfig,
    remote_up_down: RemoteUploaderDownloader | None,
    local_checkpoint_path: Path,
) -> tuple[
    WandbHistory,
    int,
    float,
    int,
    dict[str | int, ClientState],
]:
    """Retrieve server state from local storage or S3 Object Store.

    This function loads the server state checkpoint, which contains essential training
    information such as history metrics, current round, elapsed time, and client states.
    It first attempts to load the state from a local file and, if not found, downloads
    it from the S3 Object Store. The function also handles backward compatibility with
    older checkpoint formats by providing default values for missing attributes.

    Parameters
    ----------
    cfg : BaseConfig
        Configuration parameters containing resume round information and federated
        learning settings.
    remote_up_down : RemoteUploaderDownloader
        The uploader/downloader object for interacting with the S3 Object Store.
    local_checkpoint_path : Path
        The local path where checkpoint files are stored or downloaded to.

    Returns
    -------
    tuple[WandbHistory, int, float, int, dict[str | int, ClientState]]
        A tuple containing the following server state components:
        - WandbHistory: The history object containing metrics and other logging data
        - int: The starting training round number
        - float: The elapsed time offset to resume from
        - int: The cumulative server steps count
        - dict[str | int, ClientState]: A dictionary mapping client IDs to their
        respective states

    Raises
    ------
    ValueError
        If remote restoration is requested but no remote uploader/downloader
        object is set.

    Notes
    -----
    This function maintains partial compatibility across different ClientState
    implementations by dynamically determining which attributes to load based on the
    current ClientState class signature. When client state information is missing from
    the checkpoint, the function creates default client states with zero cumulative
    steps.

    """
    # Check for local paths first
    if (local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME).exists():
        with (local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME).open("rb") as f:
            server_state = custom_pickle_load(f)
        log(
            DEBUG,
            "Loaded server state from local path: %s",
            local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME,
        )
    else:
        # Download the server state from S3 Object Store
        if remote_up_down is None:
            msg = "Local files not found and no uploader object to check remote store."
            raise ValueError(msg)
        download_file_from_s3(
            remote_up_down,
            f"{cfg.mt_dao.resume_round}/{CURRENT_SERVER_STATE_FILENAME}",
            str(local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME),
        )
        with (local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME).open("rb") as f:
            server_state = custom_pickle_load(f)
        log(
            DEBUG,
            "Loaded server state from remote path: %s",
            f"{cfg.mt_dao.resume_round}/{CURRENT_SERVER_STATE_FILENAME}",
        )
    # Check the consistency of the server state
    check_server_state_lightweight_dict(
        server_state,
        raise_if_not_lightweight=True,
    )
    # Return the server state components
    return read_server_state_lightweight_dict(
        server_state,
        n_local_steps=cfg.fl.n_local_steps,
        n_total_clients=cfg.fl.n_total_clients,
    )


def get_server_momentum_checkpoint(  # noqa: C901, PLR0912
    cfg: BaseConfig,
    remote_up_down: RemoteUploaderDownloader | None,
    server_path: str,
    local_checkpoint_path: Path,
    momentum_type: str,
) -> NDArrays | None:
    """Retrieve momentum vector from local storage or S3 Object Store.

    This function attempts to load a momentum vector (either first or second) from
    available  sources. It first checks for a local server state file and, if not found,
    downloads it from  the S3 Object Store. It then tries to extract the momentum vector
    from this server state. If the vector is not present in the server state, the
    function looks for a dedicated momentum vector file locally or in S3, downloading it
    if necessary.

    Parameters
    ----------
    cfg : BaseConfig
        Configuration parameters containing S3 bucket information, run UUID, and
        resume round information.
    remote_up_down : RemoteUploaderDownloader
        The uploader/downloader object for interacting with the S3 Object Store.
    server_path : str
        The S3 path where server checkpoint files are stored.
    local_checkpoint_path : Path
        The local path where checkpoint files are stored or downloaded to.
    momentum_type : str
        The type of momentum vector to retrieve (FIRST_MOMENTUM or SECOND_MOMENTUM).

    Returns
    -------
    NDArrays | None
        The retrieved momentum vector as NDArrays, or None if no momentum vector was
        found.

    Raises
    ------
    ValueError
        If files are to be downloaded from S3 but no remote backend is found.

    Notes
    -----
    The function prioritizes finding the momentum vector in the following order:
    1. In the server state file (if it contains a "momentum" field)
    2. In a local file dedicated to the specific momentum type
    3. In a remote S3 file dedicated to the specific momentum type

    """
    # Initialize the important paths and variables
    momentum_vector: NDArrays | None = None
    local_file_name = Path(local_checkpoint_path) / f"{momentum_type}"
    remote_file_name_momentum = (
        server_path + f"{cfg.mt_dao.resume_round}/{momentum_type}"
    )

    # TODO(anon): Discontinue this clause. Check and convert checkpoints first
    # For backward compatibility, check for momentum in the state object instead of
    # separate files/folders
    if (local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME).exists():
        with (local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME).open("rb") as f:
            server_state = custom_pickle_load(f)
    else:
        # Download the server state from S3 Object Store
        if remote_up_down is None:
            msg = "No local files and no uploader object to check remote store."
            raise ValueError(msg)
        download_file_from_s3(
            remote_up_down,
            f"{cfg.mt_dao.resume_round}/{CURRENT_SERVER_STATE_FILENAME}",
            str(local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME),
        )
        log(DEBUG, "Read server state from S3")
        with (local_checkpoint_path / CURRENT_SERVER_STATE_FILENAME).open("rb") as f:
            server_state = custom_pickle_load(f)
    if "momentum" in server_state:
        log(DEBUG, "Get momentum vector from server state")
        momentum_vector = cast("NDArrays", server_state["momentum"])

    # Check the local path first or download from remote
    if not local_file_name.exists() and validate_given_remote_path(
        remote_file_name_momentum,
    ):
        # Download the files from remote
        if remote_up_down is None:
            msg = "No local files and no uploader object to check remote store."
            raise ValueError(msg)
        # Set the file names depending on the extension found
        files_to_download = list_remote_objects(
            remote_path=str(remote_file_name_momentum),
        )
        for file_name in files_to_download:
            parsed_file_name = file_name.split(momentum_type)[-1]
            file_to_download = (
                f"{cfg.mt_dao.resume_round}/{momentum_type}{parsed_file_name}"
            )
            download_file_name = (
                str(local_checkpoint_path) + f"/{momentum_type}" + parsed_file_name
            )
            download_file_from_s3(
                remote_up_down,
                file_to_download,
                download_file_name,
            )
        log(DEBUG, "Downloaded momentum files from S3 complete.")
    # If the local file already existed now we are loading it. If the local files didn't
    # exist, but the remote did, the local path now contains the remote files and we are
    # loading them.
    if local_file_name.exists():
        # Check for local paths first with backward compatibility
        for suffixes in [BIN_SUFFIX, NUMPY_ARRAY_SUFFIX]:
            path_to_check = Path(str(local_file_name) + suffixes)
            if path_to_check.exists():
                _, momentum_vector = load_model_states_from_file(
                    path_to_check,
                )
                log(
                    DEBUG,
                    "Loaded optimizer states from local path: %s",
                    path_to_check,
                )
        if not momentum_vector:
            tmp_momentum_vector: NDArrays = []
            for model_state in [
                ModelStateNames.PARAMETERS.value,
                ModelStateNames.EXP_AVG.value,
                ModelStateNames.EXP_AVG_SQ.value,
            ]:
                loading_path = Path(local_file_name) / f"{model_state}"
                if loading_path.exists():
                    model_states, _ = load_chunked_model(loading_path)
                    tmp_momentum_vector.extend(model_states)
                    log(
                        DEBUG,
                        "Loaded %s %s from path: %s",
                        momentum_type,
                        model_state,
                        loading_path,
                    )
            momentum_vector = tmp_momentum_vector or momentum_vector
    # If the momentum vector is empty now, it means that no local or remote files were
    # found and we couldn't load it
    if not momentum_vector:
        log(
            WARNING,
            "No %s momentum vector found. Initializing it with zero values",
            momentum_type,
        )
    return momentum_vector


def get_server_checkpoint(
    cfg: BaseConfig,
    remote_up_down: RemoteUploaderDownloader | None,
    mt_dao_save_path: str,
    timeout: float = 0.5,
) -> ServerState:
    """Load server checkpoint from S3 or local storage.

    This function retrieves all necessary components of a server checkpoint, including
    model parameters, server state information, and momentum vectors. It first attempts
    to find checkpoint files locally, and if not found, downloads them from the S3
    Object Store. The components are then assembled into a ServerState object for
    resuming federated learning training from a previous checkpoint.

    Parameters
    ----------
    cfg : BaseConfig
        Configuration parameters containing S3 bucket information, run UUID, and
        resume round information.
    remote_up_down : RemoteUploaderDownloader
        The uploader/downloader object for interacting with the S3 Object Store.
    mt_dao_save_path : str
        The local path where checkpoint files are stored or downloaded to.
    timeout : float, optional
        Time to wait between attempts to validate remote paths, by default 0.5.

    Returns
    -------
    ServerState
        A fully initialized ServerState object containing the model parameters,
        training history, momentum vectors, client states, and other necessary
        information to resume training.

    Notes
    -----
    The function loads multiple components:
    - Model parameters and layer names/types
    - Server state (history, start round, time offset, client state)
    - First momentum vector (if available)
    - Second momentum vector (if available)

    Each component is loaded by calling specialized helper functions that handle both
    local and remote checkpoint retrieval.

    """
    # Set the path to server checkpoints
    server_path = f"s3://{cfg.s3_comm_config.bucket_name}/{cfg.run_uuid}/server/"
    # Get the local checkpoint path
    local_checkpoint_path = (
        Path(mt_dao_save_path) / "server" / f"{cfg.mt_dao.resume_round}"
    )
    # Model states
    model_states, layer_names_and_types = get_model_parameters_checkpoint(
        cfg=cfg,
        remote_up_down=remote_up_down,
        server_path=server_path,
        local_checkpoint_path=local_checkpoint_path,
        timeout=timeout,
    )
    # Server state
    (
        history,
        start_round,
        time_offset,
        server_steps_cumulative,
        client_states,
    ) = get_server_state_checkpoint(
        cfg=cfg,
        remote_up_down=remote_up_down,
        local_checkpoint_path=local_checkpoint_path,
    )

    momentum_vector: NDArrays | None = None
    second_momentum_vector: NDArrays | None = None
    # NOTE: We should unify the state with the strategy
    # Only focusing on saving space right now
    if cfg.fl.strategy_name.lower() in {
        StrategyName.FEDMOM,
        StrategyName.NESTEROV,
        StrategyName.FEDADAM,
        StrategyName.FEDYOGI,
    }:
        # First momentum vector
        momentum_vector = get_server_momentum_checkpoint(
            cfg=cfg,
            remote_up_down=remote_up_down,
            server_path=server_path,
            local_checkpoint_path=local_checkpoint_path,
            momentum_type=FIRST_MOMENTUM,
        )
        momentum_vector = (
            [np.zeros_like(x) for x in model_states]
            if momentum_vector is None
            else momentum_vector
        )
    if cfg.fl.strategy_name.lower() in {StrategyName.FEDADAM, StrategyName.FEDYOGI}:
        # Second momentum vector
        second_momentum_vector = get_server_momentum_checkpoint(
            cfg=cfg,
            remote_up_down=remote_up_down,
            server_path=server_path,
            local_checkpoint_path=local_checkpoint_path,
            momentum_type=SECOND_MOMENTUM,
        )
        second_momentum_vector = (
            [np.zeros_like(x) for x in model_states]
            if second_momentum_vector is None
            else second_momentum_vector
        )

    log(INFO, "Server checkpoint loaded")
    return ServerState(
        model_states=model_states,
        transmission_scheduler=dispatch_model_state_scheduler(cfg),
        aggregation_mask_scheduler=dispatch_model_state_scheduler(cfg),
        history=history,
        time_offset=time_offset,
        server_steps_cumulative=server_steps_cumulative,
        client_states=client_states,
        momentum_vector=momentum_vector,
        second_momentum_vector=second_momentum_vector,
        state_keys=dispatch_strategy_state_keys(cfg),
        sampled_clients=[],
        previously_sampled_clients=[],
        rng=random.Random(cfg.seed),  # noqa: S311
        start_round=start_round,
        current_round=start_round,
        total_rounds=cfg.fl.n_rounds,
        perf_counter=time.perf_counter(),
        layer_names_and_types=layer_names_and_types,
        remote_up_down=remote_up_down,
        local_checkpoint_path=mt_dao_save_path,
        client_sampler=ClientSampler(
            number_of_clients_per_round=cfg.fl.n_clients_per_round,
            total_number_of_clients=cfg.fl.n_total_clients,
            dropout_ratio=cfg.fl.dropout_ratio,
            dropout_function_name=cfg.fl.dropout_function_name,
        ),
        max_shard_size=cfg.mt_dao.max_shard_size,
    )


def load_pretrained_model_from_path(
    pretrained_model_path: str,
    run_uuid: str,
    s3_comm_config: S3CommConfig,
    trainer: Trainer,
) -> None:
    """Load a pretrained model from a specified path and set it to the trainer.

    This function supports loading models from both local file paths and S3 URIs.
    It downloads the model if the path points to an S3 bucket, and then sets the
    parameters in the provided trainer object.

    Parameters
    ----------
    pretrained_model_path : str
        Path to pre-trained model
    run_uuid : str
        The unique identifier for the run, used for S3 operations.
    s3_comm_config : S3CommConfig
        The S3 communication configuration.
    trainer : Trainer
        The trainer object to which the pretrained model parameters are set.

    """
    log(
        INFO,
        "Loading pretrained model from %s",
        pretrained_model_path,
    )
    # Create a temporary directory for storing the downloaded parameters
    with TemporaryDirectory() as tmp_dir:
        # Load the local file path
        local_file_path = get_file_from_path(
            output_parent_path=tmp_dir,
            input_file_path=pretrained_model_path,
            run_uuid=run_uuid,
            s3_comm_config=s3_comm_config,
        )
        _, initial_parameters = load_model_states_from_file(local_file_path)
        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  # noqa: SLF001
            if isinstance(  # pyright: ignore[reportAttributeAccessIssue]
                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,
            )
        initial_parameters_names = list(
            tensors_dict.keys(),
        )
        set_trainer_tensors_dict(
            trainer=trainer,
            tensors_dict=create_tensors_dict(
                parameters=initial_parameters,
                parameters_names=initial_parameters_names,
            ),
        )


def delete_past_communication_states(
    bucket_name: str,
    prefix: str,
) -> None:
    """Delete previous communication state files from S3 storage.

    This function removes all .npz files from the communication stack folder in the
    specified S3 bucket and prefix path. It first lists all objects in the communication
    stack folder, filters for .npz files, and then deletes each matching file. This is
    typically used to clean up stale communication state files after they are no longer
    needed, preventing unnecessary storage usage and costs.

    Parameters
    ----------
    bucket_name : str
        The name of the S3 bucket containing the communication state files.
    prefix : str
        The prefix (folder path) within the bucket where communication state files
        are stored.

    Notes
    -----
    This function only removes files with the .npz extension, which are typically
    NumPy compressed array files used for storing model parameters or other numeric
    data during federated learning communications.

    """
    # Compose remote URI
    remote_uri = f"s3://{bucket_name}/{prefix}/comm_stack"
    # List all the object in the communication stack folder
    remote_objects = list_remote_objects(remote_uri)
    # Remove non-.npz paths from the list of files
    remote_objects_to_remove = [
        cro for cro in remote_objects if cro.endswith(NUMPY_ARRAY_SUFFIX)
    ]
    for object_to_remove in remote_objects_to_remove:
        # Parse the URI to extract the backend and bucket name
        backend, bucket_name, _prefix = parse_uri(remote_uri)
        delete_object(f"{backend}://{bucket_name}/{object_to_remove}")


def copy_old_checkpoints_to_new_run(  # noqa: C901, PLR0912, PLR0913, PLR0914, PLR0915
    remote_up_down: RemoteUploaderDownloader | None,
    mt_dao_save_path: str,
    bucket_uri: str,
    restore_run_data: tuple[str, str, int | None, int],
    n_total_clients: int | None,
    *,
    copy_client_checkpoints: bool = True,
) -> None:
    """Copy old checkpoints to the new run folder.

    Parameters
    ----------
    remote_up_down : RemoteUploaderDownloader | None
        The remote uploader/downloader object.
    mt_dao_save_path: str
        The path where local checkpoints are saved to by default.
    bucket_uri : str
        The URI of the bucket.
    restore_run_data : tuple[str, str, int | None, int]
        A tuple containing the run UUID, restore run UUID, restore run round, and
        restore run step.
    n_total_clients : int | None
        The total number of clients expected in the new run.
    copy_client_checkpoints : bool, optional
        A flag indicating whether to copy the client checkpoints. Defaults to True.

    Raises
    ------
        NotImplementedError: If the backend is not an S3ObjectStore.
        ValueError: If the old run folder or the new run folder is not found.

    """
    run_uuid, restore_run_uuid, restore_run_round, restore_run_step = restore_run_data
    build_from_remote = False
    if remote_up_down:
        if not isinstance(remote_up_down.remote_backend, S3ObjectStore):
            msg = "Support for resuming from non-S3 backends is not yet implemented."
            raise NotImplementedError(
                msg,
            )
        log(DEBUG, "Enabling copying on remote S3 Backend.")
        build_from_remote = True
    else:
        log(DEBUG, "Copying from locally saved checkpoints.")
    if build_from_remote:
        new_run_folder = bucket_uri + f"/{run_uuid}"
        old_run_folder = bucket_uri + f"/{restore_run_uuid}"
        old_run_val = validate_given_remote_path(old_run_folder)
        validation = (old_run_val) and (validate_given_remote_path(new_run_folder))
    else:
        # NOTE: the mt_dao_path will by convention have the run_uuid specified.
        # as such we can just replace this with the restore_uuid to get the correct
        # local path.
        new_run_folder = mt_dao_save_path
        old_run_folder = mt_dao_save_path.replace(run_uuid, restore_run_uuid)
        _, old_run_val = list_objects(old_run_folder)  # type: ignore[assignment]
        _, new_run_val = list_objects(new_run_folder)
        validation = old_run_val and new_run_val  # type: ignore[assignment]

    if validation:
        state_bin = (
            old_run_folder
            + f"/server/{restore_run_round}/{CURRENT_SERVER_STATE_FILENAME}"
        )

        # NOTE: every one of these have their own subdirectories potentially.
        first_momentum_vec = (
            old_run_folder + f"/server/{restore_run_round}/{FIRST_MOMENTUM}"
        )

        second_momentum_vec = (
            old_run_folder + f"/server/{restore_run_round}/{SECOND_MOMENTUM}"
        )

        parameters = old_run_folder + f"/server/{restore_run_round}/{MODEL_PARAMETERS}"

        # We assume that we have a directory now that has chunked the updates.
        # Extract the client and the batches
        # NOTE: (?:\d+) means a do-not-capture group
        # As such we allow any number of epochs without extracting
        # The number of epochs
        _, potential_client_paths = list_objects(old_run_folder)

        client_path_batches = sorted(
            [
                (
                    path,
                    int(reg.group(1)),
                    int(reg.group(2)),
                )
                for path in potential_client_paths
                if (reg := re.search(r"client_(\d+)/ep(?:\d+)-ba(\d+)", path))
                is not None
            ],
            key=operator.itemgetter(1, 2),
        )

        # For each client, choose the latest checkpoint
        # That is consistent with the step of the resume round
        # groupby acts like an sql groupby
        client_paths = [
            list(filter(lambda x: x[2] <= restore_run_step, group))[-1][0]
            for _, group in groupby(client_path_batches, key=operator.itemgetter(1))
        ]

        if (
            copy_client_checkpoints
            and n_total_clients is not None
            and (found_clients := len(client_paths)) != n_total_clients
        ):
            msg = (
                f"Found {found_clients} clients in the old run folder {old_run_folder},"
                f" but expected {n_total_clients}."
            )
            raise ValueError(
                msg,
            )

        paths_to_copy = [state_bin, parameters]
        for momentum_vec in [first_momentum_vec, second_momentum_vec]:
            _, contents = list_objects(momentum_vec)
            if contents:
                paths_to_copy.append(momentum_vec)
            else:
                log(
                    DEBUG,
                    f"Could not find momentum vector to copy from {momentum_vec}",
                )

        if copy_client_checkpoints:
            paths_to_copy.extend(client_paths)
        # NOTE: We have to enumerate all the objects within each directory
        # because each of the above paths that we have specified can potentially
        # have subdirectories.
        for potential_directory in paths_to_copy:
            _, files_to_download = list_objects(potential_directory)
            if not files_to_download:
                # NOTE: This occurs in the case that the potential directory was a
                # local file. The work around is acheived in this was so that it
                # works well with the remote version of the list_objects function.
                files_to_download = [potential_directory]
            for path in files_to_download:
                target_key = path.replace(restore_run_uuid, run_uuid)
                log(DEBUG, "Copying %s to %s", path, target_key)
                if remote_up_down and isinstance(
                    remote_up_down.remote_backend,
                    S3ObjectStore,
                ):
                    bucket = remote_up_down.remote_backend.bucket
                    if key_exists(
                        client=remote_up_down.remote_backend.client,
                        bucket=bucket,
                        key=target_key,
                    ):
                        log(DEBUG, f"{target_key} already exists in {bucket}")
                    else:
                        remote_up_down.remote_backend.client.copy(
                            {"Bucket": bucket, "Key": path},
                            bucket,
                            target_key,
                        )
                else:
                    target_path = Path(target_key)
                    if target_path.exists():
                        # NOTE: We do not want to inadvertantly overwrite a
                        # local directory.
                        log(DEBUG, f"The file {target_path} already exists.")
                    else:
                        target_path.parent.mkdir(parents=True, exist_ok=True)
                        shutil.copy2(path, target_path)

    else:
        if not old_run_val:
            msg = (
                f"Could not find the old run folder {old_run_folder} to copy"
                " checkpoints."
            )
            raise ValueError(
                msg,
            )
        msg = f"Could not find the new run folder {new_run_folder} to copy checkpoints."
        raise ValueError(
            msg,
        )


def cleanup_checkpoints(
    server_state: ServerState,
    run_uuid: str | None,
    end_idx: int | None = -1,
    *,
    local_cleanup: bool = False,
    remote_cleanup: bool = False,
) -> None:
    """Clean up checkpoints from S3 Object Store and local storage.

    This function handles cleanup of both client and server checkpoints. It removes old
    checkpoints from the S3 Object Store based on the provided end_idx parameter, which
    specifies which checkpoints to keep. It also cleans up the local checkpoint path to
    free up disk space. This function is useful for periodic cleanup during training to
    avoid excessive storage usage.

    Parameters
    ----------
    server_state : ServerState
        The server state object containing the local checkpoint path.
    run_uuid : str | None
        The unique identifier for the run whose checkpoints need to be cleaned up
        if supplied, else None.
    end_idx : int | None, optional
        The index up to which checkpoints should be deleted. Defaults to -1, which
        means all checkpoints except for the latest one will be removed.
    local_cleanup: bool
        Whether to clean the local checkpoint directory.
    remote_cleanup: bool
        Whether to clean up the remote checkpoint directory.

    Raises
    ------
    ValueError
        If remote deletion is requested but the server state's remote
        uploader/downloader is not set.

    Notes
    -----
    This function performs cleanup in both remote S3 storage and local storage. It uses
    the underlying functions `delete_clients_checkpoints` and `delete_rounds` to perform
    the actual deletion operations.

    """
    # Clean up checkpoints if asked to

    # Remove old clients checkpoints from the S3 Object Store
    if remote_cleanup:
        if server_state.remote_up_down:
            bucket_name = server_state.remote_up_down.remote_bucket_name  # type: ignore[union-attr]
            delete_clients_checkpoints(
                run_uuid_path=f"s3://{bucket_name}/{run_uuid}",
                end_idx=end_idx,
            )
            # Remove old server checkpoints from the S3 Object Store
            delete_rounds(
                run_uuid_path=f"s3://{bucket_name}/{run_uuid}",
                state_keys=server_state.state_keys,
                catch_all=True,
                end_idx=end_idx,
            )
        else:
            msg = "Cannot delete remote without the valid bucket name"
            raise ValueError(msg)
    if local_cleanup:
        # Remove client checkpoints locally
        delete_clients_checkpoints(
            run_uuid_path=server_state.local_checkpoint_path,
            end_idx=end_idx,
        )
        # Remove old server checkpoints from the local directory.
        delete_rounds(
            run_uuid_path=server_state.local_checkpoint_path,
            state_keys=server_state.state_keys,
            catch_all=True,
            end_idx=end_idx,
        )
