# Authors: SheepTAO <sheeptao@outlook.com>

# License: MIT
# Copyright the dpeeg contributors.

from sklearn.model_selection import StratifiedKFold
from torchmetrics.aggregation import MeanMetric, CatMetric

from .base import ClsExp
from ..trainer.classifier import BaseClassifier
from ..datasets.base import SplitEEGData
from ..transforms.base import ToEEGData, SplitTrainTest
from ..tools import Filer
from ..utils import DPEEG_SEED, get_init_args
from ..tools.docs import fill_doc


@fill_doc
class KFold(ClsExp):
    r"""K-Fold cross validation experiment.

    The KFold experiment divides the dataset into K non-overlapping subsets
    (i.e., "folds") and repeatedly trains and tests the model. The purpose is
    to reduce the dependence of the model evaluation results on the way the
    dataset is divided and to improve the stability and reliability of the
    evaluation results. However, its computational cost is high, especially for
    large datasets and complex models. It may take a long time to complete the
    training of all folds.

    Two validation methods are provided in the experiment, determined by the
    parameter ``isolate_testset``. When set to ``True``, it indicates that the
    test set is independent of the k-fold cross-validation. That is, for each
    fold, the data is divided into a training set and a validation set to find
    the optimal parameters for each fold, and then the model is evaluated on an
    independent dataset. When set to ``False``, it indicates that one fold of
    data in each fold is used as the test set, and the remaining folds are used
    to train the model. The average value of all folds' evaluations is used as
    the performance metric of the model. The specific experimental method is
    shown in the figure below, which illustrates a 3-fold cross-validation
    experiment:

    .. image:: ../_static/images/kfold_isolate_testset.png
       :align: center
       :alt: kfold isolate testset

    When the training set and test set come from different sessions, setting
    this parameter is very useful.

    Parameters
    ----------
    %(trainer)s
    %(out_folder)s
    k : int, optional
        k of k-Fold.
    isolate_testset : bool
        By default, the test set is independent, that is, the k-fold cross-
        validation at this time only divides the training set and the
        verification set based on the training set to implement an early
        stopping mechanism, and finally evaluates on the isolated test set.
        If False, the test set is for each fold of k-fold cross-validation.
    shuffle : bool
        Shuffle before kfold.
    seed : int
        Seed of random for review.
    %(timestamp)s

    Notes
    -----
    If ``isolate_testset`` False, please provide the ``transforms`` parameter
    of the ``run`` function to avoid data leakage caused by operations such
    as data augmentation in advance. When set to ``True``, it means that the
    experiment requires the `trainer` to support a validation set.
    """

    def __init__(
        self,
        trainer: BaseClassifier,
        out_folder: str | None = None,
        k: int = 5,
        isolate_testset: bool = True,
        shuffle: bool = True,
        seed: int = DPEEG_SEED,
        timestamp: bool = True,
        verbose: int | str = "INFO",
    ) -> None:
        super().__init__(
            get_init_args(self, locals(), ret_dict=True),
            trainer=trainer,
            out_folder=out_folder,
            timestamp=timestamp,
            verbose=verbose,
        )
        self.k = str(k)
        self.isolate_testset = isolate_testset
        self.skf = StratifiedKFold(k, shuffle=shuffle, random_state=seed)

    def _run_sub_classifier(self, eegdata, sub_folder):
        if self.isolate_testset:
            raise TypeError("Isolated test set is useless with a `Classifier` trainer.")

        eegdata = ToEEGData()(eegdata, verbose=False)

        result = {}
        filer = Filer(sub_folder / "summary.txt")
        train_acc_metric = MeanMetric()
        test_acc_metric = MeanMetric()
        preds_metric = CatMetric()
        target_metric = CatMetric()

        for exp_idx, (train_idx, test_idx) in enumerate(
            self.skf.split(eegdata["edata"], eegdata["label"]), start=1
        ):
            self.logger.info(f"\n# ---- {sub_folder.name}_exp{exp_idx} ---- #")

            trainset = eegdata.index(train_idx)
            testset = eegdata.index(test_idx)
            egd = self._trans_eegdata(SplitEEGData(trainset, testset))
            train_set = egd["train"]
            test_set = egd["test"]

            exp_folder = sub_folder / f"exp{exp_idx}"
            exp_result = self.trainer.fit(
                trainset=train_set,
                testset=test_set,
                log_dir=exp_folder,
            )

            result[f"exp{exp_idx}"] = exp_result
            train_acc_metric.update(exp_result["train"]["acc"])
            test_acc_metric.update(exp_result["test"]["acc"])
            preds_metric.update(exp_result["test"]["preds"])
            target_metric.update(exp_result["test"]["target"])

            filer.write(
                f"Exp_{str(exp_idx).zfill(len(self.k))} Acc: "
                f"Train={exp_result['train']['acc']:.4f} | "
                f"Test={exp_result['test']['acc']:.4f}\n"
            )

        train_acc = train_acc_metric.compute()
        test_acc = test_acc_metric.compute()
        filer.write(f"Avg Acc = {test_acc*100:.2f}%\n")
        self.logger.info("-" * 30)
        self.logger.info(f"Avg Acc: Train={train_acc:.4f} | Test={test_acc:.4f}")

        result.update(
            {
                "acc": test_acc,
                "preds": preds_metric.compute(),
                "target": target_metric.compute(),
            }
        )
        return result

    def _run_sub_classifier_two_stage(self, eegdata, sub_folder):
        if self.isolate_testset:
            eegdata = self._trans_eegdata(eegdata)
            X = eegdata["train"]["edata"]
            y = eegdata["train"]["label"]
        else:
            eegdata = ToEEGData()(eegdata, verbose=False)
            X = eegdata["edata"]
            y = eegdata["label"]

        result = {}
        filer = Filer(sub_folder / "summary.txt")
        train_acc_metric = MeanMetric()
        valid_acc_metric = MeanMetric()
        test_acc_metric = MeanMetric()
        preds_metric = CatMetric()
        target_metric = CatMetric()

        for exp_idx, (train_idx, test_idx) in enumerate(self.skf.split(X, y), start=1):
            self.logger.info(f"\n# ---- {sub_folder.name}_exp{exp_idx} ---- #")

            if self.isolate_testset:
                train_set = eegdata["train"].index(train_idx)
                valid_set = eegdata["train"].index(test_idx)
                test_set = eegdata["test"]
            else:
                trainset = eegdata.index(train_idx)  # type: ignore
                testset = eegdata.index(test_idx)  # type: ignore
                egd = self._trans_eegdata(SplitEEGData(trainset, testset))

                test_set = egd["test"]
                trainset = egd["train"]
                egd = SplitTrainTest()(trainset)
                train_set = egd["train"]
                valid_set = egd["test"]

            exp_folder = sub_folder / f"exp{exp_idx}"
            exp_result = self.trainer.fit(
                trainset=train_set,
                validset=valid_set,
                testset=test_set,
                log_dir=exp_folder,
            )

            result[f"exp{exp_idx}"] = exp_result
            train_acc_metric.update(exp_result["train"]["acc"])
            valid_acc_metric.update(exp_result["valid"]["acc"])
            test_acc_metric.update(exp_result["test"]["acc"])
            preds_metric.update(exp_result["test"]["preds"])
            target_metric.update(exp_result["test"]["target"])

            filer.write(
                f"Exp_{str(exp_idx).zfill(len(self.k))} Acc: "
                f"Train={exp_result['train']['acc']:.4f} | "
                f"Valid={exp_result['train']['acc']:.4f} | "
                f"Test={exp_result['test']['acc']:.4f}\n"
            )

        train_acc = train_acc_metric.compute()
        valid_acc = valid_acc_metric.compute()
        test_acc = test_acc_metric.compute()
        filer.write(f"Avg Acc = {test_acc*100:.2f}%\n")
        self.logger.info("-" * 30)
        self.logger.info(
            f"Avg Acc: Train={train_acc:.4f} | Valid={valid_acc:.4f} | "
            f"Test={test_acc:.4f}"
        )

        result.update(
            {
                "acc": test_acc,
                "preds": preds_metric.compute(),
                "target": target_metric.compute(),
            }
        )
        return result

    def _run_sub(self, eegdata, sub_folder):
        """Basic K-Fold cross validation function.

        Returns
        -------
        Return test_acc, test_kappa, test_preds, test_target and results :\n
        results = {
            'expNo_1' : { ... },
            'expNo_2' : { ... },
            .
            .
        }
        """
        self.timer.start("kfold")
        self.logger.info(f"\n# ---------- {sub_folder.name} ---------- #")
        result = self._run_sub_func(eegdata, sub_folder)
        h, m, s = self.timer.stop("kfold")
        self.logger.info(
            f"\n[{self.k}Fold CV Finish] - [Cost Time = {h}H:{m}M:{s:.2f}S]"
        )

        return result["acc"], result["preds"], result["target"], result
