from __future__ import annotations

import os
from pathlib import Path

MPLCONFIGDIR = Path(__file__).resolve().parents[2] / "outputs" / "logs" / "matplotlib"
MPLCONFIGDIR.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.metrics import ConfusionMatrixDisplay, PrecisionRecallDisplay, RocCurveDisplay, confusion_matrix

from evaluation.metrics import ranking_recall_curve
from utils.project import FIGURES_ROOT, ensure_dir


TASK_LABELS = {
    "a41_clean_binary": "A4.1 clean cell",
    "a42_strict_binary": "A4.2 strict morphology",
    "a43_full_realistic_binary": "A4.3 full realistic",
}
TASK_ORDER = list(TASK_LABELS.keys())
FEATURE_GROUP_ORDER = [
    "morphology",
    "color",
    "texture",
    "morphology_color",
    "morphology_texture",
    "color_texture",
    "full",
]
MODEL_ORDER = [
    "logistic_regression",
    "linear_svm",
    "rbf_svm",
    "random_forest",
    "xgboost",
]
MASK_LABELS = {
    "bbox_only": "bbox only",
    "mask_only": "mask only",
    "bbox_plus_mask": "bbox + mask",
}
MASK_ORDER = list(MASK_LABELS.keys())
REGIME_LABELS = {
    "combined_with_class_weight": "combined + class weight",
    "combined_without_class_weight": "combined without class weight",
    "combined_top30_features": "combined top-30 features",
}
REGIME_ORDER = list(REGIME_LABELS.keys())
AGGREGATION_LABELS = {
    "max_object_score": "max object score",
    "mean_top3_scores": "mean top-3 scores",
    "count_objects_above_threshold": "count > threshold",
    "soft_noisy_or": "soft noisy-or",
}
AGGREGATION_ORDER = list(AGGREGATION_LABELS.keys())
METRIC_LABELS = {
    "pr_auc": "PR-AUC",
    "roc_auc": "ROC-AUC",
    "balanced_accuracy": "Balanced accuracy",
    "recall_blast": "Blast recall",
    "specificity": "Specificity",
    "matthews_corrcoef": "MCC",
    "brier_score": "Brier score",
    "expected_calibration_error": "ECE",
    "recall_at_fpr_10pct": "Recall at FPR 10%",
    "top_10pct_recall": "Recall in top-10%",
    "top_20pct_recall": "Recall in top-20%",
}


def _save(fig: plt.Figure, path: Path) -> str:
    ensure_dir(path.parent)
    fig.tight_layout()
    fig.savefig(path, dpi=180, bbox_inches="tight")
    plt.close(fig)
    return str(path)


def _figure_path(hypothesis: str, date_tag: str, index: int, slug: str) -> Path:
    ensure_dir(FIGURES_ROOT)
    return FIGURES_ROOT / f"{hypothesis}_{date_tag}_{index:02d}_{slug}.png"


def _labeled_categories(values: list[str], mapping: dict[str, str]) -> list[str]:
    return [mapping.get(value, value) for value in values]


def _prepare_metric_frame(
    df: pd.DataFrame,
    category_col: str,
    value_col: str,
    order: list[str] | None = None,
) -> pd.DataFrame:
    frame = df[[category_col, value_col]].copy()
    frame[value_col] = frame[value_col].astype(float)
    if order is not None:
        frame[category_col] = pd.Categorical(frame[category_col], categories=order, ordered=True)
        frame = frame.sort_values(category_col)
    return frame


def _save_simple_bar(
    df: pd.DataFrame,
    category_col: str,
    value_col: str,
    path: Path,
    title: str,
    y_label: str,
    order: list[str] | None = None,
    label_map: dict[str, str] | None = None,
    color: str = "#4c72b0",
) -> str:
    frame = _prepare_metric_frame(df, category_col, value_col, order=order)
    categories = frame[category_col].astype(str).tolist()
    labels = _labeled_categories(categories, label_map or {})
    values = frame[value_col].astype(float).to_numpy()

    fig, ax = plt.subplots(figsize=(10, 6))
    bars = ax.bar(labels, values, color=color)
    ax.set_title(title)
    ax.set_ylabel(y_label)
    ax.tick_params(axis="x", rotation=20)
    for bar, value in zip(bars, values, strict=False):
        vertical_align = "bottom" if value >= 0 else "top"
        ax.text(bar.get_x() + bar.get_width() / 2.0, value, f"{value:.3f}", ha="center", va=vertical_align, fontsize=8)
    return _save(fig, path)


