"""
Generate publication-ready figures for the AURIS academic paper.
All text in English, 300 DPI, Times New Roman.

Produces:
  - paper_roc_curves.png
  - paper_confusion_matrix_lightgbm.png
  - paper_model_comparison.png
  - paper_feature_importance.png
  - paper_feature_distribution.png
  - paper_ml_vs_dl.png
  - paper_calibration.png
  - paper_precision_recall.png
  - paper_score_distribution.png
  - paper_shap_summary.png  (if shap available)
  - paper_pipeline_diagram.png
  - paper_fold_std_table.png

Usage:
    python -m app.training.generate_paper_figures
"""

from __future__ import annotations

import json
import sys
import io
from pathlib import Path

import numpy as np

try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    import matplotlib.patheffects as pe
    from matplotlib.gridspec import GridSpec
except ImportError:
    print("matplotlib required")
    sys.exit(1)

try:
    import seaborn as sns
    HAS_SEABORN = True
except ImportError:
    HAS_SEABORN = False

from sklearn.metrics import (
    roc_curve, auc, precision_recall_curve,
    confusion_matrix, average_precision_score,
)
from sklearn.calibration import calibration_curve

# ── Paths ──────────────────────────────────────────────────────────────
BASE   = Path(__file__).resolve().parents[2]
MODELS = BASE / "models"
OUT    = BASE.parent / "docs" / "academic" / "figures"
OUT.mkdir(parents=True, exist_ok=True)

# ── Style ──────────────────────────────────────────────────────────────
plt.rcParams.update({
    "font.family":        "serif",
    "font.serif":         ["Times New Roman", "DejaVu Serif"],
    "font.size":          11,
    "axes.titlesize":     13,
    "axes.labelsize":     11,
    "xtick.labelsize":    10,
    "ytick.labelsize":    10,
    "legend.fontsize":    9,
    "figure.dpi":         150,
    "savefig.dpi":        300,
    "savefig.bbox":       "tight",
    "savefig.pad_inches": 0.15,
    "axes.grid":          True,
    "grid.alpha":         0.3,
    "axes.spines.top":    False,
    "axes.spines.right":  False,
})

# Colorblind-safe palette
COLORS = {
    "Logistic Regression":       "#4363d8",
    "Random Forest":             "#3cb44b",
    "Gradient Boosting":         "#e6194b",
    "SVM (RBF)":                 "#f58231",
    "MLP Neural Network":        "#911eb4",
    "XGBoost":                   "#42d4f4",
    "LightGBM":                  "#f032e6",
    "Deep MLP (512-256-128-64)": "#e6194b",
    "1D-CNN":                    "#f58231",
    "Residual MLP (3 blocks)":   "#3cb44b",
    "Attention MLP":             "#4363d8",
}
GOLD = "#C99347"
BG   = "#faf8f4"


def _c(name: str) -> str:
    return COLORS.get(name, "#555555")


def _save(fig, name: str) -> None:
    fig.savefig(OUT / f"{name}.png")
    fig.savefig(OUT / f"{name}.pdf")
    plt.close(fig)
    print(f"  OK  {name}.png")


# ── Load data ──────────────────────────────────────────────────────────

def _load_results():
    with open(MODELS / "training_results.json") as f:
        ml = json.load(f)
    with open(MODELS / "deep_learning_results.json") as f:
        dl = json.load(f)
    # Merge: DL entries don't have y_true/y_prob from CV, so keep separate
    return ml, dl


# ══════════════════════════════════════════════════════════════════════
# 1. ROC Curves — all models
# ══════════════════════════════════════════════════════════════════════

