import logging
import itertools
from typing import Any, Dict, Tuple, Iterable, Optional, List
import numpy as np

from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    classification_report,
    roc_curve,
    roc_auc_score, auc,
)

from useckit.v2.dataset import Dataset, WindowslicedDataset
from useckit.v2.evaluator import Evaluator
from useckit.v2.keras_model_builder import BuildResult



def accuracy_evaluation(dataset: Dataset, model_graph: dict):
    """
    Compute accuracy, confusion matrix, and classification reports on the matching test set.

    Returns a dict with keys:
      - "accuracy": float
      - "confusion_matrix": ndarray
      - "classification_report_txt": string version of sklearn.metrics.classification_report
      - "classification_report": dict version via classification_report(..., output_dict=True)
    """
    # Obtain the global model for CSC paradigm
    model = model_graph.get(None)
    if model is None:
        logging.error("No global model found for accuracy evaluation.")
        raise ValueError("No model available for accuracy evaluation.")

    if isinstance(model, BuildResult):
        model = model["model"]

    X = dataset.get_matching_data()
    y_true = dataset.get_matching_labels(ohc=Evaluator.is_deeplearning_model(model))

    if X is None or y_true is None:
        raise ValueError("Matching test data or labels are not available for evaluation.")

    try:
        y_pred = model.predict(X)
    except ValueError as ve:
        if "Found array with dim 3, while dim <= 2 is required" in str(ve):
            y_pred = model.predict(X.reshape(X.shape[0], -1))

    # If this is a window-sliced dataset, aggregate per-sample via voting
    if isinstance(dataset, WindowslicedDataset) and getattr(dataset, "testset_matching_sliceorigin_mask", None) is not None:
        y_pred, y_true = dataset.apply_voting_for_testset_matching_slices(np.asarray(y_pred), np.asarray(y_true))
        # Remove any positions where voting yielded None
        mask = np.array([(p is not None) and (t is not None) for p, t in zip(y_pred, y_true)])
        y_pred = np.asarray(y_pred)[mask]
        y_true = np.asarray(y_true)[mask]
        ohc_width = dataset.get_matching_labels(ohc=True).shape[1]

        if Evaluator.is_deeplearning_model(model):
            # reformat y_pred to ohc format
            y_pred_ohc = np.zeros([y_pred.shape[0], ohc_width])
            y_pred_ohc[np.arange(y_pred.shape[0]), y_pred.astype(int)] = 1

            # reformat y_true to ohc format
            y_true_ohc = np.zeros([y_true.shape[0], ohc_width])
            y_true_ohc[np.arange(y_true.shape[0]), y_true.astype(int)] = 1

            # overwrite
            y_pred = y_pred_ohc
            y_true = y_true_ohc

    # Ensure numpy arrays
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    # Backtransform from ohc to class numbers in case of Deep Learning
    if Evaluator.is_deeplearning_model(model) and (y_pred.ndim == 2 or y_true.ndim == 2):
        one_hot2d = np.zeros_like(y_pred, dtype=int)
        one_hot2d[np.arange(y_pred.shape[0]), np.argmax(y_pred, axis=1)] = 1
        y_pred = dataset.ohc_labels_to_classes(one_hot2d)
        y_true = dataset.ohc_labels_to_classes(y_true)

    # Compute metrics
    # Use consistent label order for confusion matrix when possible
    labels_for_cm = None
    try:
        labels_for_cm = dataset.get_unique_classes()
    except Exception:
        pass

    acc = float(accuracy_score(y_true, y_pred))
    # Use a fixed label ordering for both the confusion matrix and reports
    cm = confusion_matrix(y_true, y_pred, labels=labels_for_cm)
    # Produce both string and dict forms of the classification report
    report_txt = classification_report(y_true, y_pred, labels=labels_for_cm, zero_division=0)
    report_dict = classification_report(y_true, y_pred, labels=labels_for_cm, zero_division=0, output_dict=True)

    return {
        "accuracy": acc,
        "confusion_matrix": cm,
        "confusion_matrix_labels": labels_for_cm,
        "classification_report_txt": report_txt,
        "classification_report": report_dict,
    }