def _save_grouped_bar(
    df: pd.DataFrame,
    category_col: str,
    group_col: str,
    value_col: str,
    path: Path,
    title: str,
    y_label: str,
    category_order: list[str] | None = None,
    group_order: list[str] | None = None,
    category_labels: dict[str, str] | None = None,
    group_labels: dict[str, str] | None = None,
) -> str:
    frame = df[[category_col, group_col, value_col]].copy()
    frame[value_col] = frame[value_col].astype(float)
    pivot = frame.pivot(index=category_col, columns=group_col, values=value_col)
    if category_order is not None:
        pivot = pivot.reindex(category_order)
    if group_order is not None:
        pivot = pivot.reindex(columns=group_order)
    categories = [value for value in pivot.index.tolist() if value is not None]
    groups = [value for value in pivot.columns.tolist() if value is not None]
    x_positions = np.arange(len(categories))
    width = 0.8 / max(len(groups), 1)

    fig, ax = plt.subplots(figsize=(11, 6))
    for idx, group in enumerate(groups):
        offsets = x_positions + (idx - (len(groups) - 1) / 2.0) * width
        values = pivot[group].astype(float).to_numpy()
        bars = ax.bar(offsets, values, width=width, label=(group_labels or {}).get(group, group))
        for bar, value in zip(bars, values, strict=False):
            vertical_align = "bottom" if value >= 0 else "top"
            ax.text(bar.get_x() + bar.get_width() / 2.0, value, f"{value:.3f}", ha="center", va=vertical_align, fontsize=7)
    ax.set_xticks(x_positions)
    ax.set_xticklabels(_labeled_categories(categories, category_labels or {}), rotation=15)
    ax.set_title(title)
    ax.set_ylabel(y_label)
    ax.legend(fontsize=8)
    return _save(fig, path)


def _save_heatmap(
    frame: pd.DataFrame,
    path: Path,
    title: str,
    x_label: str,
    y_label: str,
    value_fmt: str = ".3f",
    cmap: str = "Blues",
) -> str:
    matrix = frame.astype(float).to_numpy()
    fig, ax = plt.subplots(figsize=(11, 6))
    image = ax.imshow(matrix, cmap=cmap, aspect="auto")
    ax.set_xticks(np.arange(frame.shape[1]))
    ax.set_yticks(np.arange(frame.shape[0]))
    ax.set_xticklabels(frame.columns.tolist(), rotation=25, ha="right")
    ax.set_yticklabels(frame.index.tolist())
    ax.set_title(title)
    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)
    for row_idx in range(frame.shape[0]):
        for col_idx in range(frame.shape[1]):
            value = matrix[row_idx, col_idx]
            if np.isnan(value):
                continue
            ax.text(col_idx, row_idx, format(value, value_fmt), ha="center", va="center", fontsize=8)
    fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
    return _save(fig, path)


def _save_metric_heatmap(
    df: pd.DataFrame,
    index_col: str,
    column_col: str,
    value_col: str,
    path: Path,
    title: str,
    index_order: list[str] | None = None,
    column_order: list[str] | None = None,
    index_labels: dict[str, str] | None = None,
    column_labels: dict[str, str] | None = None,
) -> str:
    pivot = df.pivot(index=index_col, columns=column_col, values=value_col)
    if index_order is not None:
        pivot = pivot.reindex(index_order)
    if column_order is not None:
        pivot = pivot.reindex(columns=column_order)
    pivot.index = _labeled_categories([str(value) for value in pivot.index.tolist()], index_labels or {})
    pivot.columns = _labeled_categories([str(value) for value in pivot.columns.tolist()], column_labels or {})
    return _save_heatmap(
        pivot,
        path=path,
        title=title,
        x_label=column_col,
        y_label=index_col,
    )