def plot_roc_curves(ml: dict, dl: dict) -> None:
    fig, ax = plt.subplots(figsize=(7, 6))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    # ML models (have y_true / y_prob from CV)
    for name, data in ml.items():
        if name.startswith("_") or not isinstance(data, dict):
            continue
        if "y_true" not in data or "y_prob" not in data:
            continue
        fpr, tpr, _ = roc_curve(data["y_true"], data["y_prob"])
        roc_auc = auc(fpr, tpr)
        lw = 2.5 if name == ml.get("_best_model") else 1.5
        ax.plot(fpr, tpr, color=_c(name), lw=lw,
                label=f"{name} (AUC={roc_auc:.3f})")

    # DL models — use reported AUC value as annotation
    for name, data in dl.items():
        if not isinstance(data, dict):
            continue
        auc_val = data.get("roc_auc", 0)
        ax.annotate(f"{name}: AUC={auc_val:.3f}",
                    xy=(0.55, 0.05 + list(dl.keys()).index(name) * 0.055),
                    xycoords="axes fraction", fontsize=8, color=_c(name))

    ax.plot([0, 1], [0, 1], "k--", lw=1, alpha=0.4, label="Random (AUC=0.500)")
    ax.set_xlabel("False Positive Rate")
    ax.set_ylabel("True Positive Rate")
    ax.set_title("ROC Curves — All Models (5-Fold CV)")
    ax.legend(loc="lower right", fontsize=8)
    ax.set_xlim(0, 1); ax.set_ylim(0, 1.02)
    _save(fig, "paper_roc_curves")


# ══════════════════════════════════════════════════════════════════════
# 2. Confusion Matrix — LightGBM (best ML model)
# ══════════════════════════════════════════════════════════════════════

def plot_confusion_matrix_lightgbm(ml: dict) -> None:
    name = "LightGBM"
    data = ml[name]
    # Load pre-computed CV predictions
    y_true = np.load(MODELS / "lgbm_cv_ytrue.npy")
    y_prob = np.load(MODELS / "lgbm_cv_probs.npy")
    threshold = data.get("optimal_threshold", 0.5)
    y_pred = (y_prob >= threshold).astype(int)

    cm = confusion_matrix(y_true, y_pred)
    acc = (cm[0, 0] + cm[1, 1]) / cm.sum()
    f1  = data.get("f1", 0)
    roc = data.get("roc_auc", 0)

    fig, ax = plt.subplots(figsize=(5.5, 5))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    cmap = matplotlib.colors.LinearSegmentedColormap.from_list(
        "auris", ["#faf8f4", GOLD])
    im = ax.imshow(cm, cmap=cmap)
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    labels = ["Human", "AI"]
    for i in range(2):
        for j in range(2):
            pct = cm[i, j] / cm.sum() * 100
            ax.text(j, i, f"{cm[i, j]}\n({pct:.1f}%)",
                    ha="center", va="center", fontsize=13, fontweight="bold",
                    color="white" if cm[i, j] > cm.max() * 0.5 else "#333")

    ax.set_xticks([0, 1]); ax.set_yticks([0, 1])
    ax.set_xticklabels(labels); ax.set_yticklabels(labels)
    ax.set_xlabel("Predicted Label"); ax.set_ylabel("Actual Label")
    ax.set_title(f"Confusion Matrix — LightGBM\n"
                 f"Accuracy: {acc:.1%}  |  F1: {f1:.3f}  |  AUC: {roc:.4f}")
    _save(fig, "paper_confusion_matrix_lightgbm")


# ══════════════════════════════════════════════════════════════════════
# 3. Model Comparison Bar Chart — all 11 models
# ══════════════════════════════════════════════════════════════════════

