# Modified plotting utilities with hard dictionary access and no soft-checks.
# This defines the functions in the environment, but does not execute any plotting.
# If keys are missing or arrays are empty when you call them, Python will raise the usual errors (KeyError/ValueError).

import os
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score


def _sanitize_tag(tag: str) -> str:
    """Sanitize a tag for file names."""
    allowed = "-_.() abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    return "".join(c if c in allowed else "_" for c in str(tag)).strip().replace(" ", "_")


def _save_current_figure(save_base_path: str, tag: str, subdir: str | None = None) -> None:
    """Save the current matplotlib figure using a flexible base path.

    Behavior:
    - If `save_base_path` is a directory (exists or ends with os.sep), saves to
      `<dir>/<tag>.png` and `<dir>/<tag>.pdf`.
    - If `save_base_path` points to a file with an extension (e.g., `/tmp/out.png`),
      inserts `_<tag>` before the extension (e.g., `/tmp/out_<tag>.png`) and also
      writes a `.pdf` sibling.
    - Otherwise treats it as a base prefix and writes `<base>_<tag>.png` and `.pdf`.
    """
    tag = _sanitize_tag(tag)

    # Determine how to construct target paths
    is_dir_hint = save_base_path.endswith(os.sep) or os.path.isdir(save_base_path)
    base_dir, base_name = (save_base_path, "") if is_dir_hint else os.path.split(save_base_path)

    if is_dir_hint:
        stem = tag
    else:
        name, ext = os.path.splitext(base_name)
        if ext.lower() in {".png", ".pdf", ".svg", ".jpg", ".jpeg"}:
            stem = f"{name}_{tag}"
        else:
            # treat as prefix without a real extension
            stem = f"{base_name}_{tag}" if base_name else tag

    # Ensure directory exists if provided
    out_dir = base_dir
    if subdir:
        out_dir = os.path.join(out_dir if out_dir else ".", _sanitize_tag(subdir))
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)

    # Build final paths and save
    def _candidate_paths(sfx: str | None = None):
        s = f"{stem}{'' if sfx is None else '_' + str(sfx)}"
        if out_dir:
            return os.path.join(out_dir, f"{s}.png"), os.path.join(out_dir, f"{s}.pdf")
        return f"{s}.png", f"{s}.pdf"

    png_path, pdf_path = _candidate_paths()
    if os.path.exists(png_path) or os.path.exists(pdf_path):
        i = 1
        while True:
            png_try, pdf_try = _candidate_paths(i)
            if not os.path.exists(png_try) and not os.path.exists(pdf_try):
                png_path, pdf_path = png_try, pdf_try
                break
            i += 1

    plt.tight_layout()
    plt.savefig(png_path, bbox_inches="tight")
    plt.savefig(pdf_path, bbox_inches="tight")


def plot_roc(metrics: dict, per_model: bool = False, save_base_path: str = None) -> None:
    metrics = metrics["threshold_evaluation"]
    roc_curve = metrics["roc_curve"]
    overall_fpr = roc_curve["fpr"]
    overall_tpr = roc_curve["tpr"]

    plt.figure()

    # Overall ROC
    plt.plot(overall_fpr, overall_tpr, label=f"Overall ROC (AUC={metrics['overall_metrics']['auc']:.4f})")

    # Per-model ROC
    if per_model:
        for name, res in metrics["per_model_results"].items():
            fpr = res["roc_curve"]["fpr_curve"]
            tpr = res["roc_curve"]["tpr_curve"]
            auc_val = res.get('auc', None)
            if auc_val is None:
                try:
                    from sklearn.metrics import auc as _auc
                    auc_val = _auc(np.array(fpr), np.array(tpr))
                except Exception:
                    auc_val = float('nan')
            plt.plot(fpr, tpr, label=f"Model {name} (AUC={auc_val:.4f})")

    plt.plot([0, 1], [0, 1], linestyle="--", label="Random Guess")
    plt.xlabel("False Positive Rate (FPR)")
    plt.ylabel("True Positive Rate (TPR)")
    plt.title("ROC Curve")
    plt.legend()
    plt.grid(True)
    if save_base_path is not None:
        _save_current_figure(save_base_path, "roc")
        plt.close()
        # Additionally, save per-model individual ROC curves into subfolder
        if per_model:
            for name, res in metrics["per_model_results"].items():
                fpr = np.array(res["roc_curve"]["fpr_curve"]) if res["roc_curve"]["fpr_curve"] is not None else None
                tpr = np.array(res["roc_curve"]["tpr_curve"]) if res["roc_curve"]["tpr_curve"] is not None else None
                if fpr is None or tpr is None:
                    continue
                plt.figure()
                plt.plot(fpr, tpr, label=f"Model {name}")
                plt.plot([0, 1], [0, 1], linestyle="--", label="Random Guess")
                plt.xlabel("False Positive Rate (FPR)")
                plt.ylabel("True Positive Rate (TPR)")
                plt.title(f"ROC Curve (Model {name})")
                plt.legend(); plt.grid(True)
                _save_current_figure(save_base_path, _sanitize_tag(f"roc_model_{name}"), subdir="per_model")
                plt.close()
    else:
        plt.show()