def _save_pr_or_roc_curves(
    prediction_df: pd.DataFrame,
    run_rows: pd.DataFrame,
    path: Path,
    curve_type: str,
    title: str,
) -> str:
    fig, ax = plt.subplots(figsize=(10, 8))
    for _, row in run_rows.iterrows():
        run_df = prediction_df.loc[prediction_df["run_id"].eq(row["run_id"])].copy()
        label = f"{TASK_LABELS.get(str(row['task_name']), row['task_name'])} | {row['model_name']} | {row['feature_group']}"
        if curve_type == "pr":
            PrecisionRecallDisplay.from_predictions(run_df["y_true"], run_df["y_score"], name=label, ax=ax)
        else:
            RocCurveDisplay.from_predictions(run_df["y_true"], run_df["y_score"], name=label, ax=ax)
    ax.set_title(title)
    return _save(fig, path)


def _save_confusion_from_predictions(prediction_df: pd.DataFrame, path: Path, title: str) -> str:
    fig, ax = plt.subplots(figsize=(5, 5))
    ConfusionMatrixDisplay(
        confusion_matrix(prediction_df["y_true"], prediction_df["y_pred"], labels=[0, 1]),
        display_labels=["non-blast", "blast"],
    ).plot(ax=ax, colorbar=False)
    ax.set_title(title)
    return _save(fig, path)


def _save_metric_pair_bar(
    df: pd.DataFrame,
    category_col: str,
    metric_a: str,
    metric_b: str,
    path: Path,
    title: str,
    category_order: list[str] | None = None,
    category_labels: dict[str, str] | None = None,
) -> str:
    frame = df[[category_col, metric_a, metric_b]].copy()
    if category_order is not None:
        frame[category_col] = pd.Categorical(frame[category_col], categories=category_order, ordered=True)
        frame = frame.sort_values(category_col)
    categories = frame[category_col].astype(str).tolist()
    labels = _labeled_categories(categories, category_labels or {})
    values_a = frame[metric_a].astype(float).to_numpy()
    values_b = frame[metric_b].astype(float).to_numpy()

    x_positions = np.arange(len(labels))
    width = 0.35
    fig, ax = plt.subplots(figsize=(10, 6))
    bars_a = ax.bar(x_positions - width / 2.0, values_a, width=width, label=METRIC_LABELS.get(metric_a, metric_a))
    bars_b = ax.bar(x_positions + width / 2.0, values_b, width=width, label=METRIC_LABELS.get(metric_b, metric_b))
    for bars, values in [(bars_a, values_a), (bars_b, values_b)]:
        for bar, value in zip(bars, values, strict=False):
            vertical_align = "bottom" if value >= 0 else "top"
            ax.text(bar.get_x() + bar.get_width() / 2.0, value, f"{value:.3f}", ha="center", va=vertical_align, fontsize=8)
    ax.set_xticks(x_positions)
    ax.set_xticklabels(labels, rotation=15)
    ax.set_title(title)
    ax.set_ylabel("Score")
    ax.legend(fontsize=8)
    return _save(fig, path)


def _save_feature_barh(
    df: pd.DataFrame,
    value_col: str,
    path: Path,
    title: str,
    x_label: str,
) -> str:
    plot_df = df.head(20).iloc[::-1].copy()
    values = plot_df[value_col].astype(float).to_numpy()
    fig, ax = plt.subplots(figsize=(10, 8))
    ax.barh(plot_df["feature_name"], values, color="#55a868")
    ax.set_title(title)
    ax.set_xlabel(x_label)
    return _save(fig, path)


