"""Generate plots for Chapter 5.

Reads the CSVs produced by src/meta/train_eval.py and src/meta/logo_eval.py
and writes a set of figures to outputs/figures/. Each figure is a single
visual claim that can be dropped into Chapter 5 with a caption.
"""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.calibration import calibration_curve
from sklearn.metrics import roc_curve

sns.set_theme(context="paper", style="whitegrid", palette="colorblind")
plt.rcParams["figure.dpi"] = 120
plt.rcParams["savefig.dpi"] = 200
plt.rcParams["font.family"] = "DejaVu Sans"

SYSTEM_ORDER = [
    "binoculars_solo",
    "fast_detect_gpt_solo",
    "roberta_solo",
    "meta_logistic",
    "meta_xgboost",
    "meta_mlp",
]
SYSTEM_LABELS = {
    "binoculars_solo": "Binoculars",
    "fast_detect_gpt_solo": "Fast-DetectGPT",
    "roberta_solo": "RoBERTa",
    "meta_logistic": "Meta (logistic)",
    "meta_xgboost": "Meta (XGBoost)",
    "meta_mlp": "Meta (MLP)",
}


def plot_roc_curves(predictions: pd.DataFrame, outdir: Path) -> None:
    fig, ax = plt.subplots(figsize=(6, 5))
    for system in SYSTEM_ORDER:
        if system not in predictions.columns:
            continue
        fpr, tpr, _ = roc_curve(predictions["label"], predictions[system])
        ax.plot(fpr, tpr, label=SYSTEM_LABELS[system], linewidth=1.8)
    ax.plot([0, 1], [0, 1], "k--", alpha=0.4, linewidth=0.8, label="Random")
    ax.set_xlabel("False positive rate")
    ax.set_ylabel("True positive rate")
    ax.set_title("ROC curves on RAID test set (n = {:,})".format(len(predictions)))
    ax.legend(loc="lower right", frameon=True)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1.02)
    fig.tight_layout()
    fig.savefig(outdir / "fig_roc_curves.png")
    plt.close(fig)
    print("  fig_roc_curves.png")


def plot_fpr_at_tpr95(summary: pd.DataFrame, outdir: Path) -> None:
    df = summary.copy()
    df["display"] = df["system"].map(SYSTEM_LABELS)
    df = df.sort_values("fpr_at_tpr_95", ascending=False)

    fig, ax = plt.subplots(figsize=(7, 4))
    colours = ["#d62728" if "solo" in s else "#1f77b4" for s in df["system"]]
    bars = ax.barh(df["display"], df["fpr_at_tpr_95"], color=colours, edgecolor="black", linewidth=0.5)
    ax.set_xlabel("False positive rate at TPR = 0.95")
    ax.set_title("False positive cost of high-recall detection (lower is better)")
    ax.set_xlim(0, 1.05)
    for bar, value in zip(bars, df["fpr_at_tpr_95"]):
        ax.text(value + 0.01, bar.get_y() + bar.get_height() / 2,
                f"{value:.3f}", va="center", fontsize=9)
    fig.tight_layout()
    fig.savefig(outdir / "fig_fpr_at_tpr95.png")
    plt.close(fig)
    print("  fig_fpr_at_tpr95.png")


def plot_score_distributions(predictions: pd.DataFrame, outdir: Path) -> None:
    base = ["binoculars_solo", "fast_detect_gpt_solo", "roberta_solo"]
    fig, axes = plt.subplots(1, 3, figsize=(12, 3.5), sharey=False)
    for ax, system in zip(axes, base):
        for label, sub in predictions.groupby("label"):
            sns.histplot(sub[system], bins=40, ax=ax, alpha=0.55,
                         label="Human" if label == 0 else "AI",
                         element="step", stat="density", common_norm=False)
        ax.set_title(SYSTEM_LABELS[system])
        ax.set_xlabel("Score (normalised to [0, 1])")
        ax.set_ylabel("Density")
        ax.legend(loc="upper center")
    fig.suptitle("Base detector score distributions on test set")
    fig.tight_layout()
    fig.savefig(outdir / "fig_score_distributions.png")
    plt.close(fig)
    print("  fig_score_distributions.png")