def multiclass_auc_evaluation(dataset: Dataset, model_graph: dict):
    """
    Compute multiclass ROC-AUC metrics on the matching test set using
    One-vs-Rest (OvR) and One-vs-One (OvO) schemes.

    Returns a dict with keys similar in spirit to other evaluators:
      - 'ovr': micro and macro averaged ROC-AUC and per-class OvR AUC
      - 'ovo': macro averaged ROC-AUC
      - 'classes': numpy array of class labels used for columns in y_score

    Notes
    - Uses predict_proba when available; otherwise uses decision_function;
      as a last resort, uses predict if it returns a 2D array of scores.
    - For window-sliced datasets, probabilities are averaged per original
      sample using the slice-origin mask.
    """

    print("FIXME multiclass_auc not implemented yet; reuse initial code from above")
    return

    # Obtain the global model for CSC paradigm
    model = model_graph.get(None)
    if model is None:
        logging.error("No global model found for multiclass AUC evaluation.")
        raise ValueError("No model available for multiclass AUC evaluation.")

    X = dataset.get_matching_data()
    y_true_slices = dataset.get_matching_labels()

    if X is None or y_true_slices is None:
        raise ValueError("Matching test data or labels are not available for multiclass AUC evaluation.")

    # Get scores/probabilities
    y_score_slices = None
    if hasattr(model, "predict_proba") and callable(getattr(model, "predict_proba")):
        try:
            y_score_slices = model.predict_proba(X)
        except Exception as e:
            logging.warning(f"predict_proba failed with {e!r}; trying decision_function/predict as fallback.")
            y_score_slices = None

    if y_score_slices is None and hasattr(model, "decision_function") and callable(getattr(model, "decision_function")):
        try:
            y_score_slices = model.decision_function(X)
        except Exception:
            y_score_slices = None

    if y_score_slices is None:
        # Last resort: some models (e.g., Keras) may return probabilities via predict
        try:
            pred = model.predict(X)
            pred = np.asarray(pred)
            if pred.ndim == 2 and pred.shape[0] == len(X):
                y_score_slices = pred
        except Exception:
            pass

    if y_score_slices is None:
        raise RuntimeError("Could not obtain class scores/probabilities from the model for multiclass AUC.")

    # Ensure 2D array
    y_score_slices = np.asarray(y_score_slices)
    if y_score_slices.ndim != 2:
        raise ValueError("Model output for scores must be a 2D array of shape (n_samples, n_classes).")

    # If dataset is window-sliced, aggregate slice-level scores to sample-level by averaging
    if isinstance(dataset, WindowslicedDataset) and getattr(dataset, "testset_matching_sliceorigin_mask", None) is not None:
        origin_mask = dataset.testset_matching_sliceorigin_mask
        n_orig = origin_mask.shape[0]
        y_score = np.zeros((n_orig, y_score_slices.shape[1]), dtype=float)
        y_true = np.empty((n_orig,), dtype=y_true_slices.dtype)

        for i in range(n_orig):
            mask = origin_mask[i]
            if np.any(mask):
                y_score[i] = np.mean(y_score_slices[mask], axis=0)
                # assuming label consistency among slices belonging to same origin
                y_true[i] = y_true_slices[mask][0]
            else:
                # No slices for this origin; mark as NaN-probabilities and None label
                y_score[i] = np.nan
                y_true[i] = None

        # Filter out empty/orphaned origins
        valid = np.isfinite(y_score).all(axis=1) & (y_true != None)  # noqa: E711
        y_score = y_score[valid]
        y_true = np.asarray(y_true)[valid]
    else:
        y_score = y_score_slices
        y_true = np.asarray(y_true_slices)

    # Determine class ordering and align y_score columns if possible
    # Use the unique classes observed in y_true as the canonical order
    classes_true = np.unique(y_true)

    # Attempt to reorder columns according to model.classes_ if available
    if hasattr(model, "classes_"):
        try:
            model_classes = np.asarray(model.classes_)
            # Build index mapping from desired classes_true to model_classes columns
            idx = [int(np.where(model_classes == c)[0][0]) for c in classes_true]
            y_score = y_score[:, idx]
        except Exception:
            # Fallback: assume current order already matches classes_true
            pass

    # Compute multiclass AUCs
    # OvR micro and macro
    try:
        micro_ovr = float(roc_auc_score(y_true, y_score, multi_class="ovr", average="micro"))
    except Exception as e:
        logging.warning(f"Failed computing micro-averaged OvR AUC: {e!r}")
        micro_ovr = np.nan

    try:
        macro_ovr = float(roc_auc_score(y_true, y_score, multi_class="ovr", average="macro"))
    except Exception as e:
        logging.warning(f"Failed computing macro-averaged OvR AUC: {e!r}")
        macro_ovr = np.nan

    # OvO macro
    try:
        macro_ovo = float(roc_auc_score(y_true, y_score, multi_class="ovo", average="macro"))
    except Exception as e:
        logging.warning(f"Failed computing macro-averaged OvO AUC: {e!r}")
        macro_ovo = np.nan

    # Per-class One-vs-Rest AUC values
    per_class_auc: Dict[Any, float] = {}
    try:
        from sklearn.preprocessing import LabelBinarizer

        lb = LabelBinarizer()
        lb.fit(classes_true)
        y_onehot = lb.transform(y_true)
        # Ensure two-dimensional for binary corner-case
        if y_onehot.ndim == 1:
            y_onehot = np.column_stack([1 - y_onehot, y_onehot])
        for i, cls in enumerate(lb.classes_):
            try:
                per_class_auc[cls] = float(roc_auc_score(y_onehot[:, i], y_score[:, i]))
            except Exception:
                per_class_auc[cls] = np.nan
    except Exception:
        # If LabelBinarizer is unavailable for some reason, skip per-class
        per_class_auc = {cls: np.nan for cls in classes_true}

    return {
        "ovr": {
            "micro_roc_auc": micro_ovr,
            "macro_roc_auc": macro_ovr,
            "per_class_roc_auc": per_class_auc,
        },
        "ovo": {
            "macro_roc_auc": macro_ovo,
        },
        "classes": classes_true,
    }


