"""Centralised training script for LLMFoundry models.

Slightly adapted from the original https://github.com/mosaicml/llm-foundry/blob/25599294c942cfed2c6f8329e14791e4a2f91539/scripts/train/train.py
Copyright 2022 MosaicML LLM Foundry authors
SPDX-License-Identifier: Apache-2.0
"""

import copy
import gc
import os
from collections import OrderedDict
from logging import DEBUG, INFO
from pathlib import Path
from typing import cast

import torch
from composer import Trainer
from flwr.common import log
from llmfoundry.callbacks import EvalGauntlet
from omegaconf import OmegaConf

from mt_dao.clients.configs import CentralizedConfig
from mt_dao.clients.llm_client_functions import (
    get_trainer_object,
)
from mt_dao.composer_trainer_utils.composer_trainer_getters import (
    get_parameters_from_trainer,
    get_wte_parameters_from_trainer,
)
from mt_dao.composer_trainer_utils.composer_trainer_setters import (
    set_trainer_tensors_dict,
)
from mt_dao.conf.base_schema import BaseConfig
from mt_dao.constants import (
    CENTRALIZED_LOG_SUFFIX,
    CONFIG_FILE_YAML,
    NUMPY_ARRAY_SUFFIX,
    mt_dao_SAVE_PATH,
)
from mt_dao.file_utils import dump_model_parameters_to_file
from mt_dao.initialization_utils.helper import (
    extract_layer_names_and_experts,
    retrieve_checkpoint_states,
    retrieve_model_states_and_types,
)
from mt_dao.masks_utils import ModelStateNames
from mt_dao.utils import create_tensors_dict


def initialize_checkpoint_state(
    trainer: Trainer,
    cfg: BaseConfig,
    checkpoint_dir: str,
) -> None:
    """Initialize the checkpoint state into the Trainer object.

    A heper method that downloads checkpoints from a particular
    checkpoint directory and initialises a Trainer obkect
    with the correct state.

    Parameters
    ----------
    trainer : Trainer
        The trainer object.
    cfg : BaseConfig
        The configuration for this setup.
    checkpoint_dir : str
        The directory where the checkpoint is located.


    """
    checkpoint_information = retrieve_checkpoint_states(
        checkpoint_dir=checkpoint_dir,
        s3_comm_config=cfg.s3_comm_config,
        run_uuid=cfg.run_uuid,
        desired_steps=-1,
    )
    (local_checkpoint_directory, tmp_dir, remote_or_local, checkpoint_suffixes) = (
        checkpoint_information
    )
    param_names, n_experts = extract_layer_names_and_experts(trainer=trainer)
    model_states, model_names_and_types = retrieve_model_states_and_types(
        local_checkpoint_directory=local_checkpoint_directory,
        checkpoint_suffixes=checkpoint_suffixes,
        preserve_optimizer_states=False,
        desired_steps=-1,  # NOTE: We choose the last checkpoint if of .pt type in a dir
        param_names=param_names,
        n_experts=n_experts,
    )
    # NOTE: Filtering out the optimizer states manually.
    initial_parameters = []
    parameter_names = []
    for state, (layer_name, layer_type) in zip(
        model_states, model_names_and_types, strict=True,
    ):
        if layer_type == ModelStateNames.PARAMETERS:
            initial_parameters.append(state)
            parameter_names.append(layer_name)
    assert len(model_states) == 3 * len(initial_parameters)
    set_trainer_tensors_dict(
        trainer=trainer,
        tensors_dict=create_tensors_dict(
            parameters=initial_parameters,
            parameters_names=parameter_names,
        ),
    )
    if remote_or_local and tmp_dir:
        log(DEBUG, "Cleaning up the temporary directory.")
        tmp_dir.cleanup()