def _save_delta_vs_reference(
    df: pd.DataFrame,
    category_col: str,
    baseline_label: str,
    compare_labels: list[str],
    value_col: str,
    path: Path,
    title: str,
    category_order: list[str] | None = None,
    category_labels: dict[str, str] | None = None,
    compare_label_map: dict[str, str] | None = None,
) -> str:
    pivot = df.pivot(index=category_col, columns="label", values=value_col)
    if category_order is not None:
        pivot = pivot.reindex(category_order)
    delta_df = pd.DataFrame(index=pivot.index)
    for label in compare_labels:
        delta_df[label] = pivot[label] - pivot[baseline_label]
    delta_df = delta_df.reset_index().melt(id_vars=category_col, var_name="delta_label", value_name="delta_value")
    return _save_grouped_bar(
        delta_df,
        category_col=category_col,
        group_col="delta_label",
        value_col="delta_value",
        path=path,
        title=title,
        y_label=f"Delta {METRIC_LABELS.get(value_col, value_col)}",
        category_order=category_order,
        group_order=compare_labels,
        category_labels=category_labels,
        group_labels=compare_label_map,
    )


def _save_ranking_curve_panel(
    score_df: pd.DataFrame,
    target_column: str,
    path: Path,
    title: str,
) -> str:
    fig, ax = plt.subplots(figsize=(8, 6))
    for method in AGGREGATION_ORDER:
        method_df = score_df.loc[score_df["aggregation_method"].eq(method)].copy()
        if method_df.empty:
            continue
        screened_fraction, cumulative_recall = ranking_recall_curve(
            method_df[target_column].astype(int).to_numpy(),
            method_df["score"].astype(float).to_numpy(),
        )
        ax.plot(screened_fraction, cumulative_recall, label=AGGREGATION_LABELS.get(method, method))
    ax.set_xlabel("Screened fraction")
    ax.set_ylabel("Recall of positive groups")
    ax.set_title(title)
    ax.legend(fontsize=8)
    return _save(fig, path)


def _save_best_score_distributions(
    image_scores_df: pd.DataFrame,
    patient_scores_df: pd.DataFrame,
    image_method: str,
    patient_method: str,
    path: Path,
) -> str:
    image_frame = image_scores_df.loc[image_scores_df["aggregation_method"].eq(image_method)].copy()
    patient_frame = patient_scores_df.loc[patient_scores_df["aggregation_method"].eq(patient_method)].copy()

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    for axis, frame, target_col, method, panel_title in [
        (axes[0], image_frame, "image_contains_blast", image_method, "Best image-level method"),
        (axes[1], patient_frame, "patient_contains_blast", patient_method, "Best patient-level method"),
    ]:
        negatives = frame.loc[frame[target_col].eq(0), "score"].astype(float).to_numpy()
        positives = frame.loc[frame[target_col].eq(1), "score"].astype(float).to_numpy()
        axis.hist(negatives, bins=20, alpha=0.65, label="negative", color="#c44e52")
        axis.hist(positives, bins=20, alpha=0.65, label="positive", color="#4c72b0")
        axis.set_title(f"{panel_title}\n{AGGREGATION_LABELS.get(method, method)}")
        axis.set_xlabel("Aggregated score")
        axis.set_ylabel("Count")
        axis.legend(fontsize=8)
    return _save(fig, path)


def _save_training_curves(history_payload: dict[str, list[dict[str, float]]], path: Path) -> str:
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    for model_name, history in history_payload.items():
        history_df = pd.DataFrame(history)
        if history_df.empty:
            continue
        axes[0].plot(history_df["epoch"], history_df["train_loss"], label=f"{model_name} train")
        axes[0].plot(history_df["epoch"], history_df["val_loss"], linestyle="--", label=f"{model_name} val")
        axes[1].plot(history_df["epoch"], history_df["val_pr_auc"], label=model_name)
    axes[0].set_title("Deep baseline loss")
    axes[0].set_xlabel("Epoch")
    axes[1].set_title("Validation PR-AUC")
    axes[1].set_xlabel("Epoch")
    for axis in axes:
        axis.legend(fontsize=8)
    return _save(fig, path)


def _save_deep_pr_curve_comparison(
    handcrafted_predictions: pd.DataFrame,
    deep_predictions: pd.DataFrame,
    path: Path,
    title: str,
) -> str:
    fig, ax = plt.subplots(figsize=(10, 8))
    PrecisionRecallDisplay.from_predictions(
        handcrafted_predictions["y_true"],
        handcrafted_predictions["y_score"],
        name="handcrafted best",
        ax=ax,
    )
    for model_name, frame in deep_predictions.groupby("model_name", sort=False):
        PrecisionRecallDisplay.from_predictions(frame["y_true"], frame["y_score"], name=str(model_name), ax=ax)
    ax.set_title(title)
    return _save(fig, path)