def ___threshold_evaluation_deprecated(dataset: Dataset, model_graph: dict,
                         determine_threshold_on_enrollmentdata=False,
                         determine_threshold_on_matchingdata=True,
                         threshold: float | None = None,
                         ws_streaming_enroll_chunk_size: int = 64,
                         ws_streaming_match_chunk_size: int = 4096):
    # Enforce guard only when auto-determining threshold
    if threshold is None:
        if not (determine_threshold_on_enrollmentdata ^ determine_threshold_on_matchingdata):
            raise ValueError("Only one of determine_threshold_on_enrollmentdata or "
                             "determine_threshold_on_matchingdata must be True.")

    # Collect all (score, ground_truth) pairs across all models
    all_scores = []
    all_ground_truths = []
    all_thresholds = []
    model_results = {}

    data_accept, data_reject = None, None

    ###################
    # DATA COLLECTION #
    ###################

    for k, v in zip(model_graph.keys(), model_graph.values()):
        model = v

        # Generator function for streaming data
        def contrastive_batch_generator():
            """Generator that yields batches of (left, right, labels) for contrastive learning"""
            enroll_X = dataset.get_enrollment_data()
            enroll_y = dataset.get_enrollment_labels()
            match_X = dataset.get_matching_data()
            match_y = dataset.get_matching_labels()

            if enroll_X is None or len(enroll_X) == 0 or match_X is None or len(match_X) == 0:
                raise ValueError("No data available for contrastive evaluation.")

            n_e = len(enroll_X)
            n_m = len(match_X)

            chunk_size_e = max(1, int(ws_streaming_enroll_chunk_size))
            chunk_size_m = max(1, int(ws_streaming_match_chunk_size))

            e_start = 0
            while e_start < n_e:
                e_end = min(e_start + chunk_size_e, n_e)
                e_chunk = enroll_X[e_start:e_end]
                y_e = np.asarray(enroll_y[e_start:e_end])

                m_start = 0
                while m_start < n_m:
                    m_end = min(m_start + chunk_size_m, n_m)
                    m_chunk = match_X[m_start:m_end]
                    y_m = np.asarray(match_y[m_start:m_end])

                    # Build E×M pairs only within this chunk to bound memory
                    left_batch = np.repeat(np.asarray(e_chunk), repeats=len(m_chunk), axis=0)
                    right_batch = np.tile(np.asarray(m_chunk),
                                          reps=(len(e_chunk),) + (1,) * (m_chunk.ndim - 1))
                    labels_batch = (y_e[:, None] == y_m[None, :]).astype(np.int32).ravel()

                    yield left_batch, right_batch, labels_batch

                    m_start = m_end

                e_start = e_end
        # End of generator function

        # Variable to hold data generator/batches
        batch_generator = None
        total_batches = None

        if isinstance(k, (tuple, frozenset)):  # if OVO
            class1 = k[0]
            class2 = k[1]
            if threshold is not None:
                # When a fixed threshold is provided, evaluate only on matching/test data
                data_accept = dataset.get_matching_data(filter=lambda x: x == class1)
                data_reject = dataset.get_matching_data(filter=lambda x: x == class2)
            else:
                if determine_threshold_on_enrollmentdata:
                    data_accept = dataset.get_enrollment_data(filter=lambda x: x == class1)
                    data_reject = dataset.get_enrollment_data(filter=lambda x: x == class2)
                elif determine_threshold_on_matchingdata:
                    data_accept = dataset.get_matching_data(filter=lambda x: x == class1)
                    data_reject = dataset.get_matching_data(filter=lambda x: x == class2)
        else:  # if OVR or something else
            if len(model_graph.keys()) > 1:  # OVR: more than one model in the A vs. B graph
                class1 = k
                if threshold is not None:
                    # When a fixed threshold is provided, evaluate only on matching/test data
                    data_accept = dataset.get_matching_data(filter=lambda x: x == class1)
                    data_reject = dataset.get_matching_data(filter=lambda x: x != class1)
                else:
                    if determine_threshold_on_enrollmentdata:
                        data_accept = dataset.get_enrollment_data(filter=lambda x: x == class1)
                        data_reject = dataset.get_enrollment_data(filter=lambda x: x != class1)
                    elif determine_threshold_on_matchingdata:
                        data_accept = dataset.get_matching_data(filter=lambda x: x == class1)
                        data_reject = dataset.get_matching_data(filter=lambda x: x != class1)
            elif len(model_graph.keys()) == 1 and isinstance(model_graph[None], BuildResult):  # only a single global model provided
                if "contrastive" in model.metadata["paradigm"]:
                    if isinstance(dataset, WindowslicedDataset):
                        # streaming for WindowSlicedDataset - prepare generator
                        enroll_X = dataset.get_enrollment_data()
                        match_X = dataset.get_matching_data()

                        if (enroll_X is None or len(enroll_X) == 0 or
                                match_X is None or len(match_X) == 0):
                            raise ValueError("No data available for contrastive evaluation.")

                        n_e = len(enroll_X)
                        n_m = len(match_X)

                        chunk_size_e = max(1, int(ws_streaming_enroll_chunk_size))
                        chunk_size_m = max(1, int(ws_streaming_match_chunk_size))

                        # Calculate total number of chunk pairs
                        n_e_chunks = (n_e + chunk_size_e - 1) // chunk_size_e
                        n_m_chunks = (n_m + chunk_size_m - 1) // chunk_size_m
                        total_batches = n_e_chunks * n_m_chunks

                        batch_generator = contrastive_batch_generator()
                    else:
                        # materialized path - no generator needed
                        pass
                elif "triplet" in model.metadata["paradigm"]:
                    pass
                else:
                    raise ValueError(f"Unknown/unhandled model.metadata['paradigm']: {model.metadata['paradigm']}")
            else:
                raise ValueError(f"Could not parse (unhandled case) for model_graph: {model_graph}")

        ###################
        # PREDICTION STEP #
        ###################

        # if we have a populated model graph vs. we use a singular global distance learning model
        if len(model_graph.keys()) > 1:  # ensemble of OVR/OVO models
            # verifying that we have proper data selected
            if data_accept is None or len(data_accept) == 0:
                raise ValueError("No data marked to be accepted.")
            if data_reject is None or len(data_reject) == 0:
                raise ValueError("No data marked to be rejected.")

            # Get predictions/scores from the model
            scores_accept = get_scores(model, data_accept)
            scores_reject = get_scores(model, data_reject)

            # Combine scores and create ground truth labels
            scores = np.concatenate([scores_accept, scores_reject])  # format: 1D list with scores
            ground_truth = np.concatenate(
                [np.ones(len(scores_accept)), np.zeros(len(scores_reject))])  # 1=accept, 0=reject

            del scores_accept, scores_reject  # not needed anymore further down
            # from here on we just require: scores, ground_truth
        elif (len(model_graph.keys()) == 1) and (None in model_graph.keys()) and ("contrastive" in model.metadata["paradigm"]):
            # Window slicing path that performs predictions using the generator
            if isinstance(dataset, WindowslicedDataset):
                if batch_generator is None or total_batches is None:
                    raise ValueError("Contrastive batch generator was not initialized")

                streamed_scores = []
                streamed_labels = []

                from tqdm import tqdm
                with tqdm(total=total_batches, desc="Processing contrastive predictions", unit="batch") as pbar:
                    for left_batch, right_batch, labels_batch in batch_generator:
                        # Predict on this batch
                        scores_batch = model.model.predict([left_batch, right_batch], verbose=0).ravel()

                        streamed_scores.append(scores_batch)
                        streamed_labels.append(labels_batch)

                        del left_batch, right_batch, labels_batch, scores_batch
                        pbar.update(1)

                scores = np.concatenate(streamed_scores)
                ground_truth = np.concatenate(streamed_labels)

                del streamed_scores, streamed_labels
            else:
                # materialized path for regular predictions (non-ws)
                cartesian_pairs = list(itertools.product(dataset.get_enrollment_data(),
                                                         dataset.get_matching_data()))
                cartesian_label_pairs = list(itertools.product(dataset.get_enrollment_labels(),
                                                               dataset.get_matching_labels()))
                ground_truth = np.fromiter(
                    (1 if l1 == l2 else 0 for (l1, l2) in cartesian_label_pairs), dtype=int
                )

                if len(cartesian_pairs) == 0:
                    raise ValueError("No pairs produced from itertools.product; check dataset sizes.")

                left_list, right_list = zip(*cartesian_pairs)  # tuples of arrays
                left_batch = np.asarray(left_list)
                right_batch = np.asarray(right_list)

                # Predict scores for all pairs at once
                scores = model.model.predict([left_batch, right_batch], verbose=0).ravel()

                del cartesian_pairs, left_list, right_list, cartesian_label_pairs, left_batch, right_batch
        else:
            raise ValueError(f"Unhandled case for model_graph: {model_graph}.")

        ################
        # THRESHOLDING #
        ################

        # Determine threshold to use
        if threshold is None:
            # Calculate ROC curve to find EER threshold
            fpr, tpr, thresholds = roc_curve(ground_truth, scores)
            fnr = 1 - tpr

            # Find EER (where FPR = FNR)
            eer_idx = np.nanargmin(np.abs(fpr - fnr))
            eer_threshold = thresholds[eer_idx]
            eer = (fpr[eer_idx] + fnr[eer_idx]) / 2
        else:
            # Use provided threshold; skip ROC/EER computation
            fpr = tpr = thresholds = None
            eer_threshold = float(threshold)
            eer = None

        # Store results for this model
        all_scores.extend(scores)
        all_ground_truths.extend(ground_truth)
        all_thresholds.append(eer_threshold)

        # Apply threshold to get binary predictions
        predictions = (scores >= eer_threshold).astype(int)

        # Calculate confusion matrix
        tn, fp, fn, tp = confusion_matrix(ground_truth, predictions).ravel()

        # Calculate metrics
        if tp + fn > 0:
            tpr_at_threshold = tp / (tp + fn)
            fnr_at_threshold = fn / (tp + fn)
        else:
            tpr_at_threshold = 0
            fnr_at_threshold = 1

        if tn + fp > 0:
            tnr_at_threshold = tn / (tn + fp)
            fpr_at_threshold = fp / (tn + fp)
        else:
            tnr_at_threshold = 0
            fpr_at_threshold = 1

        # Additional metrics
        accuracy = (tp + tn) / (tp + tn + fp + fn) if (tp + tn + fp + fn) > 0 else 0
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0
        f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0

        # Matthews Correlation Coefficient
        denominator = np.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))
        mcc = (tp * tn - fp * fn) / denominator if denominator > 0 else 0

        model_results[k] = {
            'eer': eer,
            'eer_threshold': eer_threshold,
            'tp': tp,
            'tn': tn,
            'fp': fp,
            'fn': fn,
            'tpr': tpr_at_threshold,
            'tnr': tnr_at_threshold,
            'fpr': fpr_at_threshold,
            'fnr': fnr_at_threshold,
            'accuracy': accuracy,
            'precision': precision,
            'recall': recall,
            'f1_score': f1_score,
            'mcc': mcc,
            'scores': scores,
            'ground_truth': ground_truth,
            'fpr_curve': fpr,
            'tpr_curve': tpr,
            'thresholds_curve': thresholds,
            'auc': roc_auc_score(ground_truth, scores)
        }

    # Calculate overall metrics across all models
    all_scores = np.array(all_scores)
    all_ground_truths = np.array(all_ground_truths)

    if threshold is None:
        # Overall ROC curve
        overall_fpr, overall_tpr, overall_thresholds = roc_curve(all_ground_truths, all_scores)
        overall_fnr = 1 - overall_tpr

        # Overall EER
        overall_eer_idx = np.nanargmin(np.abs(overall_fpr - overall_fnr))
        overall_eer = (overall_fpr[overall_eer_idx] + overall_fnr[overall_eer_idx]) / 2
        overall_eer_threshold = overall_thresholds[overall_eer_idx]

        # Overall AUC
        overall_auc = auc(overall_fpr, overall_tpr)
    else:
        # When user supplies a fixed threshold, evaluate on matching data only
        # and avoid computing ROC/EER. We still compute AUROC since it does not
        # require a chosen threshold.
        overall_fpr = overall_tpr = overall_thresholds = None
        overall_eer = None
        overall_eer_threshold = float(threshold)
        overall_auc = roc_auc_score(all_ground_truths, all_scores)

    # Apply overall threshold
    overall_predictions = (all_scores >= overall_eer_threshold).astype(int)
    overall_tn, overall_fp, overall_fn, overall_tp = confusion_matrix(all_ground_truths, overall_predictions).ravel()

    # Calculate overall metrics
    overall_accuracy = (overall_tp + overall_tn) / len(all_ground_truths)
    overall_precision = overall_tp / (overall_tp + overall_fp) if (overall_tp + overall_fp) > 0 else 0
    overall_recall = overall_tp / (overall_tp + overall_fn) if (overall_tp + overall_fn) > 0 else 0
    overall_f1 = 2 * (overall_precision * overall_recall) / (overall_precision + overall_recall) if (
                                                                                                                overall_precision + overall_recall) > 0 else 0

    results = {
        'overall_metrics': {
            'eer': overall_eer,
            'eer_threshold': overall_eer_threshold,
            'auc': overall_auc,
            'tp': overall_tp,
            'tn': overall_tn,
            'fp': overall_fp,
            'fn': overall_fn,
            'accuracy': overall_accuracy,
            'precision': overall_precision,
            'recall': overall_recall,
            'f1_score': overall_f1,
            'fpr': overall_fp / (overall_fp + overall_tn) if (overall_fp + overall_tn) > 0 else 0,
            'fnr': overall_fn / (overall_fn + overall_tp) if (overall_fn + overall_tp) > 0 else 0,
            'tpr': overall_tp / (overall_tp + overall_fn) if (overall_tp + overall_fn) > 0 else 0,
            'tnr': overall_tn / (overall_tn + overall_fp) if (overall_tn + overall_fp) > 0 else 0
        },
        'per_model_results': model_results,
        'roc_curve': {
            'fpr': overall_fpr,
            'tpr': overall_tpr,
            'thresholds': overall_thresholds
        }
    }
    return results



















