source/gconvex_authors/scripts/plot_training_log.py

"""Parse FCOT-Separable training logs and visualize dual/grad/active metrics."""

import argparse
import os
import re
from dataclasses import dataclass, field
from datetime import datetime
from typing import List, Dict, Any

import matplotlib.pyplot as plt
import numpy as np


@dataclass
class RunRecord:
    meta: Dict[str, Any] = field(default_factory=dict)
    iters: List[int] = field(default_factory=list)
    timestamps: List[str] = field(default_factory=list)
    dual: List[float] = field(default_factory=list)
    u: List[float] = field(default_factory=list)
    uc: List[float] = field(default_factory=list)
    grad: List[float] = field(default_factory=list)
    temp: List[float] = field(default_factory=list)
    active: List[float] = field(default_factory=list)
    react_steps: List[int] = field(default_factory=list)
    refresh_steps: List[int] = field(default_factory=list)


ARCH_RE = re.compile(
    r"\[FCOT-SEP ARCH\] dim=(?P<dim>\d+), radius=(?P<radius>[\d\.]+), ny=(?P<ny>\d+)"
)

ITER_RE = re.compile(
    r"\[(?P<ts>[^]]+)\] INFO:training: "
    r"\[Iter (?P<iter>\d+)\] dual=(?P<dual>[-0-9.e+]+) "
    r"u=(?P<u>[-0-9.e+]+) uc=(?P<uc>[-0-9.e+]+) "
    r"grad=(?P<grad>[-0-9.e+]+).* temp=(?P<temp>[-0-9.e+]+) "
    r"active=(?P<active>[-0-9.e+]+)%"
)

REACT_RE = re.compile(
    r"\[FCOT-SEP\] Reactivated .* at step (?P<step>\d+)"
)

REFRESH_RE = re.compile(
    r"\[FCOT-SEP\] Full refresh updated .* at step (?P<step>\d+)"
)

TRAIN_PATH_RE = re.compile(
    r"tmp/FCOTSeparable_dim(?P<dim>\d+)_kernel-(?P<kernel>.+?)_"
)


def parse_training_log(path: str) -> List[RunRecord]:
    """
    Scan `training.log` and extract iteration stats plus reactivation/refresh markers.

    Parameters
    ----------
    path : str
        Path to the training log file generated by FCOT-Separable.

    Returns
    -------
    List[RunRecord]
        Parsed run records sorted in chronological order.
    """
    runs: List[RunRecord] = []
    current: RunRecord | None = None

    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            # Start of a new run: architecture line
            if "[FCOT-SEP ARCH]" in line:
                m = ARCH_RE.search(line)
                current = RunRecord()
                if m:
                    current.meta["dim"] = int(m.group("dim"))
                    current.meta["radius"] = float(m.group("radius"))
                    current.meta["ny"] = int(m.group("ny"))
                runs.append(current)
                continue

            if current is None:
                # Skip lines before the first ARCH
                continue

            # Iteration line with dual/grad/etc.
            if "[Iter " in line:
                m = ITER_RE.search(line)
                if not m:
                    continue
                current.timestamps.append(m.group("ts"))
                current.iters.append(int(m.group("iter")))
                current.dual.append(float(m.group("dual")))
                current.u.append(float(m.group("u")))
                current.uc.append(float(m.group("uc")))
                current.grad.append(float(m.group("grad")))
                current.temp.append(float(m.group("temp")))
                current.active.append(float(m.group("active")))
                continue

            # Reactivation events
            if "[FCOT-SEP] Reactivated" in line:
                m = REACT_RE.search(line)
                if m:
                    current.react_steps.append(int(m.group("step")))
                continue

            # Full refresh events
            if "[FCOT-SEP] Full refresh updated" in line:
                m = REFRESH_RE.search(line)
                if m:
                    current.refresh_steps.append(int(m.group("step")))
                continue

            # Lines that mention checkpoint paths (TRAIN/CACHE/SAVE) → extract kernel
            if "[TRAIN] No checkpoint found at" in line or "[CACHE] Found saved" in line or "[SAVE] Writing checkpoint to" in line:
                m = TRAIN_PATH_RE.search(line)
                if m and current is not None:
                    current.meta["kernel"] = m.group("kernel")
                continue

    return runs