def main() -> Trainer:  # noqa: C901
    """Execute centralized training.

    This function sets up the training environment, loads configurations, initializes
    the trainer, optionally loads pretrained models, performs evaluation if requested,
    and starts the training process. It also handles dumping model parameters to files
    before and after training.

    Returns
    -------
        Trainer
            The initialized trainer object.

    Raises
    ------
    ValueError
        If the environmental variable mt_dao_SAVE_PATH is not set.

    """
    # Get the environmental variable for the dump folder
    mt_dao_save_path = os.environ.get(mt_dao_SAVE_PATH, "")
    # Raise an error if the environmental variable is not set
    if not mt_dao_save_path:
        msg = "The environmental variable mt_dao_SAVE_PATH is not set."
        raise ValueError(msg)
    # Load the configuration from the config file
    cfg_ = cast("BaseConfig", OmegaConf.load(mt_dao_save_path + f"/{CONFIG_FILE_YAML}"))
    # Resolve all interpolation variables as early as possible
    OmegaConf.resolve(cfg_)
    OmegaConf.set_struct(cfg_, value=False)
    cfg = cfg_.llm_config
    # Resolve all interpolation variables as early as possible
    OmegaConf.resolve(cfg)
    OmegaConf.set_struct(cfg, value=False)
    # Creating ClientConfig object
    client_config = CentralizedConfig(
        cid=cfg_.centralized.stream_id,
        allow_unigram_metrics_failures=cfg_.fl.allow_unigram_metrics_failures,
        resize_vocab=cfg_.fl.resize_vocab,
        split_eval=cfg_.centralized.split_eval,
        use_unigram_metrics=cfg_.fl.use_unigram_metrics,
        s3_comm_config=cfg_.s3_comm_config,
    )

    # NOTE: The cid passed her is use to appoint the position for the stream used for
    # creating the streaming dataset object
    log(
        INFO,
        "Creating trainer object using stream_id: %s...",
        cfg_.centralized.stream_id,
    )
    trainer, train_config, *_ = get_trainer_object(
        client_config=client_config,
        cfg=cfg,
        log_name=CENTRALIZED_LOG_SUFFIX,
        split_eval=client_config.split_eval,
    )
    torch.cuda.empty_cache()
    gc.collect()

    wte_parameters_dict: OrderedDict[str, torch.Tensor] | None = None
    if cfg_.wte_parameters_path:
        initialize_checkpoint_state(
            trainer=trainer,
            cfg=copy.deepcopy(cfg_),
            checkpoint_dir=cfg_.wte_parameters_path,
        )
        wte_parameters_name, wte_parameters = get_wte_parameters_from_trainer(trainer)
        wte_parameters_dict = OrderedDict(
            {
                wte_parameters_name: torch.as_tensor(wte_parameters),
            },
        )
    if cfg_.pretrained_model_path:
        initialize_checkpoint_state(
            trainer=trainer,
            cfg=copy.deepcopy(cfg_),
            checkpoint_dir=cfg_.pretrained_model_path,
        )

    if wte_parameters_dict is not None:
        set_trainer_tensors_dict(trainer, wte_parameters_dict)

    # Eval first if requested
    if train_config.eval_first:
        trainer.eval()
        eval_gauntlet_callback: EvalGauntlet | None = None
        for callback in trainer.state.callbacks:
            if isinstance(callback, EvalGauntlet):
                eval_gauntlet_callback = callback
        if eval_gauntlet_callback is not None:
            assert isinstance(eval_gauntlet_callback, EvalGauntlet)
            composite_scores = eval_gauntlet_callback.eval_after_all(
                trainer.state,
                trainer.logger,
            )
            log(
                INFO,
                "Evaluated model with the Gauntlet before training: %s",
                composite_scores,
            )

    # Dump model parameters to file
    if cfg_.centralized.store_init_model:
        # Get model parameters from trainer object
        model_parameters = get_parameters_from_trainer(trainer)
        # Get number of steps executed
        n_steps = trainer.state.timestamp.batch.value
        # Dump the compressed model parameters to file
        dump_model_parameters_to_file(
            model_parameters=model_parameters,
            file_path=Path(f"{cfg_.run_uuid}-{n_steps}-checkpoint{NUMPY_ARRAY_SUFFIX}"),
        )
        log(INFO, "Dumped initial model parameters to file.")

    if not cfg_.centralized.eval_only:
        log(INFO, "Starting training...")
        trainer.fit(reset_time=cfg_.centralized.reset_timestamp)

    # Dump model parameters to file
    if cfg_.centralized.store_final_model:
        # Get model parameters from trainer object
        model_parameters = get_parameters_from_trainer(trainer)
        # Get number of steps executed
        n_steps = trainer.state.timestamp.batch.value
        # Dump the compressed model parameters to file
        dump_model_parameters_to_file(
            model_parameters=model_parameters,
            file_path=Path(f"{cfg_.run_uuid}-{n_steps}-checkpoint{NUMPY_ARRAY_SUFFIX}"),
        )
        log(INFO, "Dumped final model parameters to file.")

    log(INFO, "Done.")
    return trainer


if __name__ == "__main__":
    main()
