"""Function that trains models in the Parametric UMAP family."""

from __future__ import annotations

import datetime
import logging
import pathlib

import lightning
import pandas as pd
import torch.utils.data

from .. import datapipes, models
from ..common import validators
from ..common.paths import generate_results_dir_name_as_path
from ..datapipes.parametric_umap import Datapipe

logger = logging.getLogger(__name__)


def get_split_dur(df: pd.DataFrame, split: str) -> float:
    """Get duration of a split in a dataset from a pandas DataFrame representing the dataset."""
    return df[df["split"] == split]["duration"].sum()


def get_trainer(
    accelerator: str,
    devices: int | list[int],
    max_epochs: int,
    ckpt_root: str | pathlib.Path,
    ckpt_step: int,
    log_save_dir: str | pathlib.Path,
) -> lightning.pytorch.Trainer:
    """Returns an instance of ``lightning.pytorch.Trainer``
    with a default set of callbacks.
    Used by ``vak.core`` functions."""
    ckpt_callback = lightning.pytorch.callbacks.ModelCheckpoint(
        dirpath=ckpt_root,
        filename="checkpoint",
        every_n_train_steps=ckpt_step,
        save_last=True,
        verbose=True,
    )
    ckpt_callback.CHECKPOINT_NAME_LAST = "checkpoint"
    ckpt_callback.FILE_EXTENSION = ".pt"

    val_ckpt_callback = lightning.pytorch.callbacks.ModelCheckpoint(
        monitor="val_loss",
        dirpath=ckpt_root,
        save_top_k=1,
        mode="min",
        filename="min-val-loss-checkpoint",
        auto_insert_metric_name=False,
        verbose=True,
    )
    val_ckpt_callback.FILE_EXTENSION = ".pt"

    callbacks = [
        ckpt_callback,
        val_ckpt_callback,
    ]

    logger = lightning.pytorch.loggers.TensorBoardLogger(save_dir=log_save_dir)

    trainer = lightning.pytorch.Trainer(
        max_epochs=max_epochs,
        accelerator=accelerator,
        devices=devices,
        logger=logger,
        callbacks=callbacks,
    )
    return trainer