def generate_h1_figures(
    leaderboard_df: pd.DataFrame,
    predictions_df: pd.DataFrame,
    top_features_df: pd.DataFrame,
    date_tag: str,
) -> list[str]:
    best_rows = (
        leaderboard_df.sort_values(["task_name", "validation_pr_auc", "test_pr_auc"], ascending=[True, False, False])
        .groupby("task_name", sort=False)
        .head(1)
        .reset_index(drop=True)
    )
    top_runs = (
        leaderboard_df.sort_values(["task_name", "validation_pr_auc", "test_pr_auc"], ascending=[True, False, False])
        .groupby("task_name", sort=False)
        .head(2)
        .reset_index(drop=True)
    )
    best_full_row = leaderboard_df.loc[leaderboard_df["task_name"].eq("a43_full_realistic_binary")].sort_values(
        ["validation_pr_auc", "test_pr_auc"],
        ascending=[False, False],
    ).iloc[0]

    figure_paths: list[str] = []
    figure_paths.append(
        _save_pr_or_roc_curves(
            prediction_df=predictions_df,
            run_rows=top_runs,
            path=_figure_path("H1", date_tag, 1, "pr_curves_top_runs"),
            curve_type="pr",
            title="H1: precision-recall curves for top object-level runs",
        )
    )
    figure_paths.append(
        _save_pr_or_roc_curves(
            prediction_df=predictions_df,
            run_rows=top_runs,
            path=_figure_path("H1", date_tag, 2, "roc_curves_top_runs"),
            curve_type="roc",
            title="H1: ROC curves for top object-level runs",
        )
    )

    for idx, task_name in enumerate(TASK_ORDER, start=3):
        task_frame = leaderboard_df.loc[leaderboard_df["task_name"].eq(task_name)].copy()
        pivot = task_frame.pivot(index="model_name", columns="feature_group", values="test_pr_auc")
        pivot = pivot.reindex(index=MODEL_ORDER, columns=FEATURE_GROUP_ORDER)
        pivot.index = [str(value) for value in pivot.index.tolist()]
        pivot.columns = [str(value) for value in pivot.columns.tolist()]
        figure_paths.append(
            _save_heatmap(
                pivot,
                path=_figure_path("H1", date_tag, idx, f"{task_name}_pr_auc_heatmap"),
                title=f"H1: test PR-AUC for {TASK_LABELS.get(task_name, task_name)}",
                x_label="Feature group",
                y_label="Model",
            )
        )

    figure_paths.append(
        _save_simple_bar(
            df=best_rows,
            category_col="task_name",
            value_col="test_pr_auc",
            path=_figure_path("H1", date_tag, 6, "best_pr_auc_by_task"),
            title="H1: best object-level PR-AUC by task",
            y_label="PR-AUC",
            order=TASK_ORDER,
            label_map=TASK_LABELS,
            color="#8172b3",
        )
    )
    figure_paths.append(
        _save_metric_pair_bar(
            df=best_rows,
            category_col="task_name",
            metric_a="test_recall_blast",
            metric_b="test_specificity",
            path=_figure_path("H1", date_tag, 7, "best_recall_specificity_by_task"),
            title="H1: recall and specificity for the best run in each task",
            category_order=TASK_ORDER,
            category_labels=TASK_LABELS,
        )
    )
    best_full_predictions = predictions_df.loc[
        predictions_df["run_id"].eq(best_full_row["run_id"]) & predictions_df["split_name"].eq("test")
    ].copy()
    figure_paths.append(
        _save_confusion_from_predictions(
            prediction_df=best_full_predictions,
            path=_figure_path("H1", date_tag, 8, "best_full_realistic_confusion"),
            title=f"H1: confusion matrix for {best_full_row['model_name']} + {best_full_row['feature_group']}",
        )
    )
    figure_paths.append(
        _save_feature_barh(
            df=top_features_df,
            value_col="top20_frequency",
            path=_figure_path("H1", date_tag, 9, "top_feature_frequency"),
            title="H1: stability of top features",
            x_label="Top-20 frequency",
        )
    )
    figure_paths.append(
        _save_feature_barh(
            df=top_features_df.sort_values(["mean_importance", "top20_frequency"], ascending=[False, False]),
            value_col="mean_importance",
            path=_figure_path("H1", date_tag, 10, "top_feature_importance"),
            title="H1: mean permutation importance of top features",
            x_label="Mean importance",
        )
    )
    return figure_paths