def plot_model_comparison(ml: dict, dl: dict) -> None:
    models, accs, f1s, aucs_list = [], [], [], []

    for name, data in ml.items():
        if name.startswith("_") or not isinstance(data, dict):
            continue
        models.append(name)
        accs.append(data.get("accuracy", 0))
        f1s.append(data.get("f1", 0))
        aucs_list.append(data.get("roc_auc", 0))

    for name, data in dl.items():
        if not isinstance(data, dict):
            continue
        models.append(name)
        accs.append(data.get("accuracy", 0))
        f1s.append(data.get("f1", 0))
        aucs_list.append(data.get("roc_auc", 0))

    # Sort by AUC descending
    order = sorted(range(len(models)), key=lambda i: aucs_list[i], reverse=True)
    models   = [models[i] for i in order]
    accs     = [accs[i] for i in order]
    f1s      = [f1s[i] for i in order]
    aucs_list = [aucs_list[i] for i in order]

    x = np.arange(len(models))
    w = 0.25
    fig, ax = plt.subplots(figsize=(13, 5))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    b1 = ax.bar(x - w, accs,     w, label="Accuracy",  color="#C99347", alpha=0.85)
    b2 = ax.bar(x,     f1s,      w, label="F1 Score",  color="#6b8f7a", alpha=0.85)
    b3 = ax.bar(x + w, aucs_list, w, label="ROC-AUC",  color="#a64b3c", alpha=0.85)

    ax.set_xticks(x)
    ax.set_xticklabels(models, rotation=35, ha="right", fontsize=9)
    ax.set_ylim(0.60, 1.0)
    ax.set_ylabel("Score")
    ax.set_title("Model Performance Comparison — 11 Models (5-Fold CV, 47 Features, 5,195 Samples)")
    ax.legend(loc="lower left")

    # Annotate AUC on top of each AUC bar
    for bar, val in zip(b3, aucs_list):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.003,
                f"{val:.3f}", ha="center", va="bottom", fontsize=7.5, fontweight="bold")

    _save(fig, "paper_model_comparison")


# ══════════════════════════════════════════════════════════════════════
# 4. Fold Std Table Figure
# ══════════════════════════════════════════════════════════════════════

def plot_fold_std_table(ml: dict, dl: dict) -> None:
    rows = []

    for name, data in ml.items():
        if name.startswith("_") or not isinstance(data, dict):
            continue
        auc_val = data.get("roc_auc", 0)
        # ML has no fold_aucs stored — show mean only with dash
        rows.append((name, "ML", auc_val, None, None))

    for name, data in dl.items():
        if not isinstance(data, dict):
            continue
        fold_aucs = data.get("fold_aucs", [])
        if fold_aucs:
            mean = np.mean(fold_aucs)
            std  = np.std(fold_aucs)
            rows.append((name, "DL", mean, std, fold_aucs))
        else:
            rows.append((name, "DL", data.get("roc_auc", 0), None, None))

    rows.sort(key=lambda r: r[2], reverse=True)

    fig, ax = plt.subplots(figsize=(10, 0.45 * len(rows) + 1.5))
    fig.patch.set_facecolor(BG)
    ax.axis("off")

    col_labels = ["Model", "Type", "Mean AUC", "Std AUC", "Fold AUCs"]
    table_data = []
    for name, mtype, mean, std, folds in rows:
        std_str   = f"±{std:.4f}" if std is not None else "—"
        folds_str = ", ".join(f"{v:.4f}" for v in folds) if folds else "—"
        table_data.append([name, mtype, f"{mean:.4f}", std_str, folds_str])

    tbl = ax.table(
        cellText=table_data,
        colLabels=col_labels,
        loc="center",
        cellLoc="center",
    )
    tbl.auto_set_font_size(False)
    tbl.set_fontsize(9)
    tbl.scale(1, 1.5)

    # Header style
    for j in range(len(col_labels)):
        tbl[(0, j)].set_facecolor(GOLD)
        tbl[(0, j)].set_text_props(color="white", fontweight="bold")

    # Alternate rows
    for i in range(1, len(rows) + 1):
        color = "#f5f0e8" if i % 2 == 0 else BG
        for j in range(len(col_labels)):
            tbl[(i, j)].set_facecolor(color)

    ax.set_title("Table 1. Cross-Validation AUC Results (5-Fold) — Mean ± Std",
                 fontsize=12, fontweight="bold", pad=15)
    _save(fig, "paper_fold_std_table")


# ══════════════════════════════════════════════════════════════════════
# 5. ML vs DL Comparison — English version
# ══════════════════════════════════════════════════════════════════════

