from enum import Enum
import logging
from importlib.metadata import metadata
from os import PathLike

from sklearn.ensemble import RandomForestClassifier, IsolationForest
from sklearn.metrics import DistanceMetric

logging.basicConfig(level=logging.INFO)

import numpy as np

from useckit.v2.dataset import Dataset
from useckit.v2.model_descriptions import dl4tsc_cnn_padding_valid, dl4tsc_cnn_padding_same, dl4tsc_fcn, dl4tsc_mlp, \
    dl4tsc_resnet, dl4tsc_encoder, dl4tsc_inception, dl4tsc_tlenet, dl4tsc_mcdcnn, dl4tsc_mcnn, dl4tsc_twiesn


class ModelArchitecture(Enum):
    CNN_VALID = 0
    CNN_SAME = 1
    FCN = 2
    MLP = 3
    RESNET = 4
    ENCODER = 5
    MCNN = 6
    MCDCNN = 7
    TLENET = 8
    INCEPTION = 9
    TWIESN = 10


ModelDescriptions = {
    ModelArchitecture.CNN_VALID: dl4tsc_cnn_padding_valid,
    ModelArchitecture.CNN_SAME: dl4tsc_cnn_padding_same,
    ModelArchitecture.FCN: dl4tsc_fcn,
    ModelArchitecture.MLP: dl4tsc_mlp,
    ModelArchitecture.RESNET: dl4tsc_resnet,
    ModelArchitecture.ENCODER: dl4tsc_encoder,
    ModelArchitecture.INCEPTION: dl4tsc_inception,
    ModelArchitecture.TLENET: dl4tsc_tlenet,
    ModelArchitecture.MCDCNN: dl4tsc_mcdcnn,
    ModelArchitecture.MCNN: dl4tsc_mcnn,
    # Non-Keras
    ModelArchitecture.TWIESN: dl4tsc_twiesn,
}


class OvoOvr(Enum):
    OVO = 0
    OVR = 1


class DistanceLearningLoss(Enum):
    TRIPLET = 0
    CONTRASTIVE = 1


class Paradigm(Enum):
    CSC = CLOSED_SET_CLASSIFICATION = 0  # Formerly "Time-Series Classification"
    BV = BINARY_VERIFICATION = 1
    DL = DISTANCE_LEARNING = 2
    AD = ANOMALY_DETECTION = ONE_CLASS_CLASSIFICATION = 3
    #DM = DISTANCE_METRIC = 4  # currently not implemented


