#!/usr/bin/env python3
"""Plot checkpoint evaluation summaries produced by evaluate_sft_loss.py."""

from __future__ import annotations

import argparse
import csv
import json
import re
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt


STEP_RE = re.compile(r"step(\d+)")


def step_from_path(path: Path) -> int:
    match = STEP_RE.search(path.name)
    if not match:
        raise ValueError(f"Could not infer checkpoint step from {path}")
    return int(match.group(1))


def load_jsonl(path: Path) -> list[dict]:
    rows = []
    if not path.exists():
        return rows
    with path.open() as handle:
        for line in handle:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def word_count(text: object) -> int:
    return len(str(text or "").split())


def read_eval(summary_paths: list[Path], examples_dir: Path, generations_dir: Path) -> list[dict]:
    rows = []
    for summary_path in sorted(summary_paths, key=step_from_path):
        step = step_from_path(summary_path)
        summary = json.loads(summary_path.read_text())

        examples_path = examples_dir / summary_path.name.replace("_val_loss.json", "_val_loss_examples.jsonl")
        examples = [row for row in load_jsonl(examples_path) if row.get("status") == "scored"]
        losses = [float(row["loss"]) for row in examples if "loss" in row]

        generations_path = generations_dir / summary_path.name.replace("_val_loss.json", "_val_outputs.jsonl")
        generations = load_jsonl(generations_path)
        prediction_words = [word_count(row.get("prediction")) for row in generations]
        reference_words = [word_count(row.get("reference")) for row in generations]

        pred_ref_ratio = None
        if prediction_words and reference_words and sum(reference_words) > 0:
            pred_ref_ratio = sum(prediction_words) / sum(reference_words)

        rows.append(
            {
                "step": step,
                "loss": float(summary["loss"]),
                "perplexity": float(summary["perplexity"]),
                "examples_scored": int(summary.get("examples_scored", 0)),
                "examples_skipped": int(summary.get("examples_skipped", 0)),
                "assistant_tokens": int(summary.get("assistant_tokens", 0)),
                "per_example_losses": losses,
                "mean_prediction_words": sum(prediction_words) / len(prediction_words) if prediction_words else None,
                "mean_reference_words": sum(reference_words) / len(reference_words) if reference_words else None,
                "prediction_reference_word_ratio": pred_ref_ratio,
            }
        )
    if not rows:
        raise SystemExit("No evaluation summaries found.")
    return rows


def write_csv(rows: list[dict], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "step",
        "loss",
        "perplexity",
        "examples_scored",
        "examples_skipped",
        "assistant_tokens",
        "mean_prediction_words",
        "mean_reference_words",
        "prediction_reference_word_ratio",
    ]
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow({key: row.get(key) for key in fieldnames})


def plot(rows: list[dict], output: Path, title: str) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    steps = [row["step"] for row in rows]

    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    fig.suptitle(title, fontsize=16, y=0.985)

    ax = axes[0][0]
    ax.plot(steps, [row["loss"] for row in rows], marker="o", linewidth=2)
    best = min(rows, key=lambda row: row["loss"])
    ax.scatter([best["step"]], [best["loss"]], s=80, zorder=3, label=f"best step {best['step']}")
    ax.set_title("Validation Loss")
    ax.set_xlabel("checkpoint step")
    ax.set_ylabel("assistant-token NLL")
    ax.grid(True, alpha=0.25)
    ax.legend()

    ax = axes[0][1]
    ax.plot(steps, [row["perplexity"] for row in rows], marker="o", color="tab:orange", linewidth=2)
    ax.set_title("Validation Perplexity")
    ax.set_xlabel("checkpoint step")
    ax.set_ylabel("perplexity")
    ax.grid(True, alpha=0.25)

    ax = axes[1][0]
    loss_lists = [row["per_example_losses"] for row in rows]
    if any(loss_lists):
        ax.boxplot(loss_lists, tick_labels=[str(step) for step in steps], showmeans=True)
    ax.set_title("Per-Example Loss Distribution")
    ax.set_xlabel("checkpoint step")
    ax.set_ylabel("loss")
    ax.grid(True, axis="y", alpha=0.25)

    ax = axes[1][1]
    pred_lengths = [row["mean_prediction_words"] for row in rows]
    ref_lengths = [row["mean_reference_words"] for row in rows]
    if any(value is not None for value in pred_lengths):
        ax.plot(steps, pred_lengths, marker="o", label="prediction words", linewidth=2)
    if any(value is not None for value in ref_lengths):
        ax.plot(steps, ref_lengths, marker="o", label="reference words", linewidth=2)
    ratio = [row["prediction_reference_word_ratio"] for row in rows]
    if any(value is not None for value in ratio):
        ax2 = ax.twinx()
        ax2.plot(steps, ratio, marker="s", linestyle="--", color="tab:green", label="pred/ref ratio")
        ax2.set_ylabel("prediction/reference word ratio")
        ax2.legend(loc="lower right")
    ax.set_title("Generated Answer Length")
    ax.set_xlabel("checkpoint step")
    ax.set_ylabel("mean words")
    ax.grid(True, alpha=0.25)
    ax.legend(loc="upper left")

    summary = (
        f"best loss {best['loss']:.4f} at step {best['step']} | "
        f"perplexity {best['perplexity']:.4f} | "
        f"examples {best['examples_scored']} | assistant tokens {best['assistant_tokens']}"
    )
    fig.tight_layout(rect=[0, 0.045, 1, 0.955])
    fig.text(0.01, 0.012, summary, ha="left", va="bottom", family="monospace", fontsize=9)
    fig.savefig(output, dpi=180)
    plt.close(fig)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary-pattern", help="Glob for *_val_loss.json files.")
    parser.add_argument("--summaries", nargs="*", type=Path, help="Explicit *_val_loss.json files to plot.")
    parser.add_argument(
        "--examples-dir",
        type=Path,
        default=Path("data/hep_sft/checkpoint_loss"),
        help="Directory containing per-example loss JSONL files.",
    )
    parser.add_argument(
        "--generations-dir",
        type=Path,
        default=Path("data/hep_sft/checkpoint_eval"),
        help="Directory containing generated output JSONL files.",
    )
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--csv", type=Path)
    parser.add_argument("--title", default="Checkpoint Evaluation")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.summaries:
        summary_paths = args.summaries
    elif args.summary_pattern:
        summary_paths = sorted(Path().glob(args.summary_pattern), key=step_from_path)
    else:
        raise SystemExit("Pass either --summaries or --summary-pattern.")
    rows = read_eval(summary_paths, args.examples_dir, args.generations_dir)
    plot(rows, args.output, args.title)
    print(f"Wrote {args.output}")
    if args.csv:
        write_csv(rows, args.csv)
        print(f"Wrote {args.csv}")


if __name__ == "__main__":
    main()