def plot_far_frr(metrics: dict, per_model: bool = False, save_base_path: str = None) -> None:
    metrics = metrics["threshold_evaluation"]

    def _plot_single(fpr_arr, tpr_arr, thr_arr, title_suffix):
        fnr_arr = 1.0 - tpr_arr
        plt.figure()
        plt.plot(thr_arr, fpr_arr, label="FAR (FPR)")
        plt.plot(thr_arr, fnr_arr, label="FRR (FNR)")
        idx = int(np.nanargmin(np.abs(fpr_arr - fnr_arr)))
        eer = (fpr_arr[idx] + fnr_arr[idx]) / 2.0
        plt.scatter([thr_arr[idx]], [eer], marker="o", label=f"EER={eer:.4f}")
        plt.xlabel("Threshold")
        plt.ylabel("Error rate")
        plt.title(f"FAR/FRR vs Threshold {title_suffix}")
        plt.legend()
        plt.grid(True)
        if save_base_path is not None:
            # Title suffix may carry parentheses; keep tag readable
            tag = _sanitize_tag(f"far_frr_{title_suffix.strip('()').replace(' ', '_')}")
            # Place per-model into subfolder automatically
            sub = "per_model" if "Model" in title_suffix else None
            _save_current_figure(save_base_path, tag, subdir=sub)
            plt.close()
        else:
            plt.show()

    # Overall
    overall = metrics["roc_curve"]
    fpr_o, tpr_o, thr_o = np.array(overall["fpr"]), np.array(overall["tpr"]), np.array(overall["thresholds"])
    _plot_single(fpr_o, tpr_o, thr_o, "(Overall)")

    # Per-model
    if per_model:
        for name, res in metrics["per_model_results"].items():
            fpr = np.array(res["roc_curve"]["fpr_curve"]) if res["roc_curve"]["fpr_curve"] is not None else None
            tpr = np.array(res["roc_curve"]["tpr_curve"]) if res["roc_curve"]["tpr_curve"] is not None else None
            thr = np.array(res["roc_curve"]["thresholds"]) if res["roc_curve"]["thresholds"] is not None else None
            if fpr is None or tpr is None or thr is None:
                continue
            _plot_single(fpr, tpr, thr, f"(Model {name})")