# -------------------------------------------------------------------------------------------------------------------- #

# -----------------------------
# Helper utilities and helpers
# -----------------------------

def get_scores(model, data):
    meta = None

    try:
        if Evaluator.is_deeplearning_model(model):
            if isinstance(model, BuildResult):
                meta = model.metadata
                model = model["model"]

            pred = model.predict(data)

            if meta["paradigm"] == "autoencoder_oneclass":
                # calculate RMSE between inference input and predictions
                from sklearn.metrics import root_mean_squared_error
                return np.array([root_mean_squared_error(p, d) for p, d in zip(pred, data)])
            else:
                return pred
        else:
            try:
                return model.predict_proba(data)[:, 1]
            except ValueError as ve:
                if "Found array with dim 3, while dim <= 2 is required" in str(ve):
                    return model.predict_proba(data.reshape(data.shape[0], -1))[:, 1]
    except AttributeError as ae:
        if "has no attribute 'predict_proba'" in str(ae):
            try:
                return model.score_samples(data)
            except AttributeError:
                raise AttributeError(
                    f"{ae}. Neither 'predict_proba()' nor 'score_samples()' are available. "
                    f"Please check the documentation of your model. "
                    f"For example, for sklearn's SVMs you must set `probability=True` in their constructor "
                    f"to enable 'predict_proba()'."
                )
            except ValueError as ve:
                if "Found array with dim 3, while dim <= 2 is required" in str(ve):
                    return model.score_samples(data.reshape(data.shape[0], -1))
        else:
            raise