def generate_h2_figures(mask_ablation_df: pd.DataFrame, date_tag: str) -> list[str]:
    metric_specs = [
        (1, "pr_auc", "H2: PR-AUC for bbox and mask feature variants", "PR-AUC"),
        (2, "roc_auc", "H2: ROC-AUC for bbox and mask feature variants", "ROC-AUC"),
        (3, "balanced_accuracy", "H2: balanced accuracy for bbox and mask feature variants", "Balanced accuracy"),
        (4, "recall_blast", "H2: blast recall for bbox and mask feature variants", "Blast recall"),
        (5, "specificity", "H2: specificity for bbox and mask feature variants", "Specificity"),
        (6, "matthews_corrcoef", "H2: MCC for bbox and mask feature variants", "MCC"),
        (7, "brier_score", "H2: Brier score for bbox and mask feature variants", "Brier score"),
        (8, "expected_calibration_error", "H2: calibration error for bbox and mask feature variants", "ECE"),
    ]
    figure_paths: list[str] = []
    for index, metric, title, y_label in metric_specs:
        figure_paths.append(
            _save_grouped_bar(
                df=mask_ablation_df,
                category_col="task_name",
                group_col="label",
                value_col=metric,
                path=_figure_path("H2", date_tag, index, metric),
                title=title,
                y_label=y_label,
                category_order=TASK_ORDER,
                group_order=MASK_ORDER,
                category_labels=TASK_LABELS,
                group_labels=MASK_LABELS,
            )
        )
    figure_paths.append(
        _save_delta_vs_reference(
            df=mask_ablation_df,
            category_col="task_name",
            baseline_label="bbox_only",
            compare_labels=["mask_only", "bbox_plus_mask"],
            value_col="pr_auc",
            path=_figure_path("H2", date_tag, 9, "pr_auc_delta_vs_bbox"),
            title="H2: PR-AUC gain relative to bbox-only features",
            category_order=TASK_ORDER,
            category_labels=TASK_LABELS,
            compare_label_map={
                "mask_only": "mask only - bbox only",
                "bbox_plus_mask": "bbox + mask - bbox only",
            },
        )
    )
    figure_paths.append(
        _save_metric_heatmap(
            df=mask_ablation_df,
            index_col="task_name",
            column_col="label",
            value_col="pr_auc",
            path=_figure_path("H2", date_tag, 10, "pr_auc_heatmap"),
            title="H2: PR-AUC heatmap for mask ablation",
            index_order=TASK_ORDER,
            column_order=MASK_ORDER,
            index_labels=TASK_LABELS,
            column_labels=MASK_LABELS,
        )
    )
    return figure_paths