def train_parametric_umap_model(
    model_config: dict,
    dataset_config: dict,
    trainer_config: dict,
    batch_size: int,
    num_epochs: int,
    num_workers: int,
    checkpoint_path: str | pathlib.Path | None = None,
    root_results_dir: str | pathlib.Path | None = None,
    results_path: str | pathlib.Path | None = None,
    shuffle: bool = True,
    val_step: int | None = None,
    ckpt_step: int | None = None,
    subset: str | None = None,
) -> None:
    """Train a model from the parametric UMAP family
    and save results.

    Saves checkpoint files for model,
    label map, and spectrogram scaler.
    These are saved either in ``results_path``
    if specified, or a new directory
    made inside ``root_results_dir``.

    Parameters
    ----------
    model_config : dict
        Model configuration in a :class:`dict`.
        Can be obtained by calling :meth:`vak.config.ModelConfig.asdict`.
    dataset_config: dict
        Dataset configuration in a :class:`dict`.
        Can be obtained by calling :meth:`vak.config.DatasetConfig.asdict`.
    trainer_config: dict
        Configuration for :class:`lightning.pytorch.Trainer` in a :class:`dict`.
        Can be obtained by calling :meth:`vak.config.TrainerConfig.asdict`.
    batch_size : int
        number of samples per batch presented to models during training.
    num_epochs : int
        number of training epochs. One epoch = one iteration through the entire
        training set.
    num_workers : int
        Number of processes to use for parallel loading of data.
        Argument to torch.DataLoader.
    checkpoint_path : str, pathlib.Path, optional
        path to a checkpoint file,
        e.g., one generated by a previous run of ``vak.core.train``.
        If specified, this checkpoint will be loaded into model.
        Used when continuing training.
        Default is None, in which case a new model is initialized.
    root_results_dir : str, pathlib.Path, optional
        Root directory in which a new directory will be created
        where results will be saved.
    results_path : str, pathlib.Path, optional
        Directory where results will be saved.
        If specified, this parameter overrides ``root_results_dir``.
    val_step : int
        Computes the loss using validation set every ``val_step`` epochs.
        Default is None, in which case no validation is done.
    ckpt_step : int
        Step on which to save to checkpoint file.
        If ckpt_step is n, then a checkpoint is saved every time
        the global step / n is a whole number, i.e., when ckpt_step modulo the global step is 0.
        Default is None, in which case checkpoint is only saved at the last epoch.
    shuffle: bool
        if True, shuffle training data before each epoch. Default is True.
    split : str
        Name of split from dataset found at ``dataset_path`` to use
        when training model. Default is 'train'. This parameter is used by
        `vak.learncurve.learncurve` to specify specific subsets of the
        training set to use when training models for a learning curve.
    """
    for path, path_name in zip(
        (checkpoint_path,),
        ("checkpoint_path",),
    ):
        if path is not None:
            if not validators.is_a_file(path):
                raise FileNotFoundError(
                    f"value for ``{path_name}`` not recognized as a file: {path}"
                )

    dataset_path = pathlib.Path(dataset_config["path"])
    if not dataset_path.exists() or not dataset_path.is_dir():
        raise NotADirectoryError(
            f"`dataset_path` not found or not recognized as a directory: {dataset_path}"
        )

    logger.info(
        f"Loading dataset from path: {dataset_path}",
    )
    metadata = datapipes.parametric_umap.Metadata.from_dataset_path(
        dataset_path
    )
    dataset_csv_path = dataset_path / metadata.dataset_csv_filename
    dataset_df = pd.read_csv(dataset_csv_path)
    # ---------------- pre-conditions ----------------------------------------------------------------------------------
    if val_step and not dataset_df["split"].str.contains("val").any():
        raise ValueError(
            f"val_step set to {val_step} but dataset does not contain a validation set; "
            f"please run `vak prep` with a config.toml file that specifies a duration for the validation set."
        )

    # ---- set up directory to save output -----------------------------------------------------------------------------
    if results_path:
        results_path = pathlib.Path(results_path).expanduser().resolve()
        if not results_path.is_dir():
            raise NotADirectoryError(
                f"results_path not recognized as a directory: {results_path}"
            )
    else:
        results_path = generate_results_dir_name_as_path(root_results_dir)
        results_path.mkdir()

    # ---------------- load training data  -----------------------------------------------------------------------------
    logger.info(f"using training dataset from {dataset_path}")
    # below, if we're going to train network to predict unlabeled segments, then
    # we need to include a class for those unlabeled segments in labelmap,
    # the mapping from labelset provided by user to a set of consecutive
    # integers that the network learns to predict
    train_dur = get_split_dur(dataset_df, "train")
    logger.info(
        f"Total duration of training split from dataset (in s): {train_dur}",
    )

    dataset_params = dataset_config["params"]
    train_dataset = Datapipe.from_dataset_path(
        dataset_path=dataset_path,
        split="train",
        subset=subset,
        **dataset_params,
    )
    logger.info(
        f"Duration of ParametricUMAPDataset used for training, in seconds: {train_dataset.duration}",
    )
    train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        shuffle=shuffle,
        batch_size=batch_size,
        num_workers=num_workers,
    )

    # ---------------- load validation set (if there is one) -----------------------------------------------------------
    if val_step:
        val_dataset = Datapipe.from_dataset_path(
            dataset_path=dataset_path,
            split="val",
            **dataset_params,
        )
        logger.info(
            f"Duration of ParametricUMAPDataset used for validation, in seconds: {val_dataset.duration}",
        )
        val_loader = torch.utils.data.DataLoader(
            dataset=val_dataset,
            shuffle=False,
            batch_size=batch_size,
            num_workers=num_workers,
        )
    else:
        val_loader = None

    model_name = model_config["name"]
    model = models.get(
        model_name,
        model_config,
        input_shape=train_dataset.shape,
    )

    if checkpoint_path is not None:
        logger.info(
            f"loading checkpoint for {model_name} from path: {checkpoint_path}",
        )
        model.load_state_dict_from_path(checkpoint_path)

    results_model_root = results_path.joinpath(model_name)
    results_model_root.mkdir()
    ckpt_root = results_model_root.joinpath("checkpoints")
    ckpt_root.mkdir()
    logger.info(f"training {model_name}")
    trainer = get_trainer(
        accelerator=trainer_config["accelerator"],
        devices=trainer_config["devices"],
        max_epochs=num_epochs,
        log_save_dir=results_model_root,
        ckpt_root=ckpt_root,
        ckpt_step=ckpt_step,
    )
    train_time_start = datetime.datetime.now()
    logger.info(f"Training start time: {train_time_start.isoformat()}")
    trainer.fit(
        model=model,
        train_dataloaders=train_loader,
        val_dataloaders=val_loader,
    )
    train_time_stop = datetime.datetime.now()
    logger.info(f"Training stop time: {train_time_stop.isoformat()}")
    elapsed = train_time_stop - train_time_start
    logger.info(f"Elapsed training time: {elapsed}")
