scripts/generate_report_figures.py

"""
Generate publication-style figures for docs/experiment_report_latex/figures/.

Reads archived metrics from result/training_summary.json and result/evaluation_results.json
(no fabricated scores). Also draws a schematic pipeline diagram (no numeric claims).

Usage (from repo root):
    python scripts/generate_report_figures.py
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch

REPO_ROOT = Path(__file__).resolve().parent.parent
RESULT_DIR = REPO_ROOT / "result"
FIG_DIR = REPO_ROOT / "docs" / "experiment_report_latex" / "figures"

plt.rcParams.update(
    {
        "figure.dpi": 120,
        "savefig.dpi": 160,
        "font.size": 10,
        "axes.titlesize": 11,
        "axes.labelsize": 10,
        "axes.unicode_minus": False,
        "axes.grid": True,
        "grid.alpha": 0.25,
        "grid.linestyle": "--",
    }
)


def _load_json(path: Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def plot_results_panel(summary_path: Path, eval_path: Path, out_path: Path) -> None:
    summary = _load_json(summary_path)
    ev = _load_json(eval_path)

    train_hist = summary.get("train_loss_history") or []
    val_hist = summary.get("val_metrics_history") or []
    val_losses = [v["val_loss"] for v in val_hist if isinstance(v, dict) and "val_loss" in v]

    fig, axes = plt.subplots(2, 2, figsize=(10.5, 8.0))
    fig.suptitle("EasyTranslate — archived run (result/*.json)", fontsize=12, fontweight="bold")

    # (a) Training loss
    ax = axes[0, 0]
    if train_hist:
        ep = list(range(1, len(train_hist) + 1))
        ax.plot(ep, train_hist, "o-", color="#1f77b4", lw=2, ms=6)
        ax.set_xlabel("Epoch")
        ax.set_ylabel("Train CE loss")
        ax.set_title("(a) Training loss")
        ax.set_xticks(ep)
    else:
        ax.text(0.5, 0.5, "No train_loss_history", ha="center", va="center", transform=ax.transAxes)
        ax.set_axis_off()

    # (b) Validation loss
    ax = axes[0, 1]
    if val_losses:
        ep = list(range(1, len(val_losses) + 1))
        ax.plot(ep, val_losses, "s-", color="#d62728", lw=2, ms=6)
        ax.set_xlabel("Epoch")
        ax.set_ylabel("Validation loss")
        ax.set_title("(b) Validation loss")
        ax.set_xticks(ep)
    else:
        ax.text(0.5, 0.5, "No val loss", ha="center", va="center", transform=ax.transAxes)
        ax.set_axis_off()

    # (c) BLEU n-gram breakdown
    ax = axes[1, 0]
    keys = [("bleu_1", "BLEU-1"), ("bleu_2", "BLEU-2"), ("bleu_3", "BLEU-3"), ("bleu_4", "BLEU-4")]
    labels = [k[1] for k in keys]
    vals = [float(ev.get(k[0], 0.0)) for k in keys]
    colors = ["#2ca02c", "#98df8a", "#aec7e8", "#6baed6"]
    bars = ax.bar(labels, vals, color=colors, edgecolor="#333", linewidth=0.6)
    ax.set_ylabel("Score")
    ax.set_title("(c) N-gram BLEU breakdown")
    ax.set_ylim(0, max(vals) * 1.15 + 1e-6)
    for b, v in zip(bars, vals):
        ax.text(b.get_x() + b.get_width() / 2, v + 0.8, f"{v:.1f}", ha="center", va="bottom", fontsize=9)

    # (d) Corpus BLEU + chrF (TER annotated — different scale)
    ax = axes[1, 1]
    bleu_c = float(ev.get("bleu", 0.0))
    chrf = float(ev.get("chrf", 0.0))
    ter = float(ev.get("ter", 0.0))
    x = ["Corpus BLEU", "chrF++"]
    y = [bleu_c, chrf]
    ax.bar(x, y, color=["#9467bd", "#ff7f0e"], edgecolor="#333", linewidth=0.6)
    ax.set_ylabel("Score")
    ax.set_title("(d) Corpus BLEU & chrF++ (TER in caption)")
    ax.set_ylim(0, max(y) * 1.2 + 1e-6)
    for i, v in enumerate(y):
        ax.text(i, v + 0.4, f"{v:.2f}", ha="center", va="bottom", fontsize=9)
    ax.text(
        0.5,
        -0.22,
        f"TER = {ter:.2f} (lower is better; same run as evaluation_results.json / main metrics table)",
        transform=ax.transAxes,
        ha="center",
        fontsize=9,
        style="italic",
    )

    fig.tight_layout(rect=[0, 0.02, 1, 0.96])
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)


def plot_experiment_pipeline(out_path: Path) -> None:
    """Schematic only — English labels inside figure to avoid font issues in Matplotlib."""
    fig, ax = plt.subplots(figsize=(12.5, 3.2))
    ax.set_xlim(0, 12)
    ax.set_ylim(0, 3)
    ax.axis("off")

    def box(cx: float, cy: float, w: float, h: float, text: str) -> FancyBboxPatch:
        x, y = cx - w / 2, cy - h / 2
        p = FancyBboxPatch(
            (x, y),
            w,
            h,
            boxstyle="round,pad=0.05,rounding_size=0.12",
            linewidth=1.2,
            edgecolor="#2c3e50",
            facecolor="#ecf0f1",
        )
        ax.add_patch(p)
        ax.text(cx, cy, text, ha="center", va="center", fontsize=9, fontweight="medium", color="#2c3e50")
        return p

    def arrow(x1: float, y1: float, x2: float, y2: float) -> None:
        arr = FancyArrowPatch(
            (x1, y1),
            (x2, y2),
            arrowstyle="-|>",
            mutation_scale=12,
            linewidth=1.4,
            color="#34495e",
        )
        ax.add_patch(arr)

    y = 1.55
    specs = [
        (1.0, "Corpus\n(WMT19 zh--en)"),
        (2.85, "Preprocess\n& tokenize"),
        (4.75, "Model\n(scratch / NLLB)"),
        (6.65, "Train\n(AdamW, sched.)"),
        (8.45, "Best\nckpt"),
        (10.15, "Decode\n(beam / greedy)"),
        (11.55, "Metrics\n(SacreBLEU, …)"),
    ]
    w, h = 1.05, 0.95
    for cx, txt in specs:
        box(cx, y, w, h, txt)

    xs = [s[0] for s in specs]
    for a, b in zip(xs[:-1], xs[1:]):
        arrow(a + w / 2 + 0.02, y, b - w / 2 - 0.02, y)

    ax.text(
        6.0,
        2.55,
        "EasyTranslate evaluation pipeline (schematic)",
        ha="center",
        fontsize=11,
        fontweight="bold",
        color="#2c3e50",
    )
    fig.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)


def plot_metric_sparkline(eval_path: Path, out_path: Path) -> None:
    """Single-row horizontal bar: main metrics for slide-style summary."""
    ev = _load_json(eval_path)
    labels = ["BLEU", "chrF++", "BLEU-4"]
    vals = [float(ev.get("bleu", 0)), float(ev.get("chrf", 0)), float(ev.get("bleu_4", 0))]
    fig, ax = plt.subplots(figsize=(8.0, 3.2))
    y_pos = range(len(labels))
    ax.barh(list(y_pos), vals, color=["#1f77b4", "#ff7f0e", "#2ca02c"], height=0.55, edgecolor="#333")
    ax.set_yticks(list(y_pos))
    ax.set_yticklabels(labels)
    ax.invert_yaxis()
    ax.set_xlabel("Score")
    ax.set_title("Main automatic metrics (archived evaluation_results.json)")
    for i, v in enumerate(vals):
        ax.text(v + 0.5, i, f"{v:.2f}", va="center", fontsize=10)
    ax.set_xlim(0, max(vals) * 1.35 + 5)
    fig.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)


def main() -> int:
    summary_path = RESULT_DIR / "training_summary.json"
    eval_path = RESULT_DIR / "evaluation_results.json"
    if not summary_path.exists():
        print(f"Missing {summary_path}", file=sys.stderr)
        return 1
    if not eval_path.exists():
        print(f"Missing {eval_path}", file=sys.stderr)
        return 1

    FIG_DIR.mkdir(parents=True, exist_ok=True)
    plot_results_panel(summary_path, eval_path, FIG_DIR / "figure_results_panel.png")
    plot_experiment_pipeline(FIG_DIR / "figure_experiment_pipeline.png")
    plot_metric_sparkline(eval_path, FIG_DIR / "figure_main_metrics_horizontal.png")
    print(f"Wrote figures to {FIG_DIR}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())