def plot_auc_bars(metrics: dict, save_base_path: str = None) -> None:
    labels, aucs = [], []

    if metrics["meta"]['num_models'] > 1:
        # Collect per-model AUCs
        pairs = []
        for name, res in metrics["threshold_evaluation"]["per_model_results"].items():
            fpr, tpr = np.array(res['roc_curve']["fpr_curve"]), np.array(res['roc_curve']["tpr_curve"])
            pairs.append((f"Model {name}", auc(fpr, tpr)))
        # Human/natural sort of labels like P1, P2, ..., P10
        import re
        def _natural_key(s: str):
            parts = re.findall(r"\d+|\D+", str(s))
            return [int(p) if p.isdigit() else p.lower() for p in parts]
        pairs.sort(key=lambda x: _natural_key(x[0]))
        labels = [p[0] for p in pairs]
        aucs = [p[1] for p in pairs]
        overall_auc = metrics["threshold_evaluation"]["overall_metrics"]["auc"]
    elif metrics["meta"]['num_models'] == 1:
        labels.append("Model")
        aucs.append(metrics["threshold_evaluation"]["roc_curve"]["auc"])
        overall_auc = metrics["threshold_evaluation"]["roc_curve"]["auc"]
    else:
        raise ValueError("No `num_models` key found in metrics.")


    labels.append("Overall")
    aucs.append(overall_auc)

    plt.figure()
    x = np.arange(len(labels))
    plt.bar(x, aucs)
    plt.xticks(x, labels, rotation=45, ha="right")
    plt.ylabel("AUC")
    plt.title("AUC per Model and Overall")
    plt.tight_layout()
    if save_base_path is not None:
        _save_current_figure(save_base_path, "auc_bars")
        plt.close()
    else:
        plt.show()


def plot_precision_recall(metrics: dict, per_model: bool = False, save_base_path: str = None) -> None:
    scores_all, labels_all = [], []

    if metrics["meta"]["num_models"] > 1:
        for res in metrics["threshold_evaluation"]["per_model_results"].values():
            scores_all.extend(res["scores"])
            labels_all.extend(res["ground_truth"])
    elif metrics["meta"]["num_models"] == 1:
        scores_all = metrics["threshold_evaluation"]["per_model_results"]["scores"]
        labels_all = metrics["threshold_evaluation"]["per_model_results"]["ground_truth"]

    precision, recall, _ = precision_recall_curve(np.array(labels_all), np.array(scores_all))
    ap = average_precision_score(np.array(labels_all), np.array(scores_all))
    plt.figure()
    plt.plot(recall, precision, label=f"Overall PR (AP={ap:.4f})")
    plt.xlabel("Recall")
    plt.ylabel("Precision")
    plt.title("Precision-Recall Curve")
    plt.legend()
    plt.grid(True)
    if save_base_path is not None:
        _save_current_figure(save_base_path, "pr_overall")
        plt.close()
    else:
        plt.show()

    if per_model:
        for name, res in metrics["per_model_results"].items():
            s, y = np.array(res["scores"]), np.array(res["ground_truth"])
            precision, recall, _ = precision_recall_curve(y, s)
            ap = average_precision_score(y, s)
            plt.figure()
            plt.plot(recall, precision, label=f"Model {name} (AP={ap:.4f})")
            plt.xlabel("Recall")
            plt.ylabel("Precision")
            plt.title(f"Precision-Recall Curve (Model {name})")
            plt.legend()
            plt.grid(True)
            if save_base_path is not None:
                _save_current_figure(save_base_path, _sanitize_tag(f"pr_model_{name}"), subdir="per_model")
                plt.close()
            else:
                plt.show()


def plot_confusion_bars(metrics: dict, save_base_path: str = None) -> None:
    metrics = metrics["threshold_evaluation"]
    om = metrics["overall_metrics"]
    tp, fp, tn, fn = om["tp"], om["fp"], om["tn"], om["fn"]
    plt.figure()
    cats = ["TP", "FP", "TN", "FN"]
    vals = [tp, fp, tn, fn]
    x = np.arange(len(cats))
    plt.bar(x, vals)
    plt.xticks(x, cats)
    plt.ylabel("Count")
    plt.title("Overall Confusion Matrix Counts")
    plt.grid(True, axis="y")
    if save_base_path is not None:
        _save_current_figure(save_base_path, "confusion_bars")
        plt.close()
    else:
        plt.show()