def generate_h3_figures(dataset_ablation_df: pd.DataFrame, date_tag: str) -> list[str]:
    primary_df = dataset_ablation_df.loc[dataset_ablation_df["label"].eq("combined_with_class_weight")].copy()
    extra_df = dataset_ablation_df.loc[dataset_ablation_df["task_name"].eq("a43_full_realistic_binary")].copy()
    metric_specs = [
        (1, "pr_auc", "H3: PR-AUC across corpus regimes", "PR-AUC"),
        (2, "roc_auc", "H3: ROC-AUC across corpus regimes", "ROC-AUC"),
        (3, "balanced_accuracy", "H3: balanced accuracy across corpus regimes", "Balanced accuracy"),
        (4, "recall_blast", "H3: blast recall across corpus regimes", "Blast recall"),
        (5, "specificity", "H3: specificity across corpus regimes", "Specificity"),
        (6, "matthews_corrcoef", "H3: MCC across corpus regimes", "MCC"),
        (7, "brier_score", "H3: Brier score across corpus regimes", "Brier score"),
        (8, "expected_calibration_error", "H3: calibration error across corpus regimes", "ECE"),
    ]
    figure_paths: list[str] = []
    for index, metric, title, y_label in metric_specs:
        figure_paths.append(
            _save_simple_bar(
                df=primary_df,
                category_col="task_name",
                value_col=metric,
                path=_figure_path("H3", date_tag, index, metric),
                title=title,
                y_label=y_label,
                order=TASK_ORDER,
                label_map=TASK_LABELS,
                color="#64b5cd",
            )
        )
    heatmap_df = primary_df.set_index("task_name")[
        ["pr_auc", "roc_auc", "balanced_accuracy", "recall_blast", "specificity"]
    ].reindex(TASK_ORDER)
    heatmap_df.index = _labeled_categories([str(value) for value in heatmap_df.index.tolist()], TASK_LABELS)
    heatmap_df.columns = [METRIC_LABELS.get(column, column) for column in heatmap_df.columns.tolist()]
    figure_paths.append(
        _save_heatmap(
            heatmap_df,
            path=_figure_path("H3", date_tag, 9, "primary_metrics_heatmap"),
            title="H3: metric heatmap for the main corpus regimes",
            x_label="Metric",
            y_label="Task",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=extra_df,
            category_col="label",
            value_col="pr_auc",
            path=_figure_path("H3", date_tag, 10, "full_realistic_variants_pr_auc"),
            title="H3: full-realistic ablation variants",
            y_label="PR-AUC",
            order=REGIME_ORDER,
            label_map=REGIME_LABELS,
            color="#dd8452",
        )
    )
    return figure_paths


def generate_h4_figures(
    image_metrics_df: pd.DataFrame,
    patient_metrics_df: pd.DataFrame,
    image_scores_df: pd.DataFrame,
    patient_scores_df: pd.DataFrame,
    date_tag: str,
) -> list[str]:
    best_image_method = (
        image_metrics_df.sort_values(["pr_auc", "roc_auc"], ascending=[False, False]).iloc[0]["aggregation_method"]
    )
    best_patient_method = (
        patient_metrics_df.sort_values(["pr_auc", "roc_auc"], ascending=[False, False]).iloc[0]["aggregation_method"]
    )
    figure_paths: list[str] = []
    figure_paths.append(
        _save_simple_bar(
            df=image_metrics_df,
            category_col="aggregation_method",
            value_col="pr_auc",
            path=_figure_path("H4", date_tag, 1, "image_pr_auc"),
            title="H4: image-level PR-AUC by aggregation method",
            y_label="PR-AUC",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#4c72b0",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=image_metrics_df,
            category_col="aggregation_method",
            value_col="roc_auc",
            path=_figure_path("H4", date_tag, 2, "image_roc_auc"),
            title="H4: image-level ROC-AUC by aggregation method",
            y_label="ROC-AUC",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#55a868",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=image_metrics_df,
            category_col="aggregation_method",
            value_col="balanced_accuracy",
            path=_figure_path("H4", date_tag, 3, "image_balanced_accuracy"),
            title="H4: image-level balanced accuracy by aggregation method",
            y_label="Balanced accuracy",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#c44e52",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=image_metrics_df,
            category_col="aggregation_method",
            value_col="recall_at_fpr_10pct",
            path=_figure_path("H4", date_tag, 4, "image_recall_at_fpr10"),
            title="H4: image-level recall at FPR 10%",
            y_label="Recall at FPR 10%",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#8172b3",
        )
    )
    figure_paths.append(
        _save_ranking_curve_panel(
            score_df=image_scores_df,
            target_column="image_contains_blast",
            path=_figure_path("H4", date_tag, 5, "image_ranking_curve"),
            title="H4: image-level ranking recall curves",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=patient_metrics_df,
            category_col="aggregation_method",
            value_col="pr_auc",
            path=_figure_path("H4", date_tag, 6, "patient_pr_auc"),
            title="H4: patient-level PR-AUC by aggregation method",
            y_label="PR-AUC",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#4c72b0",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=patient_metrics_df,
            category_col="aggregation_method",
            value_col="roc_auc",
            path=_figure_path("H4", date_tag, 7, "patient_roc_auc"),
            title="H4: patient-level ROC-AUC by aggregation method",
            y_label="ROC-AUC",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#55a868",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=patient_metrics_df,
            category_col="aggregation_method",
            value_col="balanced_accuracy",
            path=_figure_path("H4", date_tag, 8, "patient_balanced_accuracy"),
            title="H4: patient-level balanced accuracy by aggregation method",
            y_label="Balanced accuracy",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#c44e52",
        )
    )
    figure_paths.append(
        _save_simple_bar(
            df=patient_metrics_df,
            category_col="aggregation_method",
            value_col="recall_at_fpr_10pct",
            path=_figure_path("H4", date_tag, 9, "patient_recall_at_fpr10"),
            title="H4: patient-level recall at FPR 10%",
            y_label="Recall at FPR 10%",
            order=AGGREGATION_ORDER,
            label_map=AGGREGATION_LABELS,
            color="#8172b3",
        )
    )
    figure_paths.append(
        _save_best_score_distributions(
            image_scores_df=image_scores_df,
            patient_scores_df=patient_scores_df,
            image_method=str(best_image_method),
            patient_method=str(best_patient_method),
            path=_figure_path("H4", date_tag, 10, "best_score_distributions"),
        )
    )
    return figure_paths