def _ensure_nonempty(X, name: str):
    if X is None or len(X) == 0:
        raise ValueError(f"No data available for {name}.")


def _contrastive_batch_generator(enroll_X: np.ndarray,
                                 enroll_y: np.ndarray,
                                 match_X: np.ndarray,
                                 match_y: np.ndarray,
                                 chunk_size_e: int,
                                 chunk_size_m: int) -> Iterable[Tuple[np.ndarray, np.ndarray, np.ndarray]]:
    """
    Generates batches of (left_batch, right_batch, labels_batch) using chunked E×M pairing.
    """
    _ensure_nonempty(enroll_X, "contrastive enrollment data")
    _ensure_nonempty(match_X, "contrastive matching data")

    n_e = len(enroll_X)
    n_m = len(match_X)

    e_start = 0
    while e_start < n_e:
        e_end = min(e_start + max(1, int(chunk_size_e)), n_e)
        e_chunk = enroll_X[e_start:e_end]
        y_e = np.asarray(enroll_y[e_start:e_end])

        m_start = 0
        while m_start < n_m:
            m_end = min(m_start + max(1, int(chunk_size_m)), n_m)
            m_chunk = match_X[m_start:m_end]
            y_m = np.asarray(match_y[m_start:m_end])

            left_batch = np.repeat(np.asarray(e_chunk), repeats=len(m_chunk), axis=0)
            right_batch = np.tile(np.asarray(m_chunk), reps=(len(e_chunk),) + (1,) * (m_chunk.ndim - 1))
            labels_batch = (y_e[:, None] == y_m[None, :]).astype(np.int32).ravel()

            yield left_batch, right_batch, labels_batch
            m_start = m_end

        e_start = e_end