def plot_score_histograms(metrics: dict, bins: int = 50, save_base_path: str = None) -> None:
    scores_all, labels_all = [], []

    if metrics["meta"]["num_models"] > 1:
        for res in metrics["threshold_evaluation"]["per_model_results"].values():
            scores_all.extend(res["scores"])
            labels_all.extend(res["ground_truth"])
    elif metrics["meta"]["num_models"] == 1:
        scores_all = metrics["threshold_evaluation"]["per_model_results"]["scores"]
        labels_all = metrics["threshold_evaluation"]["per_model_results"]["ground_truth"]

    scores_all = np.array(scores_all)
    labels_all = np.array(labels_all)
    pos = scores_all[labels_all == 1]
    neg = scores_all[labels_all == 0]

    plt.figure()
    plt.hist(pos, bins=bins, alpha=0.5, label="Positive (1)")
    plt.hist(neg, bins=bins, alpha=0.5, label="Negative (0)")
    plt.xlabel("Score")
    plt.ylabel("Frequency")
    plt.title("Score Distributions (Aggregated)")
    plt.legend()
    plt.grid(True)
    if save_base_path is not None:
        _save_current_figure(save_base_path, "score_histograms")
        plt.close()
    else:
        plt.show()


def plot_confusion_matrix(metrics: dict, normalize: bool = False, labels=None, cmap="Blues", save_base_path: str = None) -> None:
    # Access the confusion matrix from accuracy_evaluation as requested
    cm = np.array(metrics["accuracy_evaluation"]["confusion_matrix"])  # expected shape (n_classes, n_classes)

    # Determine labels if not provided; prefer labels used to compute cm
    n = cm.shape[0]
    if labels is None:
        labels = metrics.get("accuracy_evaluation", {}).get("confusion_matrix_labels", None)
        if labels is not None:
            labels = list(labels)
        else:
            if n == 2:
                labels = ["Negative", "Positive"]
            else:
                labels = [str(i) for i in range(n)]

    # Apply human/natural sort to labels for display and reorder matrix accordingly
    # Example: P1, P2, ..., P10 instead of P1, P10, P11, ..., P2
    import re

    def _natural_key(s: str):
        parts = re.findall(r"\d+|\D+", str(s))
        return [int(p) if p.isdigit() else p.lower() for p in parts]

    orig_labels = list(labels)
    sorted_labels = sorted(orig_labels, key=_natural_key)
    # Build permutation from original order to natural-sorted order
    idx_map = {val: i for i, val in enumerate(orig_labels)}
    perm = [idx_map[val] for val in sorted_labels]
    # Reorder matrix to match the sorted labels along both axes
    cm = cm[np.ix_(perm, perm)]

    # Optionally normalize rows to sum to 1
    cm_plot = cm.astype(float)
    if normalize:
        row_sums = cm_plot.sum(axis=1, keepdims=True)
        # Avoid division by zero; consistent with module style, we don't mask errors silently.
        cm_plot = np.divide(cm_plot, row_sums, where=row_sums != 0)

    plt.figure()
    im = plt.imshow(cm_plot, interpolation="nearest", cmap=cmap)
    plt.title("Confusion Matrix" + (" (normalized)" if normalize else ""))
    plt.colorbar(im)

    tick_marks = np.arange(len(sorted_labels))
    # Cast labels to string for display to avoid type issues while preserving order
    disp_labels = [str(l) for l in sorted_labels]
    plt.xticks(tick_marks, disp_labels, rotation=90)
    plt.yticks(tick_marks, disp_labels)

    # Annotate cells
    thresh = cm_plot.max() / 1.5 if normalize else cm_plot.max() / 2.0 if cm_plot.max() != 0 else 0
    for i in range(n):
        for j in range(n):
            txt = f"{cm_plot[i, j]:.4f}" if normalize else f"{int(cm[i, j])}"
            plt.text(j, i, txt, ha="center", va="center",
                     color="white" if cm_plot[i, j] > thresh else "black")

    plt.ylabel("True label")
    plt.xlabel("Predicted label")
    plt.tight_layout()
    if save_base_path is not None:
        tag = "confusion_matrix_normalized" if normalize else "confusion_matrix"
        _save_current_figure(save_base_path, tag)
        plt.close()
    else:
        plt.show()