def generate_h5_figures(
    comparison_df: pd.DataFrame,
    handcrafted_predictions: pd.DataFrame,
    deep_predictions: pd.DataFrame,
    training_history: dict[str, list[dict[str, float]]],
    date_tag: str,
) -> list[str]:
    metric_specs = [
        (1, "test_pr_auc", "H5: PR-AUC of handcrafted and deep baselines", "PR-AUC"),
        (2, "test_roc_auc", "H5: ROC-AUC of handcrafted and deep baselines", "ROC-AUC"),
        (3, "test_balanced_accuracy", "H5: balanced accuracy of handcrafted and deep baselines", "Balanced accuracy"),
        (4, "test_recall_blast", "H5: blast recall of handcrafted and deep baselines", "Blast recall"),
        (5, "test_specificity", "H5: specificity of handcrafted and deep baselines", "Specificity"),
        (6, "test_matthews_corrcoef", "H5: MCC of handcrafted and deep baselines", "MCC"),
        (7, "test_brier_score", "H5: Brier score of handcrafted and deep baselines", "Brier score"),
        (8, "test_expected_calibration_error", "H5: calibration error of handcrafted and deep baselines", "ECE"),
    ]
    figure_paths: list[str] = []
    model_order = comparison_df["model_name"].astype(str).tolist()
    label_map = {
        "handcrafted_best_full_realistic": "handcrafted best",
        "cnn_bbox": "CNN bbox",
        "cnn_mask": "CNN mask",
        "cnn_late_fusion": "CNN late fusion",
    }
    for index, metric, title, y_label in metric_specs:
        figure_paths.append(
            _save_simple_bar(
                df=comparison_df.rename(columns={metric: "metric_value"}),
                category_col="model_name",
                value_col="metric_value",
                path=_figure_path("H5", date_tag, index, metric.replace("test_", "")),
                title=title,
                y_label=y_label,
                order=model_order,
                label_map=label_map,
                color="#937860",
            )
        )
    figure_paths.append(
        _save_training_curves(
            history_payload=training_history,
            path=_figure_path("H5", date_tag, 9, "training_curves"),
        )
    )
    figure_paths.append(
        _save_deep_pr_curve_comparison(
            handcrafted_predictions=handcrafted_predictions,
            deep_predictions=deep_predictions,
            path=_figure_path("H5", date_tag, 10, "precision_recall_curves"),
            title="H5: precision-recall comparison of handcrafted and deep baselines",
        )
    )
    return figure_paths