class Evaluator:
    def __init__(self,
                 paradigm: Paradigm,
                 dataset: Dataset,
                 model_architecture: ModelArchitecture = None,
                 model_initializer_callable = None,
                 serialize_results: bool = True,
                 experiment_description: str = None,
                 **kwargs):
        self.paradigm: Paradigm = paradigm
        self.dataset: Dataset = dataset
        self.model_architecture: ModelArchitecture = model_architecture
        self.model_graph: dict = dict()
        self.fit_results = {}
        self._fitted: bool = False

        self.serialize_results = serialize_results
        self.experiment_description = experiment_description
        self.experiment_path = None

        # Start of Optionals
        self.ovo_ovr = kwargs.get("ovo_ovr", OvoOvr.OVR) if paradigm == Paradigm.BV else None
        self.distance_learning_loss = None
        # End of Optionals

        ###############################################
        # Step 1) Sanitization of Parameters (kwargs) #
        ###############################################

        # Check whether kwargs are sufficiently passed and store to local variables
        if model_architecture is None:
            if model_initializer_callable is None:
                raise ValueError("You must specify either parameter `model_architecture` (type: ModelArchitecture[enum] "
                                 "or `model_initializer_callable` (type: sklearn-model to clone or custom function9.")
        if paradigm == Paradigm.BV:
            if not isinstance(self.ovo_ovr, OvoOvr):
                raise ValueError("'ovo_ovr' must be of type `OvoOvr[enum]`.")
        if paradigm == Paradigm.DL:
            if not "distance_learning_loss" in kwargs:
                raise ValueError("For distance learning, you must specify 'distance_learning_loss' in kwargs.")
            if kwargs["distance_learning_loss"] not in ["triplet", "contrastive"]:
                raise ValueError("'distance_learning_loss' must be either 'triplet' or 'contrastive'.")
            self.distance_learning_loss = kwargs["distance_learning_loss"]

        ########################################################################
        # Step 2: Initialize Models and set Model initializer (Callable/Clone) #
        ########################################################################

        if model_initializer_callable is None:  # we use only provided model descriptions
            from useckit.v2.keras_model_builder import ModelBuilder
            mb = ModelBuilder(model_descriptions=ModelDescriptions)
            self.model_initializer_callable = lambda *args, **_kwargs: mb.build_model(paradigm=paradigm,
                                                                                     model_architecture=model_architecture,
                                                                                     dataset=dataset,
                                                                                     **kwargs)
        else:
            # Handle sklearn: we take an initialized model and clone it
            self.model_initializer_callable = model_initializer_callable
            from sklearn.base import BaseEstimator
            if isinstance(self.model_initializer_callable, BaseEstimator):
                # Capture original estimator to avoid self-referential closure
                orig_sklearn_obj = self.model_initializer_callable
                from sklearn import clone
                self.model_initializer_callable = lambda *args, **_kwargs: clone(orig_sklearn_obj)
            else:
                pass  # keep user's specified callable

        ################################
        # Step 3: Populate Model Graph #
        ################################

        if self.paradigm == Paradigm.CSC:
            # Initialize single global model at Key `None`
            self.model_graph[None] = self.model_initializer_callable(self.dataset, **kwargs)
        elif self.paradigm == Paradigm.DL:
            self.model_graph[None] = self.model_initializer_callable(self.dataset, **kwargs)
        elif self.paradigm == Paradigm.BV and self.ovo_ovr == OvoOvr.OVR:
            for c in dataset.get_unique_classes():
                self.model_graph[c] = self.model_initializer_callable(self.dataset, **kwargs)
        elif self.paradigm == Paradigm.BV and self.ovo_ovr == OvoOvr.OVO:
            for c1 in dataset.get_unique_classes():
                for c2 in dataset.get_unique_classes():
                    if str(c1) >= str(c2):
                        continue  # skip to avoid duplicate training of the same pair
                    else:
                        self.model_graph[(c1, c2)] = self.model_initializer_callable(self.dataset, **kwargs)
        elif self.paradigm == Paradigm.AD:
            for c in dataset.get_unique_classes():
                self.model_graph[c] = self.model_initializer_callable(self.dataset, **kwargs)

    def __repr__(self):
        import re
        def sanitize(v):
            if v is None:
                return None
            s = re.sub(r'[^A-Za-z0-9]+', '_', str(v)).strip('_')
            return s or None

        # Required bits
        paradigm_str = sanitize(getattr(self.paradigm, "name", self.paradigm))

        parts = []
        if paradigm_str:
            parts.append(f"paradigm-{paradigm_str}")

        # Model: prefer enum .name; else initializer callable name; omit if neither
        model_raw = None
        if getattr(self, "model_architecture", None) is not None:
            model_raw = getattr(self.model_architecture, "name", self.model_architecture)
        elif getattr(self, "model_initializer_callable", None) is not None:
            mic = self.model_initializer_callable
            model_raw = getattr(mic, "__name__", None) or getattr(mic, "__class__", type(mic)).__name__
        model_str = sanitize(model_raw)
        if model_str:
            parts.append(f"model-{model_str}")

        # Distance-learning loss only for DL
        if getattr(self, "paradigm", None) is not None and (
                (hasattr(self.paradigm, "name") and getattr(self.paradigm, "name", "") == "DL")
        ):
            loss_str = sanitize(getattr(self, "distance_learning_loss", None))
            if loss_str:
                parts.append(f"loss-{loss_str}")

        # OVO/OVR only for BV
        if getattr(self, "paradigm", None) is not None and (
                (hasattr(self.paradigm, "name") and getattr(self.paradigm, "name", "") == "BV")
        ):
            ovoovr_raw = getattr(self, "ovo_ovr", None)
            ovoovr_str = sanitize(getattr(ovoovr_raw, "name", ovoovr_raw))
            if ovoovr_str:
                parts.append(f"ovoovr-{ovoovr_str}")

        return "_".join(parts) if parts else "evaluator"

    def fit(self, **kwargs):
        """
        Trains all models depending on paradigm and stores both trained models and fit histories.
        The `kwargs` are passed to teh fit()-method of the underlying model.
        Returns self for method chaining.
        """
        if self.paradigm == Paradigm.CSC:
            self._fit_csc(self.dataset, **kwargs)
        elif self.paradigm == Paradigm.DL:
            self._fit_dl(self.dataset, **kwargs)
        elif self.paradigm == Paradigm.BV:
            self._fit_bv(self.dataset, self.ovo_ovr, **kwargs)
        elif self.paradigm == Paradigm.AD:
            self._fit_ad(self.dataset, **kwargs)
        else:
            raise ValueError(f"Unknown paradigm in fit(): {self.paradigm}")

        # serialization of models and creation of experiment folder happens after fitting
        if self.serialize_results:
            from useckit.v2.session_folder_management import ExperimentFolder
            experiment_description = f"{self}-{'' if self.experiment_description is None else self.experiment_description}".removesuffix('-')
            self.experiment_path = ExperimentFolder(experiment_description).get_path()

        self._fitted = True
        return self

    def _fit_csc(self, dataset: Dataset, **kwargs) -> None:
        """Train the single global Closed Set Classification model."""
        model = self.model_graph.get(None)
        if model is None:
            raise RuntimeError("Global model is not initialized for TSC paradigm.")

        # Select training data: prefer enrollment set; else use train set
        x_train = dataset.get_enrollment_data()
        y_train = dataset.get_enrollment_labels(ohc=self.is_deeplearning_model(model))

        if x_train is None or y_train is None:
            raise ValueError("No training data available for TSC paradigm: enrollment/train data or labels are None.")

        trained_model, history = self._train_model(model, x_train, y_train, **kwargs)
        self.model_graph[None] = trained_model
        self.fit_results[None] = {"model": trained_model, "history": history}

    def _fit_dl(self, dataset: Dataset, *, num_pos_per_anchor: int = 1, num_neg_per_anchor: int = 3,
                seed: int | None = None, **kwargs):
        """Train the single global Distance Learning (DL) model."""
        x_train = dataset.get_trainset_data()
        y_train = dataset.get_trainset_labels()

        if x_train is None or y_train is None:
            raise ValueError("No training data available for DL paradigm.")

        model = self.model_graph.get(None)
        if model is None:
            raise RuntimeError("Global model is not initialized for DL paradigm.")

        unique_classes = dataset.get_unique_classes()
        label_to_int = {cls: idx for idx, cls in enumerate(unique_classes)}
        y_int = np.asarray([label_to_int[lbl] for lbl in y_train], dtype=np.int32)

        def _make_contrastive_pairs(X, y_idx):
            X = np.asarray(X)
            n = X.shape[0]
            by_cls = {cls: np.where(y_idx == cls)[0] for cls in np.unique(y_idx)}
            not_cls = {cls: np.where(y_idx != cls)[0] for cls in np.unique(y_idx)}

            Xa, Xb, yb = [], [], []
            rng = np.random.default_rng(seed)
            for i in range(n):
                cls_id = y_idx[i]
                anchor = X[i]

                pos_candidates = by_cls[cls_id]
                if pos_candidates.size > 1:
                    excl_self = pos_candidates[pos_candidates != i]
                    pos_choices = rng.choice(excl_self, size=num_pos_per_anchor,
                                             replace=excl_self.size < num_pos_per_anchor)
                    for pos_idx in np.atleast_1d(pos_choices):
                        Xa.append(anchor); Xb.append(X[int(pos_idx)]); yb.append(0.0)

                neg_candidates = not_cls[cls_id]
                neg_choices = rng.choice(neg_candidates, size=num_neg_per_anchor,
                                         replace=neg_candidates.size < num_neg_per_anchor)
                for neg_idx in np.atleast_1d(neg_choices):
                    Xa.append(anchor); Xb.append(X[int(neg_idx)]); yb.append(1.0)

            perm = rng.permutation(len(yb))
            Xa, Xb, yb = np.array(Xa)[perm], np.array(Xb)[perm], np.array(yb, dtype=np.float32)[perm].reshape(-1, 1)
            return [Xa, Xb], yb

        if self.distance_learning_loss == "contrastive":
            x_train_dl, y_train_dl = _make_contrastive_pairs(x_train, y_int)
        elif self.distance_learning_loss == "triplet":
            x_train_dl, y_train_dl = x_train, y_int
        else:
            raise ValueError(f"Unknown distance_learning_loss: {self.distance_learning_loss}")

        trained_model, history = self._train_model(model, x_train_dl, y_train_dl, **kwargs)
        self.model_graph[None] = trained_model
        self.fit_results[None] = {"model": trained_model, "history": history}

    def _fit_bv(self, dataset: Dataset, ovo_ovr, **kwargs):
        unique_labels = dataset.get_unique_enrollment_classes()

        fit_counter = 1
        for c1 in unique_labels:
            if ovo_ovr == OvoOvr.OVR:
                pos = dataset.get_enrollment_data(filter=lambda x: x == c1)
                neg = dataset.get_enrollment_data(filter=lambda x: x != c1)
                x_train = np.concatenate((pos, neg))
                y_train = np.concatenate((np.ones(len(pos)), np.zeros(len(neg))))
                trained_model, history = self._train_model(self.model_graph[c1], x_train, y_train, **kwargs)
                self.model_graph[c1] = trained_model
                self.fit_results[c1] = {"model": trained_model, "history": history}
                logging.info(f"Fitted BV (OVR) model {fit_counter}/{len(self.model_graph)} for label {c1}.")
                fit_counter += 1

            elif ovo_ovr == OvoOvr.OVO:
                for c2 in unique_labels:
                    if str(c1) >= str(c2):
                        continue
                    pos = dataset.get_enrollment_data(filter=lambda x: x == c1)
                    neg = dataset.get_enrollment_data(filter=lambda x: x == c2)
                    x_train = np.concatenate((pos, neg))
                    y_train = np.concatenate((np.ones(len(pos)), np.zeros(len(neg))))
                    trained_model, history = self._train_model(self.model_graph[(c1, c2)], x_train, y_train, **kwargs)
                    self.model_graph[(c1, c2)] = trained_model
                    self.fit_results[(c1, c2)] = {"model": trained_model, "history": history}
                    logging.info(f"Fitted BV (OVO) model {fit_counter}/{len(self.model_graph)} for labels {c1}/{c2}.")
                    fit_counter += 1
            else:
                raise ValueError(f"Unknown ovo_ovr: {ovo_ovr}")

    def _fit_ad(self, dataset: Dataset, **kwargs):
        unique_classes = dataset.get_unique_enrollment_classes()

        for i, c in enumerate(unique_classes):
            x_train = dataset.get_enrollment_data(lambda x: x == c)
            trained_model, history = self._train_model(self.model_graph[c], x_train, None, **kwargs)
            self.model_graph[c] = trained_model
            self.fit_results[c] = {"model": trained_model, "history": history}
            logging.info(f"Fitted AD model {i + 1}/{len(unique_classes)} for label {c}.")

    @staticmethod
    def _train_model(model, x_train, y_train, **kwargs):
        """
        Train a model with type-aware return semantics.

        Returns:
          - (BuildResult, history) for useckit BuildResult wrappers
          - (model, history) for raw Keras models
          - (model, None) for scikit-learn estimators
          - (model, fit_return) for custom estimators/callables
        """

        def _is_keras_model(obj):
            # Prefer true Keras type checks if available
            import tensorflow as tf
            from keras import Model as KerasModel
            return isinstance(obj, (tf.keras.Model, KerasModel))

        def _is_useckit_model(obj):
            from useckit.v2.keras_model_builder import BuildResult
            return isinstance(obj, BuildResult)

        def _is_sklearn_estimator(obj):
            # Prefer sklearn BaseEstimator if available
            from sklearn.base import BaseEstimator
            return isinstance(obj, BaseEstimator)

        def _fit_and_return(model):
            # 1) useckit BuildResult: update in place and return BuildResult
            if _is_useckit_model(model):
                build_result = model
                _model = build_result.model
                _callbacks = build_result.callbacks
                _metadata = build_result.metadata  # kept for completeness if you need it

                if "callbacks" in kwargs:
                    _callbacks = _callbacks + kwargs.pop("callbacks")

                history = _model.fit(
                    x_train,
                    y_train if y_train is not None else x_train,  # y_train can be None in OCC
                    callbacks=_callbacks,
                    **kwargs,
                )

                # update BuildResult in place so model_graph keeps the wrapper
                build_result.model = _model
                # optionally store history on the BuildResult (not in the dataclass fields, but as an attr)
                setattr(build_result, "history", history)

                return build_result, history

            # 2) Raw Keras: return (model, history)
            if _is_keras_model(model):
                history = model.fit(x_train, y_train, **kwargs)
                return model, history

            # 3) scikit-learn: return fitted model only
            if _is_sklearn_estimator(model):
                try:
                    fitted = model.fit(x_train, y_train, **kwargs)
                except ValueError as ve:
                    if "Found array with dim 3, while dim <= 2 is required" in str(ve):
                        fitted = model.fit(x_train.reshape(x_train.shape[0], -1), y_train, **kwargs)
                return fitted, None

            # 4) Custom estimator: return (model, fit_return)
            if hasattr(model, "fit") and callable(getattr(model, "fit")):
                fit_ret = model.fit(x_train, y_train, **kwargs)
                return model, fit_ret

            raise TypeError("Provided model is not a recognized estimator and has no callable fit().")

        return _fit_and_return(model)

    def evaluate(self):
        if not self._fitted:
            raise RuntimeError("Models must be fitted before running evaluation.")

        from useckit.v2.evaluation_methods import (accuracy_evaluation,
                                                   multiclass_auc_evaluation,
                                                   threshold_evaluation_via_global_triplet_model,
                                                   threshold_evaluation_via_global_contrastive_model,
                                                   threshold_evaluation_via_one_model_per_class,
                                                   threshold_evaluation_via_class_vs_class_graph_ensemble)

        evaluation_function_names = []

        if self.paradigm == Paradigm.CSC:
            # Return flat metrics dict for CSC (accuracy, confusion_matrix, classification_report)
            r = {"accuracy_evaluation": accuracy_evaluation(self.dataset, self.model_graph),
                 "multiclass_auc_evaluation": multiclass_auc_evaluation(self.dataset, self.model_graph)}
            evaluation_function_names.append({"multiclass_auc_evaluation": "multiclass_auc_evaluation"})
            evaluation_function_names.append({"accuracy_evaluation": "accuracy_evaluation"})
        elif self.paradigm == Paradigm.DL and self.distance_learning_loss == "contrastive":
            r = {'threshold_evaluation': threshold_evaluation_via_global_contrastive_model(self.dataset, self.model_graph)}
            evaluation_function_names.append({'threshold_evaluation': 'threshold_evaluation_via_global_contrastive_model'})
        elif self.paradigm == Paradigm.DL and self.distance_learning_loss == "triplet":
            r = {'threshold_evaluation': threshold_evaluation_via_global_triplet_model(self.dataset, self.model_graph)}
            evaluation_function_names.append({'threshold_evaluation_': 'threshold_evaluation_via_global_triplet_model'})
        elif self.paradigm == Paradigm.BV and self.ovo_ovr == OvoOvr.OVO:
            r = {'threshold_evaluation': threshold_evaluation_via_class_vs_class_graph_ensemble(self.dataset, self.model_graph)}
            evaluation_function_names.append({'threshold_evaluation': 'threshold_evaluation_via_class_vs_class_graph_ensemble'})
        elif self.paradigm == Paradigm.BV and self.ovo_ovr == OvoOvr.OVR:
            r = {'threshold_evaluation': threshold_evaluation_via_one_model_per_class(self.dataset, self.model_graph)}
            evaluation_function_names.append({'threshold_evaluation': 'threshold_evaluation_via_one_model_per_class'})
        elif self.paradigm == Paradigm.AD:
            r = {'threshold_evaluation': threshold_evaluation_via_one_model_per_class(self.dataset, self.model_graph)}
            evaluation_function_names.append({'threshold_evaluation': 'threshold_evaluation_via_one_model_per_class'})

        # add meta information
        r['meta'] = {'num_models': len(self.model_graph),
                     'evaluation_function_names': evaluation_function_names}

        self._plot_metrics(r)

        return r

    @staticmethod
    def is_deeplearning_model(obj):
        import tensorflow as tf
        from keras import Model as KerasModel
        from useckit.v2.keras_model_builder import BuildResult
        return isinstance(obj, (tf.keras.Model, KerasModel, BuildResult))

    def get_fit_summary(self):
        """Return a lightweight dict of trained model names and history availability."""
        return {
            str(k): {
                "type": type(v["model"]).__name__,
                "has_history": v["history"] is not None
            }
            for k, v in self.fit_results.items()
        }

    def _save_model(self, model_obj, out_base_path: str):
        """Persist a trained model to disk.

        - For scikit-learn estimators: writes `<out_base_path>.pkl` via joblib/pickle.
        - For Keras models (including BuildResult wrappers): writes `<out_base_path>.keras` via model.save().
        """
        try:
            # Unwrap BuildResult if provided
            from useckit.v2.keras_model_builder import BuildResult
            if isinstance(model_obj, BuildResult):
                model_obj = model_obj.model
        except Exception:
            pass

        # Try Keras first
        try:
            import tensorflow as tf
            from keras import Model as KerasModel
            if isinstance(model_obj, (tf.keras.Model, KerasModel)):
                out_path = f"{out_base_path}.keras"
                try:
                    model_obj.save(out_path)
                    return out_path
                except Exception:
                    # Fallback to H5 if needed
                    out_path = f"{out_base_path}.h5"
                    model_obj.save(out_path)
                    return out_path
        except Exception:
            pass

        # Then try scikit-learn (or any pickle-able estimator)
        try:
            try:
                import joblib
                out_path = f"{out_base_path}.pkl"
                joblib.dump(model_obj, out_path)
                return out_path
            except Exception:
                import pickle
                out_path = f"{out_base_path}.pkl"
                with open(out_path, 'wb') as f:
                    pickle.dump(model_obj, f)
                return out_path
        except Exception:
            # Last resort: do nothing if cannot persist
            return None

    def _plot_metrics(self, r):
        if not self.serialize_results:
            return

        # Save Training Histories
        self._save_fit_histories()

        # General metrics saving:
        import orjson
        import pandas as pd
        import os

        # safe json that can be saved as string (e.g., no numpy arrays in it)
        safe_r = make_json_safe(r)
        # JSON (strict, fast)
        json_path = os.path.join(self.experiment_path, "evaluation_results.json")
        with open(json_path, "wb") as f:
            f.write(orjson.dumps(safe_r, option=orjson.OPT_INDENT_2))
        # CSV (flat)
        csv_path = os.path.join(self.experiment_path, "evaluation_results.csv")
        pd.json_normalize(safe_r).T.to_csv(csv_path, index=False, encoding="utf-8")

        # create and save plots
        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)
        viz_functions = (plot_roc, plot_far_frr, plot_auc_bars, plot_precision_recall,
                         plot_confusion_bars, plot_score_histograms, plot_confusion_matrix)
        for viz_func in viz_functions:
            try:
                viz_func(r, save_base_path=self.experiment_path)
            except KeyError:
                pass

        evaluation_function_names = r['meta']['evaluation_function_names']

        # Additionally, for threshold-based evaluations with per-model results (e.g., OVR/OVO),
        # create a two-level folder hierarchy <evaluation_function_name>/<model_name>
        # and write per-model plots there (ROC, FAR/FRR, PR, and score histograms).
        try:
            # Determine evaluation function name string for subfolder
            eval_name = None
            if isinstance(evaluation_function_names, list):
                for m in evaluation_function_names:
                    if isinstance(m, dict) and 'threshold_evaluation' in m:
                        eval_name = m['threshold_evaluation'] or 'threshold_evaluation'
                        break
            if not eval_name:
                eval_name = 'threshold_evaluation'

            thr = r.get('threshold_evaluation', {})
            per_model = thr.get('per_model_results', None)
            if isinstance(per_model, dict) and len(per_model) > 0:
                import os
                from copy import deepcopy

                for model_name, model_res in per_model.items():
                    # Construct per-model save directory
                    model_dir = os.path.join(self.experiment_path, str(eval_name), str(model_name))
                    os.makedirs(model_dir, exist_ok=True)

                    # Build a minimal metrics dict focused on this model that matches viz.py expectations
                    # Map per-model roc_curve keys to overall-like structure
                    roc = model_res.get('roc_curve', {})
                    fpr = roc.get('fpr_curve')
                    tpr = roc.get('tpr_curve')
                    thresholds = roc.get('thresholds')
                    auc_val = roc.get('auc_score')
                    if auc_val is None and (fpr is not None) and (tpr is not None):
                        try:
                            from sklearn.metrics import auc as _auc
                            import numpy as _np
                            auc_val = float(_auc(_np.array(fpr), _np.array(tpr)))
                        except Exception:
                            auc_val = None

                    # Per-model overall-style metrics (include counts for confusion bars)
                    om = {
                        'auc': None,
                        'tp': model_res.get('tp'),
                        'tn': model_res.get('tn'),
                        'fp': model_res.get('fp'),
                        'fn': model_res.get('fn'),
                    }
                    if auc_val is not None:
                        om['auc'] = float(auc_val)

                    single_metrics = {
                        'threshold_evaluation': {
                            'overall_metrics': om,
                            'roc_curve': {
                                'fpr': fpr,
                                'tpr': tpr,
                                'thresholds': thresholds,
                                'auc': float(auc_val) if auc_val is not None else None,
                            },
                            # Keep the single model's results for PR/Histograms
                            'per_model_results': {
                                model_name: model_res
                            }
                        },
                        # Hint to PR/histogram functions to treat as one model
                        'meta': {
                            'num_models': 1
                        }
                    }

                    # Generate per-model plots into the model-specific folder
                    try:
                        plot_roc(single_metrics, save_base_path=model_dir)
                    except Exception:
                        pass
                    try:
                        plot_far_frr(single_metrics, save_base_path=model_dir)
                    except Exception:
                        pass
                    try:
                        plot_auc_bars(single_metrics, save_base_path=model_dir)
                    except Exception:
                        pass
                    try:
                        plot_confusion_bars(single_metrics, save_base_path=model_dir)
                    except Exception:
                        pass
                    try:
                        plot_precision_recall(single_metrics, save_base_path=model_dir)
                    except Exception:
                        pass
                    try:
                        plot_score_histograms(single_metrics, save_base_path=model_dir)
                    except Exception:
                        pass

                    # Save the corresponding per-model classifier
                    try:
                        # Find the backing model object by matching key or its string form
                        model_key = None
                        if model_name in self.model_graph:
                            model_key = model_name
                        else:
                            for k in self.model_graph.keys():
                                if str(k) == str(model_name):
                                    model_key = k
                                    break
                        if model_key is not None:
                            self._save_model(self.model_graph[model_key], os.path.join(model_dir, 'model'))
                    except Exception:
                        pass

                    # Persist per-model results JSON into the model folder
                    try:
                        import orjson as _orjson
                        pm_json_path = os.path.join(model_dir, 'per_model_results.json')
                        with open(pm_json_path, 'wb') as fpm:
                            fpm.write(_orjson.dumps(make_json_safe({
                                'evaluation_function': eval_name,
                                'model_name': str(model_name),
                                'results': model_res,
                            }), option=_orjson.OPT_INDENT_2))
                    except Exception:
                        pass

            # Save global model if there is a single global model (CSC, DL)
            try:
                if len(self.model_graph) == 1 and (None in self.model_graph):
                    self._save_model(self.model_graph[None], os.path.join(self.experiment_path, 'model'))
            except Exception:
                pass
        except Exception:
            # Do not fail evaluation due to plotting issues
            pass

    def _save_fit_histories(self):
        """Serialize Keras fit histories to CSV and save per-model history plots.

        - Writes `<experiment_path>/fit_history.csv` aggregating all available histories.
        - Saves `fit_history.png/.pdf` for a single model, or one image per model
          under `<experiment_path>/fit_history/per_model/fit_history_<model>.{png,pdf}` when multiple.
        """
        import os
        import numpy as np
        import pandas as pd
        import matplotlib.pyplot as plt

        # Helper: turn arbitrary key into a readable tag
        def _sanitize_tag(tag: str) -> str:
            allowed = "-_.() abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
            return "".join(c if c in allowed else "_" for c in str(tag)).strip().replace(" ", "_")

        # Collect per-model DataFrames
        per_model_dfs: dict[str, pd.DataFrame] = {}
        plotted_any = False

        for k, v in self.fit_results.items():
            history = None if v is None else v.get("history", None)
            if history is None:
                continue

            # Keras History has .history dict; accept a plain dict as well
            hist_dict = None
            if hasattr(history, "history") and isinstance(history.history, dict):
                hist_dict = history.history
            elif isinstance(history, dict):
                hist_dict = history
            else:
                continue

            if not hist_dict:
                continue

            df = pd.DataFrame(dict(hist_dict))
            # Add 1-based epoch for readability
            df.insert(0, "epoch", np.arange(1, len(df) + 1))

            model_name = _sanitize_tag(k)
            per_model_dfs[str(model_name)] = df.set_index("epoch")

        if not per_model_dfs:
            return  # nothing to persist

        # Aggregate CSV: align by epoch index, flatten columns as `<model>:<metric>`
        combined = pd.concat(per_model_dfs, axis=1)  # MultiIndex (model, metric)
        # Flatten columns
        combined.columns = [f"{m}:{c}" for m, c in combined.columns]
        combined.reset_index(names=["epoch"]).to_csv(os.path.join(self.experiment_path, "fit_history.csv"), index=False)

        # Plotting: single-model => root; multi-model => per_model subfolder
        single_model = len(per_model_dfs) == 1
        out_dir = self.experiment_path
        per_model_dir = os.path.join(out_dir, "fit_history", "per_model") if not single_model else out_dir
        if not os.path.exists(per_model_dir):
            os.makedirs(per_model_dir, exist_ok=True)

        for model_name, df_idx in per_model_dfs.items():
            df = df_idx.reset_index()
            # Derive base metrics (pair train/val where available)
            cols = [c for c in df.columns if c != "epoch"]
            base_metrics = {}
            for c in cols:
                if c.startswith("val_"):
                    base = c[4:]
                    base_metrics.setdefault(base, {"train": None, "val": None})["val"] = c
                else:
                    base = c
                    base_metrics.setdefault(base, {"train": None, "val": None})["train"] = c

            # Create plot
            plt.figure()
            for base, pair in sorted(base_metrics.items()):
                tr_col = pair.get("train")
                va_col = pair.get("val")
                if tr_col is not None:
                    plt.plot(df["epoch"], df[tr_col], label=base)
                if va_col is not None:
                    plt.plot(df["epoch"], df[va_col], linestyle="--", label=f"val_{base}")

            plt.xlabel("Epoch")
            plt.ylabel("Metric value")
            title_suffix = "" if single_model else f" - {model_name}"
            plt.title(f"Training History{title_suffix}")
            plt.legend()
            plt.grid(True)

            # Save figure
            if single_model:
                stem = os.path.join(out_dir, "fit_history")
                plt.tight_layout()
                plt.savefig(f"{stem}.png", bbox_inches="tight")
                plt.savefig(f"{stem}.pdf", bbox_inches="tight")
                plt.close()
            else:
                stem = os.path.join(per_model_dir, f"fit_history_{model_name}")
                plt.tight_layout()
                plt.savefig(f"{stem}.png", bbox_inches="tight")
                plt.savefig(f"{stem}.pdf", bbox_inches="tight")
                plt.close()

            plotted_any = True

        # Optionally: when multiple models, also draw a compact overview of `loss` across models, if present
        if not single_model:
            try:
                plt.figure()
                any_loss = False
                for model_name, df_idx in per_model_dfs.items():
                    df = df_idx.reset_index()
                    if "loss" in df.columns:
                        plt.plot(df["epoch"], df["loss"], label=f"{model_name}:loss")
                        any_loss = True
                if any_loss:
                    plt.xlabel("Epoch")
                    plt.ylabel("Loss")
                    plt.title("Training Loss Overview")
                    plt.legend(); plt.grid(True); plt.tight_layout()
                    stem = os.path.join(out_dir, "fit_history_overview")
                    plt.savefig(f"{stem}.png", bbox_inches="tight")
                    plt.savefig(f"{stem}.pdf", bbox_inches="tight")
                plt.close()
            except Exception:
                pass