def plot_ml_vs_dl(ml: dict, dl: dict) -> None:
    ml_names, ml_accs, ml_aucs, ml_f1s = [], [], [], []
    dl_names, dl_accs, dl_aucs, dl_f1s = [], [], [], []

    for name, data in ml.items():
        if name.startswith("_") or not isinstance(data, dict):
            continue
        ml_names.append(name); ml_accs.append(data.get("accuracy", 0))
        ml_aucs.append(data.get("roc_auc", 0)); ml_f1s.append(data.get("f1", 0))

    for name, data in dl.items():
        if not isinstance(data, dict):
            continue
        dl_names.append(name); dl_accs.append(data.get("accuracy", 0))
        dl_aucs.append(data.get("roc_auc", 0)); dl_f1s.append(data.get("f1", 0))

    fig, axes = plt.subplots(1, 3, figsize=(14, 5))
    fig.patch.set_facecolor(BG)
    fig.suptitle(
        "ML vs DL Model Comparison — 47 Features, 5,195 Samples, 5-Fold CV",
        fontsize=13, fontweight="bold"
    )

    for ax, ml_vals, dl_vals, title in zip(
        axes,
        [ml_accs, ml_aucs, ml_f1s],
        [dl_accs, dl_aucs, dl_f1s],
        ["Accuracy", "ROC-AUC", "F1 Score"],
    ):
        ax.set_facecolor(BG)
        ml_y = range(len(ml_names))
        dl_y = range(len(dl_names))
        offset = len(ml_names) + 1

        bars_ml = ax.barh(list(ml_y), ml_vals, color="#C99347", alpha=0.85)
        bars_dl = ax.barh([i + offset for i in dl_y], dl_vals, color="#6b4a1e", alpha=0.85)

        ax.set_yticks(
            list(ml_y) + [i + offset for i in dl_y]
        )
        ax.set_yticklabels(ml_names + dl_names, fontsize=8)
        # Extend the right limit so the value labels printed past each bar
        # are not clipped at the axis edge.
        ax.set_xlim(0.65, 1.06)
        ax.set_xlabel(title)

        # ML / DL labels
        ax.axhline(len(ml_names) - 0.5, color="#aaa", lw=1, ls="--")
        ax.text(0.66, len(ml_names) / 2 - 0.3, "ML", fontsize=9,
                color="#C99347", fontweight="bold")
        ax.text(0.66, offset + len(dl_names) / 2 - 0.3, "DL", fontsize=9,
                color="#6b4a1e", fontweight="bold")

        for bar, val in zip(list(bars_ml) + list(bars_dl),
                            ml_vals + dl_vals):
            ax.text(bar.get_width() + 0.002, bar.get_y() + bar.get_height() / 2,
                    f"{val:.3f}", va="center", fontsize=7.5)

    plt.tight_layout()
    _save(fig, "paper_ml_vs_dl")


# ══════════════════════════════════════════════════════════════════════
# 6. Pipeline Diagram
# ══════════════════════════════════════════════════════════════════════

