#!/usr/bin/env python3
"""Plot SAC training progress from live campaign logs.

Pure-stdlib SVG output (matplotlib not available on this box).
Generates:
  docs/graph_eval_winrates.svg  - eval win rate per opponent (+ rolling-20 mean, v1 overlay)
  docs/graph_losses.svg         - critic/actor/alpha loss curves from training_metrics.jsonl

Usage:
  python3 tools/plot_progress.py [campaign_log] [metrics_jsonl] [v1_log] [outdir]
All args optional; defaults relative to the SAC_LSTM_Bot/ root (parent of tools/).
  python3 tools/plot_progress.py --selftest   # tiny built-in sanity check
"""
import json
import math
import re
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
EVAL_RE = re.compile(r">>> \[eval\] win rate: (\d+)/(\d+) \(([\d.]+)%\) vs (\S+)")
ROLL_WINDOW = 20
COLORS = {"Corners": "#1f77b4", "Crazy": "#2ca02c", "Target": "#d62728"}
V1_COLOR = "#888888"
W, M_L, M_R, M_T, M_B = 1000, 70, 20, 40, 45


def parse_eval_series(path):
    """Return {opponent: [win% per eval, in file order]}."""
    series = {}
    if not path.is_file():
        print(f"[skip] eval log not found: {path}")
        return series
    for line in path.read_text(errors="replace").splitlines():
        m = EVAL_RE.search(line)
        if not m:
            continue
        try:
            pct = float(m.group(3))
        except ValueError:
            continue
        series.setdefault(m.group(4), []).append(pct)
    return series


def parse_metrics(path):
    """Return list of metric dicts, skipping malformed lines."""
    rows = []
    if not path.is_file():
        print(f"[skip] metrics file not found: {path}")
        return rows
    for line in path.read_text(errors="replace").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            rows.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    return rows


def rolling(vals, w=ROLL_WINDOW):
    out, s = [], 0.0
    for i, v in enumerate(vals):
        s += v
        if i >= w:
            s -= vals[i - w]
        out.append(s / min(i + 1, w))
    return out


# ---------- tiny SVG helpers ----------

def esc(s):
    return str(s).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def svg_open(w, h, title):
    return (f'<svg xmlns="http://www.w3.org/2000/svg" width="{w}" height="{h}" '
            f'viewBox="0 0 {w} {h}" font-family="sans-serif">\n'
            f'<rect width="{w}" height="{h}" fill="white"/>\n'
            f'<text x="{w // 2}" y="22" text-anchor="middle" font-size="17" font-weight="bold">'
            f"{esc(title)}</text>\n")


def polyline(pts, color, width=1.5, dash=None, opacity=1.0):
    if len(pts) < 2:
        return ""
    d = f' stroke-dasharray="{dash}"' if dash else ""
    p = " ".join(f"{x:.1f},{y:.1f}" for x, y in pts)
    return (f'<polyline points="{p}" fill="none" stroke="{color}" '
            f'stroke-width="{width}" opacity="{opacity}"{d}/>\n')


def axis(x0, y0, x1, y1, xt, yt, xlabel, ylabel, ylog=False):
    """Draw axes + ticks + labels. xt/yt are (value, px) tick lists."""
    s = (f'<line x1="{x0}" y1="{y0}" x2="{x1}" y2="{y0}" stroke="black"/>\n'
         f'<line x1="{x0}" y1="{y0}" x2="{x0}" y2="{y1}" stroke="black"/>\n')
    for v, px in xt:
        s += (f'<line x1="{px:.1f}" y1="{y0}" x2="{px:.1f}" y2="{y0 + 4}" stroke="black"/>\n'
              f'<text x="{px:.1f}" y="{y0 + 17}" text-anchor="middle" font-size="11">'
              f"{esc(v)}</text>\n")
    for v, py in yt:
        s += (f'<line x1="{x0 - 4}" y1="{py:.1f}" x2="{x0}" y2="{py:.1f}" stroke="black"/>\n'
              f'<text x="{x0 - 7}" y="{py + 4:.1f}" text-anchor="end" font-size="11">'
              f"{esc(v)}</text>\n")
    s += (f'<text x="{(x0 + x1) // 2}" y="{y0 + 33}" text-anchor="middle" font-size="12">'
          f"{esc(xlabel)}</text>\n"
          f'<text x="16" y="{(y0 + y1) // 2}" text-anchor="middle" font-size="12" '
          f'transform="rotate(-90 16 {(y0 + y1) // 2})">{esc(ylabel)}</text>\n')
    return s


def ticks_linear(vmin, vmax, p0, p1, n=5, fmt="{:.0f}"):
    return [(fmt.format(vmin + (vmax - vmin) * i / (n - 1)),
             p0 + (p1 - p0) * i / (n - 1)) for i in range(n)]