def _compute_threshold_and_metrics(scores: np.ndarray,
                                   ground_truth: np.ndarray,
                                   threshold: Optional[float]):
    """
    Compute EER (if threshold None), or use provided threshold.
    Returns: metrics_dict: dict
    """
    if threshold is None:
        fpr_curve, tpr_curve, thresholds = roc_curve(ground_truth, scores)
        fnr = 1 - tpr_curve
        eer_idx = np.nanargmin(np.abs(fpr_curve - fnr))
        used_threshold = float(thresholds[eer_idx])
        eer = float((fpr_curve[eer_idx] + fnr[eer_idx]) / 2)
        info = "Metrics calculated via sklearn.roc_curve(ground_truth, scores)."
    else:
        used_threshold = float(threshold)
        eer = fpr_curve = tpr_curve = thresholds = None
        info = ("You specified threshold {threshold} to be used, so I did not calculate a sklearn.roc_curve. "
                "Setting variables to None.")

    predictions = (scores >= used_threshold).astype(int)
    tn, fp, fn, tp = confusion_matrix(ground_truth, predictions).ravel()

    # Rates
    try:
        tpr = tp / (tp + fn)
    except ZeroDivisionError:
        print("TPR undefined (division by zero), setting to 0.0.")
        tpr = 0.0

    try:
        tnr = tn / (tn + fp)
    except ZeroDivisionError:
        print("TNR undefined (division by zero), setting to 0.0.")
        tnr = 0.0

    try:
        fnr = fn / (tp + fn)
    except ZeroDivisionError:
        print("FNR undefined (division by zero), setting to 1.0.")
        fnr = 1.0

    try:
        fpr = fp / (tn + fp)
    except ZeroDivisionError:
        print("FPR undefined (division by zero), setting to 1.0.")
        fpr = 1.0

    # Additional metrics
    accuracy = (tp + tn) / (tp + tn + fp + fn)
    precision = tp / (tp + fp)
    recall = tp / (tp + fn)
    f1_score = 2 * (precision * recall) / (precision + recall)

    mcc = (tp * tn - fp * fn) / ((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))**0.5

    auc_score = roc_auc_score(ground_truth, scores)

    return {
        'eer': eer,
        'eer_threshold': used_threshold,
        'tp': int(tp),
        'tn': int(tn),
        'fp': int(fp),
        'fn': int(fn),
        'tpr': float(tpr),
        'tnr': float(tnr),
        'fpr': float(fpr),
        'fnr': float(fnr),
        'accuracy': float(accuracy),
        'precision': float(precision),
        'recall': float(recall),
        'f1_score': float(f1_score),
        'mcc': float(mcc),
        'roc_curve': {
            'fpr_curve': fpr_curve,
            'tpr_curve': tpr_curve,
            'thresholds': thresholds,
            'auc_score': auc_score,
            'info': info,
        },
        'scores': scores,
        'ground_truth': ground_truth,
        'supplied_threshold': threshold
    }


def _overall_results(model_results: Dict[Any, Dict[str, Any]],
                     all_scores: List[float],
                     all_ground_truths: List[int],
                     used_overall_threshold: Optional[float]) -> Dict[str, Any]:
    all_scores_np = np.asarray(all_scores)
    all_gt_np = np.asarray(all_ground_truths)

    if used_overall_threshold is None:
        overall_fpr, overall_tpr, overall_thresholds = roc_curve(all_gt_np, all_scores_np)
        overall_fnr = 1 - overall_tpr
        idx = np.nanargmin(np.abs(overall_fpr - overall_fnr))
        overall_eer = float((overall_fpr[idx] + overall_fnr[idx]) / 2)
        overall_eer_threshold = float(overall_thresholds[idx])
        overall_auc = float(auc(overall_fpr, overall_tpr))
    else:
        overall_fpr = overall_tpr = overall_thresholds = None
        overall_eer = None
        overall_eer_threshold = float(used_overall_threshold)
        overall_auc = float(roc_auc_score(all_gt_np, all_scores_np))

    overall_predictions = (all_scores_np >= overall_eer_threshold).astype(int)
    overall_tn, overall_fp, overall_fn, overall_tp = confusion_matrix(all_gt_np, overall_predictions).ravel()

    overall_accuracy = (overall_tp + overall_tn) / len(all_gt_np) if len(all_gt_np) > 0 else 0.0
    overall_precision = overall_tp / (overall_tp + overall_fp) if (overall_tp + overall_fp) > 0 else 0.0
    overall_recall = overall_tp / (overall_tp + overall_fn) if (overall_tp + overall_fn) > 0 else 0.0
    overall_f1 = 2 * (overall_precision * overall_recall) / (overall_precision + overall_recall) if (overall_precision + overall_recall) > 0 else 0.0

    results = {
        'overall_metrics': {
            'eer': overall_eer,
            'eer_threshold': overall_eer_threshold,
            'auc': overall_auc,
            'tp': int(overall_tp),
            'tn': int(overall_tn),
            'fp': int(overall_fp),
            'fn': int(overall_fn),
            'accuracy': float(overall_accuracy),
            'precision': float(overall_precision),
            'recall': float(overall_recall),
            'f1_score': float(overall_f1),
            'fpr': float(overall_fp / (overall_fp + overall_tn)) if (overall_fp + overall_tn) > 0 else 0.0,
            'fnr': float(overall_fn / (overall_fn + overall_tp)) if (overall_fn + overall_tp) > 0 else 0.0,
            'tpr': float(overall_tp / (overall_tp + overall_fn)) if (overall_tp + overall_fn) > 0 else 0.0,
            'tnr': float(overall_tn / (overall_tn + overall_fp)) if (overall_tn + overall_fp) > 0 else 0.0
        },
        'per_model_results': model_results,
        'roc_curve': {
            'fpr': overall_fpr,
            'tpr': overall_tpr,
            'thresholds': overall_thresholds,
            'auc': overall_auc,
        }
    }
    return results


# -----------------------------
# 1) Class-vs-Class graph ensemble (OVO)
# -----------------------------

