# This unittest checks the Evaluator which is the high-level API
# It is the most important test file in useckit.v2.

import unittest, gc

import numpy as np
from matplotlib import pyplot as plt
from sklearn.ensemble import RandomForestClassifier, IsolationForest
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC, OneClassSVM

from useckit.v2.evaluator import Paradigm, Evaluator, OvoOvr, ModelArchitecture
from useckit.v2.utils.example_dataset import get_chi2021_dataset
from useckit.v2.viz import plot_roc, plot_far_frr, plot_auc_bars, plot_precision_recall, plot_confusion_bars, \
    plot_score_histograms, plot_confusion_matrix

from tensorflow.keras import backend as K


class TestEvaluatorV2(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Use feature vectors so that sklearn models can be trained easily
        cls.dataset_featurevector = get_chi2021_dataset(feature_vector=True)
        cls.dataset_padded = get_chi2021_dataset(padding=True)
        cls.dataset_ws = get_chi2021_dataset(ws_enable=True)

    def tearDown(self):
        """Aggressive per-test cleanup: close figures, clear TF session, collect garbage."""
        try:
            plt.close("all")
        except Exception:
            pass
        if K is not None:
            try:
                K.clear_session()
            except Exception:
                pass
        gc.collect()

    def _check_and_plot_accuracy_evaluator_results(self, evaluator_results):
        metrics = evaluator_results["accuracy_evaluation"]
        self.assertIsInstance(evaluator_results, dict)
        self.assertIn('accuracy_evaluation', evaluator_results)

        # Basic structure
        self.assertIsInstance(metrics, dict)
        self.assertEqual(set(metrics.keys()), {"accuracy", "confusion_matrix", "confusion_matrix_labels", "classification_report_txt", "classification_report"})

        # Accuracy checks
        self.assertIsInstance(metrics["accuracy"], float)
        self.assertTrue(np.isfinite(metrics["accuracy"]))
        self.assertGreaterEqual(metrics["accuracy"], 0.2)
        self.assertLessEqual(metrics["accuracy"], 0.8   )

        # Confusion matrix checks
        cm = metrics["confusion_matrix"]
        self.assertIsInstance(cm, np.ndarray)
        self.assertEqual(cm.ndim, 2)
        self.assertEqual(cm.shape[0], cm.shape[1])  # must be square
        self.assertTrue(np.issubdtype(cm.dtype, np.integer) or np.issubdtype(cm.dtype, np.number))

        # Classification report checks (string)
        report_txt = metrics["classification_report_txt"]
        self.assertIsInstance(report_txt, str)
        self.assertIn("precision", report_txt)
        
        # Classification report checks (dict)
        report_dict = metrics["classification_report"]
        self.assertIsInstance(report_dict, dict)
        self.assertIn("accuracy", report_dict)
        # Ensure all labels appear in the report somewhere (string contains label text)
        for cls in np.unique(self.dataset_featurevector.get_unique_classes()):
            self.assertIn(str(cls), report_txt)

        # Viz
        plot_confusion_matrix(evaluator_results)

        # Print to console
        print(f"Accuracy: {metrics['accuracy']:.2f}")

    def _check_and_plot_threshold_evaluator_results(self, evaluator_results):
        """Function checks everything that threshold evaluator results should contain."""
        metrics = evaluator_results['threshold_evaluation']['overall_metrics']
        self.assertIsInstance(evaluator_results, dict)
        self.assertIn('threshold_evaluation', evaluator_results)

        self.assertIsInstance(metrics, dict)

        expected_keys = {
            'eer', 'eer_threshold', 'auc',
            'tp', 'tn', 'fp', 'fn', 'accuracy',
            'precision', 'recall', 'f1_score',
            'fpr', 'fnr', 'tpr', 'tnr'
        }
        self.assertTrue(expected_keys.issubset(metrics.keys()))

        # Type and range checks
        self.assertIsInstance(metrics['tp'], (int, np.integer))
        self.assertIsInstance(metrics['fp'], (int, np.integer))
        self.assertIsInstance(metrics['fn'], (int, np.integer))
        self.assertIsInstance(metrics['tn'], (int, np.integer))

        # Viz
        plot_roc(evaluator_results)
        plot_far_frr(evaluator_results)
        plot_auc_bars(evaluator_results)
        plot_precision_recall(evaluator_results)
        plot_confusion_bars(evaluator_results)
        plot_score_histograms(evaluator_results)

        # Print to console
        print(f"Accuracy: {metrics['accuracy']:.2f}")
        print(f"AUC: {metrics['auc']:.2f}")
        print(f"EER: {metrics['eer']:.2f}")
        print(f"EER Threshold: {metrics['eer_threshold']:.2f}")

        for k in ['eer', 'eer_threshold', 'auc', 'accuracy', 'fpr', 'fnr', 'tpr', 'tnr']:
            self.assertIsInstance(metrics[k], (float, np.floating))
            # In typical cases, values should be finite and within bounds
            self.assertTrue(np.isfinite(metrics[k]))
        self.assertGreater(metrics['accuracy'], 0.1)
        self.assertLessEqual(metrics['accuracy'], 1.0)
        self.assertGreater(metrics['auc'], 0.1)
        self.assertLessEqual(metrics['auc'], 1.0)

    #######
    # CSC #
    #######

    def test_csc_fit_and_accuracy_sklearn_rf(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_featurevector,
            model_initializer_callable=RandomForestClassifier(),
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with accuracy, confusion_matrix, and both reports
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    def test_csc_fit_and_accuracy_sklearn_knn(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_featurevector,
            model_initializer_callable=KNeighborsClassifier(),
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with accuracy, confusion_matrix, and both reports
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    def test_csc_fit_and_accuracy_sklearn_svm(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_featurevector,
            model_initializer_callable=SVC(C=100, probability=True),
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with accuracy, confusion_matrix, and both reports
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    def test_csc_fit_and_accuracy_keras(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_padded,
            model_architecture=ModelArchitecture.MLP
        )
        evaluator.fit(epochs=10, verbose=2)

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    ######
    # BV #
    ######

    def test_bv_threshold_evaluation_ovr_sklearn_rf(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_featurevector,
            model_initializer_callable=RandomForestClassifier(),
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovr_sklearn_svm(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_featurevector,
            model_initializer_callable=SVC(C=100, probability=True),
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovr_sklearn_knn(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_featurevector,
            model_initializer_callable=KNeighborsClassifier(),
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_sklearn_rf(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_featurevector,
            model_initializer_callable=RandomForestClassifier(),
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_sklearn_svm(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_featurevector,
            model_initializer_callable=SVC(C=100, probability=True),
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_sklearn_knn(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_featurevector,
            model_initializer_callable=KNeighborsClassifier(),
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovr_keras(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_padded,
            model_architecture=ModelArchitecture.MLP,
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_keras(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_padded,
            model_architecture=ModelArchitecture.MLP,
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    ######
    # AD #
    ######

    def test_ad_threshold_evaluation_sklearn_if(self):
        evaluator = Evaluator(
            paradigm=Paradigm.AD,
            dataset=self.dataset_featurevector,
            model_initializer_callable=IsolationForest(),
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_ad_threshold_evaluation_sklearn_oneclasssvm(self):
        evaluator = Evaluator(
            paradigm=Paradigm.AD,
            dataset=self.dataset_featurevector,
            model_initializer_callable=OneClassSVM(),
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_ad_threshold_evaluation_keras(self):
        evaluator = Evaluator(
            paradigm=Paradigm.AD,
            dataset=self.dataset_padded,
            model_architecture=ModelArchitecture.MLP
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    ######
    # DL #
    ######

    def test_dl_threshold_evaluation_keras_contrastive(self):
        evaluator = Evaluator(
            paradigm=Paradigm.DL,
            dataset=self.dataset_padded,
            model_architecture=ModelArchitecture.MLP,
            distance_learning_loss="contrastive"
        )
        evaluator.fit(epochs=100)

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_dl_threshold_evaluation_keras_triplet(self):
        evaluator = Evaluator(
            paradigm=Paradigm.DL,
            dataset=self.dataset_padded,
            model_architecture=ModelArchitecture.MLP,
            distance_learning_loss="triplet"
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    ############
    # CSC + WS #
    ############

    def test_csc_fit_and_accuracy_sklearn_rf_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_ws,
            model_initializer_callable=RandomForestClassifier(),
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    def test_csc_fit_and_accuracy_sklearn_knn_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_ws,
            model_initializer_callable=KNeighborsClassifier(),
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    def test_csc_fit_and_accuracy_sklearn_svm_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_ws,
            model_initializer_callable=SVC(C=100, probability=True),
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    def test_csc_fit_and_accuracy_keras_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.CSC,
            dataset=self.dataset_ws,
            model_architecture=ModelArchitecture.MLP
        )
        evaluator.fit(epochs=10, verbose=2)

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_accuracy_evaluator_results(evaluator_results)

    ###########
    # BV + WS #
    ###########

    def test_bv_threshold_evaluation_ovr_sklearn_rf_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_initializer_callable=RandomForestClassifier(),
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovr_sklearn_svm_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_initializer_callable=SVC(C=100, probability=True),
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovr_sklearn_knn_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_initializer_callable=KNeighborsClassifier(),
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_sklearn_rf_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_initializer_callable=RandomForestClassifier(),
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_sklearn_svm_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_initializer_callable=SVC(C=100, probability=True),
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_sklearn_knn_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_initializer_callable=KNeighborsClassifier(),
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovr_keras_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_architecture=ModelArchitecture.MLP,
            ovo_ovr=OvoOvr.OVR,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_bv_threshold_evaluation_ovo_keras_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.BV,
            dataset=self.dataset_ws,
            model_architecture=ModelArchitecture.MLP,
            ovo_ovr=OvoOvr.OVO,
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    ######
    # AD #
    ######

    def test_ad_threshold_evaluation_sklearn_if_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.AD,
            dataset=self.dataset_ws,
            model_initializer_callable=IsolationForest(),
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_ad_threshold_evaluation_sklearn_oneclasssvm_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.AD,
            dataset=self.dataset_ws,
            model_initializer_callable=OneClassSVM(),
        )
        evaluator.fit()

        # Evaluate threshold-based metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_ad_threshold_evaluation_keras_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.AD,
            dataset=self.dataset_ws,
            model_architecture=ModelArchitecture.MLP
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    ###########
    # DL + WS #
    ###########

    def test_dl_threshold_evaluation_keras_contrastive_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.DL,
            dataset=self.dataset_ws,
            model_architecture=ModelArchitecture.MLP,
            distance_learning_loss="contrastive"
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)

    def test_dl_threshold_evaluation_keras_triplet_ws(self):
        evaluator = Evaluator(
            paradigm=Paradigm.DL,
            dataset=self.dataset_ws,
            model_architecture=ModelArchitecture.MLP,
            distance_learning_loss="triplet"
        )
        evaluator.fit()

        # Evaluate accuracy; Evaluator.evaluate() should return a flat dict with the three metrics
        evaluator_results = evaluator.evaluate()
        self._check_and_plot_threshold_evaluator_results(evaluator_results)



if __name__ == '__main__':
    unittest.main()