def make_json_safe(obj):
    """Recursively convert NumPy/Pandas objects to JSON-serializable Python types and ensure all dict keys are strings."""
    import numpy as np
    import pandas as pd
    from datetime import datetime, date
    from enum import Enum
    from pathlib import Path

    # Dict: ensure str keys + recurse
    if isinstance(obj, dict):
        return {str(k): make_json_safe(v) for k, v in obj.items()}
    # Lists / Tuples
    elif isinstance(obj, list):
        return [make_json_safe(v) for v in obj]
    elif isinstance(obj, tuple):
        return tuple(make_json_safe(v) for v in obj)
    # NumPy arrays / scalars
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, (np.generic,)):
        return obj.item()
    # Pandas structures
    elif isinstance(obj, pd.Series):
        return obj.to_dict()
    elif isinstance(obj, pd.DataFrame):
        return obj.to_dict(orient="records")
    elif isinstance(obj, (pd.Timestamp, pd.Timedelta)):
        return str(obj)
    # Other standard types
    elif isinstance(obj, (datetime, date)):
        return obj.isoformat()
    elif isinstance(obj, (Enum, Path, set)):
        if isinstance(obj, set):
            return [make_json_safe(v) for v in obj]
        return str(obj)
    # Fallback
    else:
        return obj


