#!/usr/bin/env python3
"""Plot VERL SFT training metrics from captured console logs."""

from __future__ import annotations

import argparse
import csv
import math
import re
from collections import defaultdict
from pathlib import Path

import matplotlib

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


ANSI_RE = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]")
STEP_RE = re.compile(r"\bstep:(\d+)\b")
METRIC_RE = re.compile(
    r"(?P<key>[A-Za-z0-9_./()]+):(?P<value>[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:e[+-]?\d+)?)",
    re.IGNORECASE,
)


def strip_ansi(text: str) -> str:
    return ANSI_RE.sub("", text)


def parse_log(path: Path) -> list[dict[str, float]]:
    """Parse lines like `step:30 - train/loss:... - val/loss:...`."""
    by_step: dict[int, dict[str, float]] = defaultdict(dict)

    for raw_line in path.read_text(errors="replace").splitlines():
        line = strip_ansi(raw_line)
        step_match = STEP_RE.search(line)
        if not step_match:
            continue

        step = int(step_match.group(1))
        row = by_step[step]
        row["step"] = float(step)

        for match in METRIC_RE.finditer(line):
            key = match.group("key")
            if key == "step":
                continue
            try:
                row[key] = float(match.group("value"))
            except ValueError:
                continue

    return [by_step[step] for step in sorted(by_step)]


def write_csv(rows: list[dict[str, float]], path: Path) -> None:
    keys = sorted({key for row in rows for key in row}, key=lambda k: (k != "step", k))
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=keys)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def finite_xy(rows: list[dict[str, float]], key: str) -> tuple[list[float], list[float]]:
    xs: list[float] = []
    ys: list[float] = []
    for row in rows:
        value = row.get(key)
        step = row.get("step")
        if value is None or step is None or not math.isfinite(value):
            continue
        xs.append(step)
        ys.append(value)
    return xs, ys


def plot_keys(
    ax: plt.Axes,
    rows: list[dict[str, float]],
    keys: list[str],
    title: str,
    ylabel: str | None = None,
) -> bool:
    plotted = False
    for key in keys:
        xs, ys = finite_xy(rows, key)
        if not xs:
            continue
        ax.plot(xs, ys, marker="o", linewidth=1.8, markersize=4, label=key)
        plotted = True

    ax.set_title(title)
    ax.set_xlabel("step")
    if ylabel:
        ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.25)
    if plotted and len([key for key in keys if finite_xy(rows, key)[0]]) > 1:
        ax.legend(fontsize=8)
    return plotted


def build_summary_text(rows: list[dict[str, float]], log_path: Path) -> str:
    last = rows[-1]

    best_val = None
    for row in rows:
        if "val/loss" not in row:
            continue
        pair = (row["val/loss"], int(row["step"]))
        if best_val is None or pair[0] < best_val[0]:
            best_val = pair

    lines = [
        f"log: {log_path}",
        f"steps parsed: {len(rows)}",
        f"first step: {int(rows[0]['step'])}",
        f"last step: {int(last['step'])}",
    ]
    if "train/loss" in last:
        lines.append(f"last train/loss: {last['train/loss']:.4g}")
    if best_val is not None:
        lines.append(f"best val/loss: {best_val[0]:.4g} at step {best_val[1]}")
    if "perf/max_memory_allocated_gb" in last:
        lines.append(f"last max GPU alloc: {last['perf/max_memory_allocated_gb']:.2f} GB")
    if "perf/cpu_memory_used_gb" in last:
        lines.append(f"last CPU memory: {last['perf/cpu_memory_used_gb']:.2f} GB")

    return " | ".join(lines)


def make_plot(rows: list[dict[str, float]], log_path: Path, output: Path, title: str) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)

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

    panels = [
        (
            axes[0][0],
            ["train/loss", "val/loss"],
            "Loss",
            "loss",
        ),
        (
            axes[0][1],
            ["train/grad_norm"],
            "Gradient Norm",
            "norm",
        ),
        (
            axes[1][0],
            ["perf/max_memory_allocated_gb", "perf/max_memory_reserved_gb"],
            "GPU Memory",
            "GB",
        ),
        (
            axes[1][1],
            ["perf/cpu_memory_used_gb"],
            "CPU Memory",
            "GB",
        ),
        (
            axes[2][0],
            ["train/global_tokens"],
            "Global Tokens per Step",
            "tokens",
        ),
        (
            axes[2][1],
            ["train/total_tokens(B)"],
            "Cumulative Tokens",
            None,
        ),
        (
            axes[3][0],
            ["train/lr"],
            "Learning Rate",
            "lr",
        ),
        (
            axes[3][1],
            ["train/mfu"],
            "MFU",
            "mfu",
        ),
    ]

    for ax, keys, panel_title, ylabel in panels:
        plotted = plot_keys(ax, rows, keys, panel_title, ylabel)
        if not plotted:
            ax.set_title(panel_title)
            ax.axis("off")
            ax.text(0.5, 0.5, "metric not found", ha="center", va="center")

    fig.tight_layout(rect=[0, 0.035, 1, 0.965])
    fig.text(
        0.01,
        0.01,
        build_summary_text(rows, log_path),
        ha="left",
        va="bottom",
        family="monospace",
        fontsize=8,
    )
    fig.savefig(output, dpi=180)
    plt.close(fig)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Plot VERL SFT training metrics from a stdout/stderr log captured with tee."
    )
    parser.add_argument("--log", required=True, type=Path, help="Captured training log.")
    parser.add_argument("--output", required=True, type=Path, help="Output PNG path.")
    parser.add_argument("--title", default=None, help="Figure title.")
    parser.add_argument("--csv", type=Path, help="Optional parsed metrics CSV path.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    rows = parse_log(args.log)
    if not rows:
        raise SystemExit(
            f"No VERL metric lines found in {args.log}. Capture stdout/stderr with tee, for example:\n"
            f"  ... bash run_verl_sft.sh ... 2>&1 | tee plotting/logs/run.log"
        )

    title = args.title or args.log.stem.replace("_", " ")
    make_plot(rows, args.log, args.output, title)
    if args.csv:
        write_csv(rows, args.csv)
    print(f"Wrote {args.output}")
    if args.csv:
        print(f"Wrote {args.csv}")


if __name__ == "__main__":
    main()