def plot_runs(
    runs: List[RunRecord],
    out_dir: str = "tmp/training_plots",
) -> None:
    os.makedirs(out_dir, exist_ok=True)

    for idx, run in enumerate(runs):
        fig, axes = plot_run(run, idx=idx, show=False)
        dim = run.meta.get("dim", "?")
        fname = os.path.join(out_dir, f"run_{idx}_dim{dim}.png")
        fig.savefig(fname, dpi=150)
        plt.close(fig)


def plot_run(run: RunRecord, idx: int | None = None, show: bool = True):
    """
    Plot a single RunRecord (dual, grad, active) and optionally show inline.
    Returns (fig, axes) so it can be used easily from notebooks.
    """
    if not run.iters:
        raise ValueError("Run has no iteration data to plot.")

    iters = np.array(run.iters)
    dual = np.array(run.dual)
    grad = np.array(run.grad)
    active = np.array(run.active)  # already in percent

    dim = run.meta.get("dim", "?")
    kernel = run.meta.get("kernel", "?")

    fig, axes = plt.subplots(3, 1, sharex=True, figsize=(12, 8))

    # Duration estimate from first/last timestamp (if available)
    duration_str = ""
    if run.timestamps:
        try:
            t0 = datetime.strptime(run.timestamps[0], "%Y-%m-%d %H:%M:%S,%f")
            t1 = datetime.strptime(run.timestamps[-1], "%Y-%m-%d %H:%M:%S,%f")
            dt_min = (t1 - t0).total_seconds() / 60.0
            duration_str = f", Δt≈{dt_min:.1f} min"
        except Exception:
            duration_str = ""

    # Dual
    title_suffix = f" (run {idx})" if idx is not None else ""
    axes[0].plot(iters, dual, label="dual")
    axes[0].set_ylabel("dual")
    axes[0].set_title(f"dim={dim}, kernel={kernel}{duration_str}{title_suffix}")

    # Gradient norm
    axes[1].plot(iters, grad, label="grad_norm", color="tab:orange")
    axes[1].set_ylabel("grad norm")

    # Active fraction
    axes[2].plot(iters, active, label="active %", color="tab:green")
    axes[2].set_ylabel("active (%)")
    axes[2].set_xlabel("iteration")

    # Reactivation / refresh markers
    for ax in axes:
        for s in run.react_steps:
            ax.axvline(s, color="red", linewidth=0.6, alpha=0.6)
        for s in run.refresh_steps:
            ax.axvline(s, color="blue", linewidth=0.6, alpha=0.6)

    axes[0].legend(loc="best")
    axes[1].legend(loc="best")
    axes[2].legend(loc="best")

    fig.tight_layout()
    if show:
        plt.show()
    return fig, axes


def main():
    """Entry point to parse `training.log` and dump summary plots to disk."""
    parser = argparse.ArgumentParser(
        description="Parse training.log and plot dual/grad/active with reactivation/refresh markers."
    )
    parser.add_argument(
        "--log",
        type=str,
        default="training.log",
        help="Path to training log (default: training.log)",
    )
    parser.add_argument(
        "--out",
        type=str,
        default="tmp/training_plots",
        help="Output directory for plots (default: tmp/training_plots)",
    )
    args = parser.parse_args()

    runs = parse_training_log(args.log)
    if not runs:
        print(f"No runs found in {args.log}")
        return

    print(f"Parsed {len(runs)} runs from {args.log}")
    plot_runs(runs, out_dir=args.out)
    print(f"Saved plots to {args.out}")


if __name__ == "__main__":
    main()