def plot_calibration(predictions: pd.DataFrame, outdir: Path) -> None:
    fig, ax = plt.subplots(figsize=(6, 5))
    for system in SYSTEM_ORDER:
        if system not in predictions.columns:
            continue
        try:
            prob_true, prob_pred = calibration_curve(
                predictions["label"], predictions[system], n_bins=15, strategy="uniform"
            )
            ax.plot(prob_pred, prob_true, marker="o", linewidth=1.5,
                    label=SYSTEM_LABELS[system])
        except Exception as exc:  # pragma: no cover
            print(f"    skip {system}: {exc}")
    ax.plot([0, 1], [0, 1], "k--", alpha=0.4, label="Perfect calibration")
    ax.set_xlabel("Predicted probability of AI")
    ax.set_ylabel("Empirical frequency of AI")
    ax.set_title("Reliability diagram (15-bin)")
    ax.legend(loc="upper left", fontsize=8)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    fig.tight_layout()
    fig.savefig(outdir / "fig_calibration.png")
    plt.close(fig)
    print("  fig_calibration.png")


def plot_logo_results(logo: pd.DataFrame, outdir: Path) -> None:
    df = logo.sort_values("auroc")
    x = np.arange(len(df))
    width = 0.4

    fig, ax1 = plt.subplots(figsize=(8, 4.5))
    ax2 = ax1.twinx()

    bars1 = ax1.bar(x - width / 2, df["auroc"], width, color="#1f77b4",
                    label="AUROC", edgecolor="black", linewidth=0.4)
    bars2 = ax2.bar(x + width / 2, df["fpr_at_tpr_95"], width, color="#d62728",
                    label="FPR@TPR=0.95", edgecolor="black", linewidth=0.4)

    ax1.set_xticks(x)
    ax1.set_xticklabels(df["held_out_generator"], rotation=35, ha="right")
    ax1.set_ylim(0.5, 1.02)
    ax2.set_ylim(0, 1.05)
    ax1.set_ylabel("AUROC", color="#1f77b4")
    ax2.set_ylabel("FPR at TPR = 0.95", color="#d62728")
    ax1.tick_params(axis="y", labelcolor="#1f77b4")
    ax2.tick_params(axis="y", labelcolor="#d62728")
    ax1.set_title("Leave-one-generator-out: meta-classifier (XGBoost) on unseen generator")
    ax1.legend(loc="upper left")
    ax2.legend(loc="upper right")
    fig.tight_layout()
    fig.savefig(outdir / "fig_logo_per_generator.png")
    plt.close(fig)
    print("  fig_logo_per_generator.png")


def plot_per_domain_heatmap(per_domain: pd.DataFrame, outdir: Path) -> None:
    pivot = per_domain.pivot(index="system", columns="domain", values="auroc")
    pivot = pivot.reindex([s for s in SYSTEM_ORDER if s in pivot.index])
    pivot.index = [SYSTEM_LABELS[s] for s in pivot.index]

    fig, ax = plt.subplots(figsize=(8, 4.5))
    sns.heatmap(pivot, annot=True, fmt=".3f", cmap="RdYlGn", vmin=0.5, vmax=1.0,
                cbar_kws={"label": "AUROC"}, ax=ax, linewidths=0.4, linecolor="white")
    ax.set_title("AUROC per (system, domain)")
    ax.set_xlabel("Domain")
    ax.set_ylabel("System")
    plt.setp(ax.get_xticklabels(), rotation=30, ha="right")
    fig.tight_layout()
    fig.savefig(outdir / "fig_per_domain_heatmap.png")
    plt.close(fig)
    print("  fig_per_domain_heatmap.png")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--outdir", type=str, default="outputs")
    args = parser.parse_args()

    outdir = Path(args.outdir)
    figdir = outdir / "figures"
    figdir.mkdir(parents=True, exist_ok=True)

    summary = pd.read_csv(outdir / "metrics_summary.csv")
    predictions = pd.read_csv(outdir / "predictions_test.csv")
    per_domain = pd.read_csv(outdir / "per_domain.csv")
    logo = pd.read_csv(outdir / "logo_results.csv")

    print("Generating figures:")
    plot_roc_curves(predictions, figdir)
    plot_fpr_at_tpr95(summary, figdir)
    plot_score_distributions(predictions, figdir)
    plot_calibration(predictions, figdir)
    plot_logo_results(logo, figdir)
    plot_per_domain_heatmap(per_domain, figdir)
    print(f"\nAll figures saved to {figdir}")


if __name__ == "__main__":
    main()