def threshold_evaluation_via_class_vs_class_graph_ensemble(
    dataset: Any,
    model_graph: Dict[Tuple[Any, Any], Any],
    determine_threshold_on_enrollmentdata: bool = False,
    determine_threshold_on_matchingdata: bool = True,
    threshold: Optional[float] = None
) -> Dict[str, Any]:
    """
    Evaluate an ensemble where each key in model_graph is a tuple (class1, class2).
    Uses accept=class1 data, reject=class2 data for each model.
    """
    if threshold is None:
        if not (determine_threshold_on_enrollmentdata ^ determine_threshold_on_matchingdata):
            raise ValueError("Only one of determine_threshold_on_enrollmentdata or determine_threshold_on_matchingdata must be True.")

    model_results: Dict[Any, Dict[str, Any]] = {}
    all_scores: List[float] = []
    all_gt: List[int] = []
    used_thresholds: List[float] = []

    for (class1, class2), model in model_graph.items():
        if threshold is not None:
            data_accept = dataset.get_matching_data(filter=lambda x: x == class1)
            data_reject = dataset.get_matching_data(filter=lambda x: x == class2)
        else:
            if determine_threshold_on_enrollmentdata:
                data_accept = dataset.get_enrollment_data(filter=lambda x: x == class1)
                data_reject = dataset.get_enrollment_data(filter=lambda x: x == class2)
            else:
                data_accept = dataset.get_matching_data(filter=lambda x: x == class1)
                data_reject = dataset.get_matching_data(filter=lambda x: x == class2)

        _ensure_nonempty(data_accept, "accepted class data")
        _ensure_nonempty(data_reject, "rejected class data")

        scores_accept = get_scores(model, np.asarray(data_accept))
        scores_reject = get_scores(model, np.asarray(data_reject))

        scores = np.concatenate([scores_accept, scores_reject])
        gt = np.concatenate([np.ones(len(scores_accept)), np.zeros(len(scores_reject))]).astype(int)

        metrics = _compute_threshold_and_metrics(scores, gt, threshold)

        model_results[(class1, class2)] = {
            **metrics,
            'scores': scores,
            'ground_truth': gt,
            'threshold': threshold,
            'determine_threshold_on_enrollmentdata': determine_threshold_on_enrollmentdata,
            'determine_threshold_on_matchingdata': determine_threshold_on_matchingdata,
        }

        all_scores.extend(scores.tolist())
        all_gt.extend(gt.tolist())
        used_thresholds.append(threshold)

    # Use mean of per-model thresholds as "overall" fixed threshold if user provided a fixed one; otherwise None to compute EER globally
    overall_used_threshold = float(threshold) if threshold is not None else None
    return _overall_results(model_results, all_scores, all_gt, overall_used_threshold)


# -----------------------------
# 2) One-model-per-class graph ensemble (OVR)
# -----------------------------

def threshold_evaluation_via_one_model_per_class(
    dataset: Dataset,
    model_graph: Dict[Any, Any],
    determine_threshold_on_enrollmentdata: bool = False,
    determine_threshold_on_matchingdata: bool = True,
    threshold: Optional[float] = None
) -> Dict[str, Any]:
    """
    Evaluate an ensemble where each key in model_graph is a class label (OVR).
    For each model keyed by class k: accept = class k, reject = not class k.
    """
    if threshold is None:
        if not (determine_threshold_on_enrollmentdata ^ determine_threshold_on_matchingdata):
            raise ValueError("Only one of determine_threshold_on_enrollmentdata or determine_threshold_on_matchingdata must be True.")

    model_results = dict()
    all_scores = []
    all_ground_truths = []

    for k, model in model_graph.items():
        if threshold is not None:
            data_accept = dataset.get_matching_data(filter=lambda x: x == k)
            data_reject = dataset.get_matching_data(filter=lambda x: x != k)
        else:
            if determine_threshold_on_enrollmentdata:
                data_accept = dataset.get_enrollment_data(filter=lambda x: x == k)
                data_reject = dataset.get_enrollment_data(filter=lambda x: x != k)
            else:
                data_accept = dataset.get_matching_data(filter=lambda x: x == k)
                data_reject = dataset.get_matching_data(filter=lambda x: x != k)

        _ensure_nonempty(data_accept, "accepted class data")
        _ensure_nonempty(data_reject, "rejected class data")

        scores_accept = get_scores(model, np.asarray(data_accept))
        scores_reject = get_scores(model, np.asarray(data_reject))

        scores = np.concatenate([scores_accept, scores_reject])
        gt = np.concatenate([np.ones(len(scores_accept)), np.zeros(len(scores_reject))]).astype(int)

        metrics = _compute_threshold_and_metrics(scores, gt, threshold)

        model_results[k] = {
            **metrics,
            'scores': scores,
            'ground_truth': gt,
            'auc': float(roc_auc_score(gt, scores))
        }

        all_scores.extend(scores.tolist())
        all_ground_truths.extend(gt.tolist())

    overall_used_threshold = float(threshold) if threshold is not None else None

    return _overall_results(model_results, all_scores, all_ground_truths, overall_used_threshold)


# -----------------------------
# 3) Single global contrastive model (pairwise similarity)
# -----------------------------