def plot_pipeline_diagram() -> None:
    fig, ax = plt.subplots(figsize=(14, 4))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)
    ax.axis("off")
    ax.set_xlim(0, 14)
    ax.set_ylim(0, 4)

    # Define stages
    stages = [
        (0.6,  "Audio Input\n(WAV / MP3 / FLAC)",    "#4363d8"),
        (2.6,  "Feature\nExtraction\n(librosa)",       "#3cb44b"),
        (4.6,  "47 Audio\nFeatures",                   "#f58231"),
        (6.6,  "Standard\nScaler",                     "#911eb4"),
        (8.6,  "ML / DL\nEnsemble\n(11 Models)",       "#C99347"),
        (10.6, "Probability\nFusion\n(Youden Thresh.)",  "#e6194b"),
        (12.6, "Decision\nAI / Human",                 "#1a1a1a"),
    ]

    box_w = 1.6
    box_h = 1.8
    box_y = 1.1

    for x, label, color in stages:
        rect = mpatches.FancyBboxPatch(
            (x - box_w / 2, box_y), box_w, box_h,
            boxstyle="round,pad=0.1",
            facecolor=color, edgecolor="white",
            linewidth=2, alpha=0.88,
        )
        ax.add_patch(rect)
        ax.text(x, box_y + box_h / 2, label,
                ha="center", va="center",
                fontsize=9, color="white", fontweight="bold",
                wrap=True)

    # Arrows
    for i in range(len(stages) - 1):
        x1 = stages[i][0]   + box_w / 2
        x2 = stages[i+1][0] - box_w / 2
        y  = box_y + box_h / 2
        ax.annotate("", xy=(x2, y), xytext=(x1, y),
                    arrowprops=dict(arrowstyle="->", color="#555", lw=2))

    # Sub-labels under boxes
    sub = [
        (0.6,  "File upload /\nmicrophone"),
        (2.6,  "Spectral · Temporal\nVocal · Rhythmic"),
        (4.6,  "Normalized\nvector"),
        (6.6,  "Zero mean\nUnit variance"),
        (8.6,  "7 ML + 4 DL\nmodels"),
        (10.6, "Optimal\nthreshold"),
        (12.6, "P(AI) score\n0–100%"),
    ]
    for x, txt in sub:
        ax.text(x, box_y - 0.35, txt,
                ha="center", va="top", fontsize=7.5, color="#555",
                style="italic")

    ax.set_title(
        "Figure 1. AURIS System Pipeline — End-to-End AI Music Detection",
        fontsize=12, fontweight="bold", pad=10
    )
    _save(fig, "paper_pipeline_diagram")


# ══════════════════════════════════════════════════════════════════════
# 7. Feature Importance — English
# ══════════════════════════════════════════════════════════════════════

def plot_feature_importance(ml: dict) -> None:
    imp = ml.get("_feature_importance", {})
    if not imp:
        print("  ⚠  No feature importance data")
        return

    items = sorted(imp.items(), key=lambda x: x[1], reverse=True)[:20]
    names = [k for k, _ in items]
    vals  = [v for _, v in items]

    colors_bar = [
        "#6b4a1e" if v > 0.08 else
        "#C99347" if v > 0.04 else
        "#e6c97a"
        for v in vals
    ]

    fig, ax = plt.subplots(figsize=(8, 7))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    bars = ax.barh(names[::-1], vals[::-1], color=colors_bar[::-1], edgecolor="white")
    for bar, val in zip(bars, vals[::-1]):
        ax.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height() / 2,
                f"{val:.4f}", va="center", fontsize=8)

    ax.set_xlabel("Normalized Importance")
    ax.set_title("Top 20 Feature Importances — LightGBM (Best Model)")
    ax.set_xlim(0, max(vals) * 1.18)
    _save(fig, "paper_feature_importance")


# ══════════════════════════════════════════════════════════════════════
# 8. Score Distribution — English
# ══════════════════════════════════════════════════════════════════════

def plot_score_distribution(ml: dict) -> None:
    best = ml.get("_best_model", "LightGBM")
    data = ml.get(best, {})
    y_prob_path = MODELS / "lgbm_cv_probs.npy"
    y_true_path = MODELS / "lgbm_cv_ytrue.npy"
    if not y_prob_path.exists():
        print("  SKIP  No score distribution data")
        return
    y_true = np.load(y_true_path)
    y_prob = np.load(y_prob_path)
    threshold = data.get("optimal_threshold", 0.5)

    fig, ax = plt.subplots(figsize=(7, 4.5))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    ax.hist(y_prob[y_true == 0], bins=40, alpha=0.65,
            color="#3cb44b", label=f"Human (n={int((y_true==0).sum())})", density=False)
    ax.hist(y_prob[y_true == 1], bins=40, alpha=0.65,
            color="#e6194b", label=f"AI (n={int((y_true==1).sum())})", density=False)
    ax.axvline(threshold, color="#555", ls="--", lw=1.5,
               label=f"Decision threshold ({threshold:.2f})")

    ax.set_xlabel("Predicted Probability P(AI)")
    ax.set_ylabel("Count")
    ax.set_title(f"Predicted Probability Distribution — {best}")
    ax.legend()
    _save(fig, "paper_score_distribution")


