#!/usr/bin/env python3
"""Live L15 diagnostic plots from training logs.

Parses `logs/diag_L15_*.log` and writes/overwrites figures under `figs/diag_L15/`:

  A_train_loss.png          train CE loss vs epoch (per arm)
  B_stage_timeline.png      scheduled stage vs epoch + promote markers
  C_frontier_heatmap.png    per-hop frontier accuracy over eval epochs
  D_ce_score_heatmap.png    per-hop ce_score (backtracking metric) over evals
  E_time_to_threshold.png   epochs-to-first-hit of frontier>=thr per hop
  F_latest_bars.png         latest per-hop frontier vs ce_score

Also writes `figs/diag_L15/summary.json` with time-to-threshold tables.

Usage:
  PYTHONPATH=. python scripts/plot_diag_L15.py
  PYTHONPATH=. python scripts/plot_diag_L15.py --watch 120   # replot every 120s
"""
from __future__ import annotations

import argparse
import ast
import json
import re
import time
from pathlib import Path

import matplotlib

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

ROOT = Path(__file__).resolve().parents[1]
LOG_DIR = ROOT / "logs"
OUT_DIR = ROOT / "figs" / "diag_L15"

# Labels encode BOTH gates: promX = stage advance, btY = backtrack check.
ARMS = [
    ("promF.95 / btCE.50", "diag_L15_promF095_btCE050", "#1b7837"),
    ("promF.95 / btCE.90", "diag_L15_promF095_btCE090", "#00441b"),
    ("promF.95 / btCE.95", "diag_L15_promF095_btCE095", "#006d2c"),
    ("promF.99 / btCE.90", "diag_L15_promF099_btCE090", "#762a83"),
    ("promF.99 / btCE.50", "diag_L15_promF099_btCE050", "#c51b7d"),
    ("promF.95 / btF.95", "diag_L15_promF095_btF095", "#1f78b4"),
    ("promCE.90 / btCE.90", "diag_L15_promCE090_btCE090", "#d95f02"),
    ("promF.95 / btNONE", "diag_L15_promF095_btNONE", "#a6611a"),
]


def parse_log(path: Path) -> dict:
    if not path.exists():
        return {}
    # stream large logs; keep last 12MB for speed (full history still usually fits)
    with open(path, "rb") as f:
        f.seek(0, 2)
        size = f.tell()
        f.seek(0 if size < 12_000_000 else size - 12_000_000)
        text = f.read().decode("utf-8", errors="ignore")

    # If we truncated mid-file, also pull promote history from a head+grep via
    # a second full regex on a larger window when small enough.
    if size < 40_000_000:
        text = path.read_text(errors="ignore")

    train = [
        (int(e), int(s), float(l))
        for e, s, l in re.findall(
            r"train epoch (\d+)/\d+ stage=(\d+) loss=([0-9.]+)", text
        )
    ]
    promotes = [
        (int(a), int(b), int(ep))
        for a, b, ep in re.findall(
            r"PROMOTE stage (\d+) -> (\d+) \|.*? in (\d+) epochs", text
        )
    ]
    backtracks = [
        (m, None if t == "None" else int(t), float(acc))
        for m, t, acc in re.findall(
            r"backtrack\[([^\]>=]+)[^]]*\]: target_stage=(\S+) min_mastered[_a-z]*=([0-9.]+)",
            text,
        )
    ]
    evals = []
    for ep_m, metric, raw in re.findall(
        r"train epoch (\d+)/\d+.*?\n.*?eval per-hop \(staging_metric=(\w+)\): (\{.*?\})",
        text,
        flags=re.S,
    ):
        # The above may be too greedy; fall back below.
        pass

    # Robust: pair nearest preceding train epoch with each eval line
    lines = text.splitlines()
    cur_ep, cur_stage = None, None
    for line in lines:
        m = re.match(r"train epoch (\d+)/\d+ stage=(\d+) loss=([0-9.]+)", line)
        if m:
            cur_ep, cur_stage = int(m.group(1)), int(m.group(2))
            continue
        # Compact format:
        # eval (prom=frontier@0.95 bt=ce_score@0.5) frontier=[1:1.00 2:0.98 ...]  ce_score=[1:0.95 ...]
        m = re.match(
            r"eval \(prom=(\w+)@([0-9.]+) bt=(\w+)@([0-9.]+)\) "
            r"frontier=\[([^\]]*)\]\s+ce_score=\[([^\]]*)\]",
            line,
        )
        if m and cur_ep is not None:
            def _parse_pairs(s):
                out = {}
                for tok in s.split():
                    if ":" not in tok:
                        continue
                    h, v = tok.split(":", 1)
                    try:
                        out[int(h)] = float(v)
                    except ValueError:
                        pass
                return out

            fr = _parse_pairs(m.group(5))
            ce = _parse_pairs(m.group(6))
            hops = {
                h: {
                    "frontier": fr.get(h, float("nan")),
                    "ce_score": ce.get(h, float("nan")),
                    "reachable": float("nan"),
                    "optimal": float("nan"),
                    "superposition": float("nan"),
                }
                for h in sorted(set(fr) | set(ce))
            }
            evals.append(
                {
                    "epoch": cur_ep,
                    "stage": cur_stage,
                    "metric": m.group(1),
                    "hops": hops,
                }
            )
            continue
        # Legacy full-dict format
        m = re.match(
            r"eval per-hop \((?:staging_metric=(\w+)|promote=(\w+)@([0-9.]+), "
            r"backtrack=(\w+)@([0-9.]+))\): (\{.*\})",
            line,
        )
        if m and cur_ep is not None:
            try:
                d = ast.literal_eval(m.group(6))
            except Exception:
                continue
            metric = m.group(1) or m.group(2) or "?"
            evals.append(
                {
                    "epoch": cur_ep,
                    "stage": cur_stage,
                    "metric": metric,
                    "hops": d,
                }
            )

    return {
        "train": train,
        "promotes": promotes,
        "backtracks": backtracks,
        "evals": evals,
    }