def threshold_evaluation_via_global_contrastive_model(
    dataset: Dataset,
    model_graph: dict,
    threshold: Optional[float] = None,
    ws_streaming_enroll_chunk_size: int = 64,
    ws_streaming_match_chunk_size: int = 4096
) -> Dict[str, Any]:
    """
    Evaluate a single global contrastive model.
    Supports window slicing via chunked E×M batch generation.
    Assumes model takes [left_batch, right_batch] and returns similarity scores.
    """
    enroll_X = dataset.get_enrollment_data()
    enroll_y = dataset.get_enrollment_labels()
    match_X = dataset.get_matching_data()
    match_y = dataset.get_matching_labels()

    _ensure_nonempty(enroll_X, "contrastive enrollment data")
    _ensure_nonempty(match_X, "contrastive matching data")

    gen = _contrastive_batch_generator(enroll_X, enroll_y, match_X, match_y,
                                       ws_streaming_enroll_chunk_size,
                                       ws_streaming_match_chunk_size)

    all_scores: List[float] = []
    all_gt: List[int] = []

    # unpack global model
    model = model_graph[None].model

    for left_batch, right_batch, labels_batch in gen:
        if hasattr(model, "predict"):
            pred = model.predict([left_batch, right_batch], verbose=0).ravel()
        else:
            raise AttributeError("Contrastive model does not expose a recognizable predict interface.")
        all_scores.extend(pred.tolist())
        all_gt.extend(labels_batch.tolist())

    scores_np = np.asarray(all_scores)
    gt_np = np.asarray(all_gt).astype(int)

    metrics = _compute_threshold_and_metrics(scores_np, gt_np, threshold)

    overall_used_threshold = float(threshold) if threshold is not None else None
    return _overall_results(metrics, all_scores, all_gt, overall_used_threshold)


# -----------------------------
# 4) Single global triplet model (embedding-based scoring)
# -----------------------------

def _embed_any(model: dict, X: np.ndarray) -> np.ndarray:
    """
    Get embeddings for X using common patterns. Returns 2D array (n_samples, dim).
    """
    if isinstance(model, dict):
        model = model[None].model

    if hasattr(model, "embed"):
        E = model.embed(X)
        return np.asarray(E)
    if hasattr(model, "predict"):
        E = model.predict(X)
        return np.asarray(E)
    raise AttributeError("Triplet model does not expose an embedding/predict interface.")


def _pairwise_cosine_blocks(A: np.ndarray, B: np.ndarray, block_a: int = 512, block_b: int = 4096) -> Iterable[np.ndarray]:
    """
    Yields flattened cosine similarity blocks for all pairs in A×B.
    """
    # Normalize once to unit vectors
    A_norm = A / (np.linalg.norm(A, axis=1, keepdims=True) + 1e-12)
    B_norm = B / (np.linalg.norm(B, axis=1, keepdims=True) + 1e-12)

    nA = A_norm.shape[0]
    nB = B_norm.shape[0]

    a0 = 0
    while a0 < nA:
        a1 = min(a0 + block_a, nA)
        Aa = A_norm[a0:a1]
        b0 = 0
        while b0 < nB:
            b1 = min(b0 + block_b, nB)
            Bb = B_norm[b0:b1]
            # cosine similarity is dot product since vectors are normalized
            S = Aa @ Bb.T  # shape (a_block, b_block)
            yield S.ravel()
            b0 = b1
        a0 = a1


def threshold_evaluation_via_global_triplet_model(
    dataset: Dataset,
    model_graph: dict,
    threshold: Optional[float] = None,
    ws_streaming_enroll_chunk_size: int = 2048,
    ws_streaming_match_chunk_size: int = 16384
) -> Dict[str, Any]:
    """
    Evaluate a single global triplet-trained model by treating it as an embedding model.
    We compute embeddings for enrollment and matching data, then use cosine similarity
    for all E×M pairs (chunked to support window slicing / memory constraints).
    """
    enroll_X = dataset.get_enrollment_data()
    enroll_y = dataset.get_enrollment_labels()
    match_X = dataset.get_matching_data()
    match_y = dataset.get_matching_labels()

    _ensure_nonempty(enroll_X, "triplet enrollment data")
    _ensure_nonempty(match_X, "triplet matching data")

    # Compute embeddings (can itself be heavy; done in batch if user wraps model.embed)
    E = _embed_any(model_graph, np.asarray(enroll_X))
    M = _embed_any(model_graph, np.asarray(match_X))

    # Build labels for pairwise comparison in a streaming-friendly way
    y_e = np.asarray(enroll_y)
    y_m = np.asarray(match_y)

    # Stream pairwise cosine similarities to avoid holding full E×M matrix
    all_scores: List[float] = []
    all_gt: List[int] = []

    n_e = len(y_e)
    n_m = len(y_m)

    # Precompute class-equality label blocks in sync with score blocks
    # We reuse the same block sizes used for scoring to keep alignment.
    block_e = max(1, int(ws_streaming_enroll_chunk_size))
    block_m = max(1, int(ws_streaming_match_chunk_size))

    e0 = 0
    for S_block in _pairwise_cosine_blocks(E, M, block_a=block_e, block_b=block_m):
        # Determine corresponding label block coordinates
        e1 = min(e0 + block_e, n_e)
        m0 = 0
        while m0 < n_m:
            m1 = min(m0 + block_m, n_m)

            lbl_block = (y_e[e0:e1][:, None] == y_m[m0:m1][None, :]).astype(np.int32).ravel()

            # Extract matching slice length from the current S_block segment.
            # Since _pairwise_cosine_blocks yields per (a_block,b_block),
            # we consume S_block exactly once here and reset m0/m1 accordingly.
            # The generator above yields in the same nested loop order, so we can align.
            # Append and break to advance e0/m0 windows coherently.
            all_scores.extend(S_block[:lbl_block.size].tolist())
            all_gt.extend(lbl_block.tolist())

            m0 = m1
        e0 = e1

    scores_np = np.asarray(all_scores)
    gt_np = np.asarray(all_gt).astype(int)

    metrics = _compute_threshold_and_metrics(scores_np, gt_np, threshold)

    overall_used_threshold = float(threshold) if threshold is not None else None
    return _overall_results(metrics, all_scores, all_gt, overall_used_threshold)