def ticks_log(vmin, vmax, p0, p1, n=5):
    vals = [10 ** (vmin + (vmax - vmin) * i / (n - 1)) for i in range(n)]
    return [("{:.3g}".format(v), p0 + (p1 - p0) * i / (n - 1)) for i, v in enumerate(vals)]


def legend(items, x, y):
    """items: [(color, dash, label)]"""
    s = ""
    for i, (color, dash, label) in enumerate(items):
        yy = y + i * 18
        d = f' stroke-dasharray="{dash}"' if dash else ""
        s += (f'<line x1="{x}" y1="{yy}" x2="{x + 28}" y2="{yy}" stroke="{color}" '
              f'stroke-width="3"{d}/>\n'
              f'<text x="{x + 34}" y="{yy + 4}" font-size="12">{esc(label)}</text>\n')
    return s


def map_fn(p0, p1, vmin, vmax, log=False):
    def f(v):
        t = (math.log10(v) - vmin) / (vmax - vmin) if log else (v - vmin) / (vmax - vmin)
        return p1 - max(0.0, min(1.0, t)) * (p1 - p0)
    return f


# ---------- graph 1: eval win rates ----------

def graph_eval(series_v2, series_v1, out):
    if not series_v2 and not series_v1:
        print("[skip] no eval data at all -> no win-rate graph")
        return 0
    H = 560
    x0, y0, x1, y1 = M_L, H - M_B, W - M_R, M_T
    xmax = max(max((len(v) for v in s.values()), default=1) for s in (series_v2, series_v1))
    xmax = max(xmax, 2)
    xm = map_fn(x0, x1, 1, xmax)
    ym = map_fn(y0, y1, 0, 100)
    s = svg_open(W, H, "SAC-LSTM campaign: eval win rate vs opponents (v2 live, v1 archived)")
    s += axis(x0, y0, x1, y1, ticks_linear(1, xmax, x0, x1),
              ticks_linear(0, 100, y0, y1), f"eval number (per opponent)",
              "win rate (%)")
    counts = []
    items = []
    for name in ("Corners", "Crazy", "Target"):
        vals = series_v2.get(name, [])
        if not vals:
            continue
        c = COLORS[name]
        s += polyline([(xm(i + 1), ym(v)) for i, v in enumerate(vals)], c, 1, opacity=0.30)
        s += polyline([(xm(i + 1), ym(v)) for i, v in enumerate(rolling(vals))], c, 3)
        counts.append(f"{name}: {len(vals)} evals")
        items.append((c, None, f"{name} (thick = rolling-{ROLL_WINDOW} mean)"))
    v1 = series_v1.get("Corners", [])
    if v1:
        s += polyline([(xm(i + 1), ym(v)) for i, v in enumerate(v1)], V1_COLOR, 1, dash="6,4")
        s += polyline([(xm(i + 1), ym(v)) for i, v in enumerate(rolling(v1))], V1_COLOR, 2, dash="6,4")
        counts.append(f"v1 Corners (old run): {len(v1)} evals")
        items.append((V1_COLOR, "6,4", "v1 Corners (old run)"))
    s += legend(items, x1 - 260, y1 + 14)
    s += (f'<text x="{x0 + 8}" y="{y1 + 14}" font-size="11" fill="#555">'
          f'{esc("; ".join(counts))}</text>\n</svg>\n')
    out.write_text(s)
    return 1


# ---------- graph 2: loss curves ----------