if __name__ == "__main__":
    from sklearn.neighbors import KNeighborsClassifier
    from useckit.v2.utils.example_dataset import get_chi2021_dataset, get_chi2021_windowsliced_dataset

    # _dataset = get_chi2021_dataset()
    # logging.info(f"_dataset.shape {_dataset.shape}:")

    # original 3D dataset (varying lengths)
    # dataset_regular = copy(_dataset)
    # logging.info(f"dataset_regular.shape {dataset_regular.shape}:")

    # original 3D dataset (with padding)
    #dataset_padded = get_chi2021_dataset(padding=True)
    #logging.info(f"dataset_padded.shape {dataset_padded.shape}:")

    # 2D dataset for sklearn models using feature vector method
    dataset_flattened = get_chi2021_dataset(feature_vector=True)
    logging.info(f"dataset_flattened.shape {dataset_flattened.shape}:")

    # Test sklearn models first:
    for p in Paradigm:
        logging.info(f"Testing Evaluator for paradigm {p}")

        sklearn_estimator = RandomForestClassifier(n_estimators=25)
        if p == Paradigm.AD:
            sklearn_estimator = IsolationForest(n_estimators=25)
        if p == Paradigm.DL:
            continue

        evaluator = Evaluator(
            paradigm=p,
            dataset=dataset_flattened,
            model_initializer_callable=sklearn_estimator,
        )

        evaluator.fit()

        evaluator.evaluate()