# ══════════════════════════════════════════════════════════════════════
# 9. Calibration — English
# ══════════════════════════════════════════════════════════════════════

def plot_calibration(ml: dict) -> None:
    best = ml.get("_best_model", "LightGBM")
    data = ml.get(best, {})
    y_prob_path = MODELS / "lgbm_cv_probs.npy"
    y_true_path = MODELS / "lgbm_cv_ytrue.npy"
    if not y_prob_path.exists():
        print("  SKIP  No calibration data")
        return

    y_true = np.load(y_true_path)
    y_prob = np.load(y_prob_path)

    from sklearn.metrics import brier_score_loss
    brier = brier_score_loss(y_true, y_prob)

    prob_true, prob_pred = calibration_curve(y_true, y_prob, n_bins=10)

    fig, ax = plt.subplots(figsize=(5.5, 5))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    ax.plot(prob_pred, prob_true, "o-", color=GOLD, lw=2, ms=6, label=best)
    ax.fill_between(prob_pred, prob_true, prob_pred,
                    alpha=0.12, color=GOLD)
    ax.plot([0, 1], [0, 1], "k--", lw=1, alpha=0.5, label="Perfect calibration")

    ax.text(0.05, 0.88,
            f"Brier Score = {brier:.4f}\nN = {len(y_true)} (5-Fold CV)",
            transform=ax.transAxes, fontsize=9,
            bbox=dict(boxstyle="round,pad=0.4", facecolor="white", alpha=0.7))

    ax.set_xlabel("Mean Predicted Probability")
    ax.set_ylabel("Fraction of Positives")
    ax.set_title(f"Calibration Curve — {best}")
    ax.legend()
    _save(fig, "paper_calibration")


# ══════════════════════════════════════════════════════════════════════
# 10. Precision-Recall — English
# ══════════════════════════════════════════════════════════════════════

def plot_precision_recall(ml: dict) -> None:
    best = ml.get("_best_model", "LightGBM")
    y_prob_path = MODELS / "lgbm_cv_probs.npy"
    y_true_path = MODELS / "lgbm_cv_ytrue.npy"
    if not y_prob_path.exists():
        print("  SKIP  No precision-recall data")
        return

    y_true = np.load(y_true_path)
    y_prob = np.load(y_prob_path)
    prec, rec, _ = precision_recall_curve(y_true, y_prob)
    ap = average_precision_score(y_true, y_prob)
    baseline = y_true.mean()

    fig, ax = plt.subplots(figsize=(5.5, 5))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(BG)

    ax.plot(rec, prec, color=GOLD, lw=2, label=f"{best} (AP={ap:.4f})")
    ax.fill_between(rec, prec, alpha=0.12, color=GOLD)
    ax.axhline(baseline, color="#888", ls="--", lw=1,
               label=f"Baseline = {baseline:.3f}")

    ax.set_xlabel("Recall")
    ax.set_ylabel("Precision")
    ax.set_title(f"Precision-Recall Curve — {best}")
    ax.legend()
    _save(fig, "paper_precision_recall")


# ══════════════════════════════════════════════════════════════════════
# MAIN
# ══════════════════════════════════════════════════════════════════════

def main() -> None:
    print(f"Output dir: {OUT}\n")
    ml, dl = _load_results()

    print("Generating figures...")
    plot_roc_curves(ml, dl)
    plot_confusion_matrix_lightgbm(ml)
    plot_model_comparison(ml, dl)
    plot_fold_std_table(ml, dl)
    plot_ml_vs_dl(ml, dl)
    plot_pipeline_diagram()
    plot_feature_importance(ml)
    plot_score_distribution(ml)
    plot_calibration(ml)
    plot_precision_recall(ml)

    print(f"\nDone. All figures saved to {OUT}")


if __name__ == "__main__":
    main()