def graph_losses(rows, out):
    if not rows:
        print("[skip] no metric rows -> no loss graph")
        return 0

    def col(key, log=False):
        vals = []
        for r in rows:
            try:
                v = abs(float(r[key]))  # abs: log panels plot magnitude
            except (KeyError, TypeError, ValueError):
                continue
            if log and v <= 0:
                continue
            vals.append(v)
        return vals

    critic = col("critic_loss", log=True)
    actor = col("actor_loss", log=True)
    alpha = col("alpha")
    if not (critic or actor or alpha):
        print("[skip] metric rows lack critic_loss/actor_loss/alpha -> no loss graph")
        return 0

    PH, GAP = 210, 55
    H = M_T + 3 * PH + 2 * GAP + M_B
    x0, x1 = M_L, W - M_R
    n = len(rows)
    s = svg_open(W, H, "SAC-LSTM training losses (x = metric line number)")

    def panel(top, vals, title, log, ylab, fixed_range=None):
        nonlocal s
        y0, y1 = top + PH, top
        if not vals:
            s += (f'<text x="{x0 + 10}" y="{top + 20}" font-size="12" fill="#a00">'
                  f"{esc(title)}: no valid points</text>\n")
            return
        if log:
            lo = math.floor(math.log10(min(vals)))
            hi = math.ceil(math.log10(max(vals)))
            if lo == hi:
                hi = lo + 1
            yt = ticks_log(lo, hi, y0, y1)
        else:
            lo, hi = fixed_range or (min(vals), max(vals))
            if lo == hi:
                hi = lo + 1
            yt = ticks_linear(lo, hi, y0, y1, n=5, fmt="{:.3g}")
        s += (f'<text x="{x0 + 6}" y="{top + 16}" font-size="13" font-weight="bold">'
              f"{esc(title)}</text>\n")
        s += axis(x0, y0, x1, y1, ticks_linear(1, n, x0, x1, fmt="{:.0f}"), yt,
                  "", ylab)
        ym = map_fn(y0, y1, lo, hi, log=log)
        pts = []
        for i, r in enumerate(rows):
            try:
                v = abs(float(r[title_key(title)]))
            except (KeyError, TypeError, ValueError):
                continue
            if log and v <= 0:
                continue
            pts.append((x0 + (x1 - x0) * i / max(n - 1, 1), ym(v)))
        s += polyline(pts, "#1f77b4", 1.6)

    def title_key(title):
        return {"critic_loss (log scale)": "critic_loss",
                "actor_loss |abs| (log scale)": "actor_loss",
                "alpha (temperature)": "alpha"}[title]

    panel(M_T, critic, "critic_loss (log scale)", True, "critic_loss")
    panel(M_T + PH + GAP, actor, "actor_loss |abs| (log scale)", True, "|actor_loss|")
    panel(M_T + 2 * (PH + GAP), alpha, "alpha (temperature)", False, "alpha",
          fixed_range=(0, max(1.0, max(alpha))))
    s += (f'<text x="{x0 + (x1 - x0) // 2}" y="{H - 8}" text-anchor="middle" '
          f'font-size="12">metric line number ({n} rows)</text>\n</svg>\n')
    out.write_text(s)
    return 1


# ---------- selftest ----------

def selftest():
    with tempfile.TemporaryDirectory() as td:
        td = Path(td)
        (td / "log").write_text(
            ">>> [eval] win rate: 3/10 (30%) vs Corners\n"
            "garbage line\n"
            ">>> [eval] win rate: 7/10 (70%) vs Crazy\n"
            ">>> [eval] win rate: broken\n"
            ">>> [eval] win rate: 5/10 (50%) vs Corners\n")
        ser = parse_eval_series(td / "log")
        assert ser == {"Corners": [30.0, 50.0], "Crazy": [70.0]}, ser
        assert rolling([10] * 25, 20)[-1] == 10.0
        assert rolling([1, 2, 3], 20) == [1.0, 1.5, 2.0]
        (td / "m.jsonl").write_text(
            '{"critic_loss": 10, "actor_loss": -2, "alpha": 0.5}\n'
            "not json\n"
            '{"critic_loss": 100, "actor_loss": -4, "alpha": 0.25}\n')
        rows = parse_metrics(td / "m.jsonl")
        assert len(rows) == 2 and rows[1]["critic_loss"] == 100
        ok = graph_losses(rows, td / "g2.svg") and graph_eval(ser, {"Corners": [0, 10]}, td / "g1.svg")
        assert ok and (td / "g1.svg").stat().st_size > 500
        assert graph_eval({}, {}, td / "none.svg") == 0  # missing data handled
    print("selftest OK")


def main():
    if "--selftest" in sys.argv:
        selftest()
        return
    args = [a for a in sys.argv[1:] if not a.startswith("-")]
    campaign = Path(args[0]) if len(args) > 0 else ROOT / "campaign_v2_stdout.log"
    metrics = Path(args[1]) if len(args) > 1 else ROOT / "training_metrics.jsonl"
    v1log = Path(args[2]) if len(args) > 2 else ROOT / "weights_v1_archive" / "campaign_stdout.log"
    outdir = Path(args[3]) if len(args) > 3 else ROOT / "docs"
    outdir.mkdir(parents=True, exist_ok=True)

    made = 0
    series_v2 = parse_eval_series(campaign)
    print(f"[info] v2 evals parsed: " +
          ", ".join(f"{k}={len(v)}" for k, v in sorted(series_v2.items())) or "(none)")
    series_v1 = parse_eval_series(v1log)
    print(f"[info] v1 evals parsed: " +
          ", ".join(f"{k}={len(v)}" for k, v in sorted(series_v1.items())) or "(none)")
    made += graph_eval(series_v2, series_v1, outdir / "graph_eval_winrates.svg")

    rows = parse_metrics(metrics)
    print(f"[info] metric rows parsed: {len(rows)}")
    made += graph_losses(rows, outdir / "graph_losses.svg")

    if made:
        print(f"[done] {made} graph(s) written to {outdir}")
    else:
        print("[error] nothing plotted - check paths above")
        sys.exit(1)


if __name__ == "__main__":
    main()