def time_to_threshold(evals, key: str, thr: float, max_hop: int = 15):
    """First eval epoch where hop h's `key` >= thr, for each hop."""
    first = {}
    for ev in evals:
        for h, c in ev["hops"].items():
            if h in first:
                continue
            if c.get(key, 0) >= thr:
                first[h] = ev["epoch"]
    return {h: first.get(h) for h in range(1, max_hop + 1)}


def plot_all():
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    parsed = {}
    for label, name, color in ARMS:
        p = parse_log(LOG_DIR / f"{name}.log")
        if p.get("train") or p.get("evals"):
            parsed[name] = {**p, "label": label, "color": color}

    if not parsed:
        print("no diag logs yet")
        return

    # A: train loss
    fig, ax = plt.subplots(figsize=(9, 4))
    for name, p in parsed.items():
        if not p["train"]:
            continue
        xs = [t[0] for t in p["train"]]
        ys = [t[2] for t in p["train"]]
        ax.plot(xs, ys, color=p["color"], label=p["label"], lw=1.2, alpha=0.9)
    ax.set_xlabel("epoch")
    ax.set_ylabel("train CE loss")
    ax.set_title("L15 diagnostic — training loss")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(OUT_DIR / "A_train_loss.png", dpi=140)
    plt.close(fig)

    # B: stage timeline
    fig, ax = plt.subplots(figsize=(9, 4))
    for name, p in parsed.items():
        if not p["train"]:
            continue
        xs = [t[0] for t in p["train"]]
        ys = [t[1] for t in p["train"]]
        ax.step(xs, ys, where="post", color=p["color"], label=p["label"], lw=1.5)
        for a, b, ep in p["promotes"]:
            ax.axvline(ep, color=p["color"], alpha=0.15, lw=0.8)
    ax.set_xlabel("epoch")
    ax.set_ylabel("scheduled stage")
    ax.set_title("L15 diagnostic — stage timeline (promote markers faint)")
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(OUT_DIR / "B_stage_timeline.png", dpi=140)
    plt.close(fig)

    # C/D heatmaps per arm
    for key, fname, title, vmin, vmax, cmap in [
        ("frontier", "C_frontier_heatmap", "frontier accuracy", 0, 1, "viridis"),
        ("ce_score", "D_ce_score_heatmap", "ce_score (BT metric)", 0, 1, "magma"),
    ]:
        n_arms = len(parsed)
        fig, axes = plt.subplots(1, n_arms, figsize=(4 * n_arms, 4), squeeze=False)
        for ax, (name, p) in zip(axes[0], parsed.items()):
            evs = p["evals"]
            if not evs:
                ax.set_title(p["label"] + " (no eval)")
                continue
            hops = sorted({h for e in evs for h in e["hops"]})
            mat = np.full((len(hops), len(evs)), np.nan)
            for j, e in enumerate(evs):
                for i, h in enumerate(hops):
                    if h in e["hops"]:
                        mat[i, j] = e["hops"][h].get(key, np.nan)
            im = ax.imshow(
                mat,
                aspect="auto",
                origin="lower",
                vmin=vmin,
                vmax=vmax,
                cmap=cmap,
                interpolation="nearest",
            )
            ax.set_yticks(range(len(hops)))
            ax.set_yticklabels([str(h) for h in hops], fontsize=7)
            # x ticks: a few epochs
            xt = np.linspace(0, max(len(evs) - 1, 0), num=min(6, len(evs)), dtype=int)
            ax.set_xticks(xt)
            ax.set_xticklabels([str(evs[j]["epoch"]) for j in xt], fontsize=7, rotation=45)
            ax.set_xlabel("epoch")
            ax.set_ylabel("hop")
            ax.set_title(p["label"], fontsize=9)
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        fig.suptitle(f"L15 diagnostic — {title} over training")
        fig.tight_layout()
        fig.savefig(OUT_DIR / f"{fname}.png", dpi=140)
        plt.close(fig)

    # E: time to frontier thresholds
    fig, ax = plt.subplots(figsize=(9, 4))
    summary = {}
    for name, p in parsed.items():
        summary[name] = {}
        for thr, ls in [(0.85, "--"), (0.95, "-")]:
            tt = time_to_threshold(p["evals"], "frontier", thr)
            summary[name][f"frontier>={thr}"] = tt
            xs = [h for h, e in tt.items() if e is not None]
            ys = [tt[h] for h in xs]
            if xs:
                ax.plot(
                    xs,
                    ys,
                    ls,
                    color=p["color"],
                    marker="o",
                    ms=3,
                    label=f"{p['label']} thr={thr}",
                )
        tt_ce = time_to_threshold(p["evals"], "ce_score", 0.5)
        summary[name]["ce_score>=0.5"] = tt_ce
    ax.set_xlabel("hop")
    ax.set_ylabel("first epoch reaching threshold")
    ax.set_title("L15 diagnostic — time-to-threshold (frontier)")
    ax.legend(fontsize=7, ncol=2)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(OUT_DIR / "E_time_to_threshold.png", dpi=140)
    plt.close(fig)

    # F: latest bars frontier vs ce
    n_arms = len(parsed)
    fig, axes = plt.subplots(1, n_arms, figsize=(4 * n_arms, 3.5), squeeze=False)
    for ax, (name, p) in zip(axes[0], parsed.items()):
        if not p["evals"]:
            continue
        last = p["evals"][-1]["hops"]
        hops = sorted(last)
        fr = [last[h]["frontier"] for h in hops]
        ce = [last[h]["ce_score"] for h in hops]
        x = np.arange(len(hops))
        ax.bar(x - 0.2, fr, 0.4, label="frontier", color="#1b7837")
        ax.bar(x + 0.2, ce, 0.4, label="ce_score", color="#d95f02")
        ax.axhline(0.95, color="gray", ls="--", lw=0.8)
        ax.axhline(0.5, color="#d95f02", ls=":", lw=0.8)
        ax.set_xticks(x)
        ax.set_xticklabels([str(h) for h in hops], fontsize=7)
        ax.set_ylim(0, 1.05)
        ax.set_title(
            f"{p['label']}\nep{p['evals'][-1]['epoch']} stage={p['evals'][-1]['stage']}",
            fontsize=8,
        )
        ax.legend(fontsize=7)
        ax.grid(True, axis="y", alpha=0.3)
    fig.suptitle("L15 diagnostic — latest per-hop frontier vs ce_score")
    fig.tight_layout()
    fig.savefig(OUT_DIR / "F_latest_bars.png", dpi=140)
    plt.close(fig)

    (OUT_DIR / "summary.json").write_text(json.dumps(summary, indent=2, default=str))
    print(f"wrote plots -> {OUT_DIR}")
    for name, p in parsed.items():
        last_ep = p["train"][-1][0] if p["train"] else "?"
        last_st = p["train"][-1][1] if p["train"] else "?"
        n_prom = len(p["promotes"])
        print(f"  {p['label']}: epoch={last_ep} stage={last_st} promotes={n_prom} evals={len(p['evals'])}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--watch", type=int, default=0, help="replot every N seconds")
    args = ap.parse_args()
    if args.watch <= 0:
        plot_all()
        return
    while True:
        try:
            plot_all()
        except Exception as e:
            print("plot error:", e)
        time.sleep(args.watch)


if __name__ == "__main__":
    main()