#!/usr/bin/env python3
# =============================================================================
#  Project   : localvm-research
#  File      : tools/make_pub_figures.py
#  Purpose   : Generate publication SVG figures from committed results JSON
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64)
#  License   : All rights reserved (research code)
# =============================================================================
"""Regenerates docs/publications/figures/*.svg from results/ (reproducible).

Palette: validated categorical slots (see dataviz reference; CVD-checked):
blue #2a78d6, orange #eb6834, aqua #1baf7a. Ink/grid tokens match the site.
"""

from __future__ import annotations

import glob
import json
import math
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "docs" / "publications" / "figures"
OUT.mkdir(parents=True, exist_ok=True)

INK, INK2, MUTED, GRID, BASE = "#0b0b0b", "#52514e", "#898781", "#e1e0d9", "#c3c2b7"
BLUE, ORANGE, AQUA = "#2a78d6", "#eb6834", "#1baf7a"
FONT = "font-family='system-ui, -apple-system, sans-serif'"
MONO = "font-family='ui-monospace, Menlo, monospace'"


def latest(pattern: str) -> dict:
    return json.load(open(sorted(glob.glob(str(ROOT / pattern)))[-1]))


def svg_open(w: int, h: int) -> list[str]:
    return [f"<svg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 {w} {h}' "
            f"width='{w}' height='{h}'>",
            f"<rect width='{w}' height='{h}' fill='#fcfcfb'/>"]


def text(x, y, s, size=12, fill=INK2, anchor="start", weight="400", mono=False):
    f = MONO if mono else FONT
    return (f"<text x='{x:.1f}' y='{y:.1f}' {f} font-size='{size}' fill='{fill}' "
            f"text-anchor='{anchor}' font-weight='{weight}'>{s}</text>")


# ---------------------------------------------------------------- F1: SSD
def fig1_ssd() -> None:
    d = latest("results/expH_ssd_feasibility/*/results.json")
    cells = [c for c in d["cells"] if c["nocache"] and not c["gpu_load"]
             and c["pattern"] == "random" and c["block_bytes"] <= 1 << 20]
    series: dict[int, list] = {}
    for c in cells:
        series.setdefault(c["threads"], []).append((c["block_bytes"], c["mb_per_s_mean"]))
    W, H, ML, MR, MT, MB_ = 640, 330, 64, 84, 20, 46
    iw, ih = W - ML - MR, H - MT - MB_
    xs = lambda b: ML + (math.log2(b) - 12) / 8 * iw
    ys = lambda v: MT + ih - (math.log10(v) - math.log10(50)) / (math.log10(20000) - math.log10(50)) * ih
    out = svg_open(W, H)
    for v in (100, 1000, 10000):
        out.append(f"<line x1='{ML}' y1='{ys(v):.1f}' x2='{ML+iw}' y2='{ys(v):.1f}' stroke='{GRID}'/>")
        out.append(text(ML - 8, ys(v) + 4, f"{v//1000} GB/s" if v >= 1000 else f"{v} MB/s",
                        11, MUTED, "end", mono=True))
    for b in (4096, 16384, 65536, 262144, 1048576):
        lbl = "1 MiB" if b >= 1048576 else f"{b//1024} KiB"
        out.append(text(xs(b), MT + ih + 18, lbl, 11, MUTED, "middle", mono=True))
    out.append(f"<line x1='{ML}' y1='{MT+ih}' x2='{ML+iw}' y2='{MT+ih}' stroke='{BASE}'/>")
    colors = {1: BLUE, 4: ORANGE, 8: AQUA}
    for qd, pts in sorted(series.items()):
        pts.sort()
        c = colors[qd]
        pd = " ".join(f"{'M' if i == 0 else 'L'}{xs(b):.1f},{ys(v):.1f}" for i, (b, v) in enumerate(pts))
        out.append(f"<path d='{pd}' fill='none' stroke='{c}' stroke-width='2'/>")
        for b, v in pts:
            out.append(f"<circle cx='{xs(b):.1f}' cy='{ys(v):.1f}' r='4' fill='{c}' "
                       f"stroke='#fcfcfb' stroke-width='1.5'/>")
        out.append(text(xs(pts[-1][0]) + 10, ys(pts[-1][1]) + 4, f"QD {qd}", 12, c, weight="600"))
    out.append(text(ML + iw / 2, H - 8, "random read block size (F_NOCACHE cold, log–log)", 12, INK2, "middle"))
    out.append("</svg>")
    (OUT / "f1_ssd_envelope.svg").write_text("\n".join(out))


# ---------------------------------------------------------------- F2: expG scale
def fig2_escalation() -> None:
    runs = sorted(glob.glob(str(ROOT / "results/expG_decision_stability/*/results.json")))
    data = {}
    for r in runs:
        d = json.load(open(r))
        if d["n_trajectories"] < 40:
            continue
        scale = "8B" if "8B" in d["config"]["model"] else "1.7B"
        data[scale] = {int(b): s.get("escalation_frac_for_99pct")
                       for b, s in d["results_by_bits"].items() if int(b) >= 3}
    W, H, ML, MR, MT, MB_ = 560, 320, 64, 80, 20, 46
    iw, ih = W - ML - MR, H - MT - MB_
    xs = {3: ML + iw * 0.12, 4: ML + iw * 0.5, 8: ML + iw * 0.88}
    ys = lambda v: MT + ih - v / 0.7 * ih
    out = svg_open(W, H)
    for v in (0.2, 0.4, 0.6):
        out.append(f"<line x1='{ML}' y1='{ys(v):.1f}' x2='{ML+iw}' y2='{ys(v):.1f}' stroke='{GRID}'/>")
        out.append(text(ML - 8, ys(v) + 4, f"{int(v*100)}%", 11, MUTED, "end", mono=True))
    out.append(f"<line x1='{ML}' y1='{MT+ih}' x2='{ML+iw}' y2='{MT+ih}' stroke='{BASE}'/>")
    for b in (3, 4, 8):
        out.append(text(xs[b], MT + ih + 18, f"{b}-bit base", 12, MUTED, "middle"))
    for scale, color in (("1.7B", BLUE), ("8B", ORANGE)):
        pts = [(b, v) for b, v in sorted(data.get(scale, {}).items()) if v]
        pd = " ".join(f"{'M' if i == 0 else 'L'}{xs[b]:.1f},{ys(v):.1f}" for i, (b, v) in enumerate(pts))
        out.append(f"<path d='{pd}' fill='none' stroke='{color}' stroke-width='2'/>")
        for b, v in pts:
            out.append(f"<circle cx='{xs[b]:.1f}' cy='{ys(v):.1f}' r='4.5' fill='{color}' "
                       f"stroke='#fcfcfb' stroke-width='1.5'/>")
            dy = 20 if (scale == "8B" and b == 8) else -10  # avoid series-label collision
            out.append(text(xs[b], ys(v) + dy, f"{v*100:.0f}%", 11, color, "middle", "600", mono=True))
        out.append(text(xs[pts[-1][0]] + 12, ys(pts[-1][1]) + 4, f"Qwen3-{scale}", 12, color, weight="600"))
    out.append(text(ML + iw / 2, H - 8, "tokens needing escalation for 99% greedy agreement", 12, INK2, "middle"))
    out.append("</svg>")
    (OUT / "f2_escalation_scale.svg").write_text("\n".join(out))


# ---------------------------------------------------------------- F3: continuum
def fig3_continuum() -> None:
    runs = [json.load(open(p)) for p in sorted(glob.glob(str(ROOT / "results/candidate_01/*/results.json")))]
    d = [r for r in runs if "quality_bf16_judge" in r][-1]
    q = d["quality_bf16_judge"]
    pts = [("pure q4", 0.0, q["pure_q4"]["mean_logprob_bf16"]),
           ("τ=1.0", 150, q["margin_tau1.0"]["mean_logprob_bf16"]),
           ("τ=2.0", 211, q["margin_tau2.0"]["mean_logprob_bf16"]),
           ("verify-all", 237, q["verify-all_tau2.0"]["mean_logprob_bf16"])]
    ceiling = q["pure_q8"]["mean_logprob_bf16"]
    W, H, ML, MR, MT, MB_ = 620, 330, 74, 40, 26, 46
    iw, ih = W - ML - MR, H - MT - MB_
    xs = lambda v: ML + v / 260 * iw
    ys = lambda v: MT + ih - (v + 0.42) / 0.34 * ih
    out = svg_open(W, H)
    for v in (-0.4, -0.3, -0.2, -0.1):
        out.append(f"<line x1='{ML}' y1='{ys(v):.1f}' x2='{ML+iw}' y2='{ys(v):.1f}' stroke='{GRID}'/>")
        out.append(text(ML - 8, ys(v) + 4, f"{v:.1f}", 11, MUTED, "end", mono=True))
    for v in (0, 100, 200):
        out.append(text(xs(v), MT + ih + 18, f"{v} MB", 11, MUTED, "middle", mono=True))
    out.append(f"<line x1='{ML}' y1='{MT+ih}' x2='{ML+iw}' y2='{MT+ih}' stroke='{BASE}'/>")
    out.append(f"<line x1='{ML}' y1='{ys(ceiling):.1f}' x2='{ML+iw}' y2='{ys(ceiling):.1f}' "
               f"stroke='{MUTED}' stroke-dasharray='5,4'/>")
    out.append(text(ML + 4, ys(ceiling) + 15, "resident-q8 ceiling (2.15 GB in RAM)", 11, MUTED))
    pd = " ".join(f"{'M' if i == 0 else 'L'}{xs(x):.1f},{ys(y):.1f}" for i, (_, x, y) in enumerate(pts))
    out.append(f"<path d='{pd}' fill='none' stroke='{BLUE}' stroke-width='2'/>")
    for name, x, y in pts:
        out.append(f"<circle cx='{xs(x):.1f}' cy='{ys(y):.1f}' r='5' fill='{BLUE}' "
                   f"stroke='#fcfcfb' stroke-width='2'/>")
        if name == "verify-all":
            out.append(text(xs(x) - 11, ys(y) + 4, name, 12, INK, "end", "600"))
        else:
            out.append(text(xs(x) + 9, ys(y) + (14 if name == "pure q4" else -9), name, 12, INK, weight="600"))
    out.append(text(ML + iw / 2, H - 8, "verification bytes streamed per token (Qwen3-1.7B, W=32)", 12, INK2, "middle"))
    out.append(text(16, MT + ih / 2, "bf16-judge mean logprob →", 12, INK2, "middle")
               .replace(">", f" transform='rotate(-90 16 {MT+ih/2})'>", 1))
    out.append("</svg>")
    (OUT / "f3_quality_bytes_continuum.svg").write_text("\n".join(out))


# ---------------------------------------------------------------- F4: 32B
def fig4_scale32b() -> None:
    d = latest("results/candidate_01_scale32b/20260812T055506Z/results.json")
    q = d["quality_8b_judge"]
    rows = [("pure q4 — only config that fits (17.5 GB)", q["pure_q4"]["mean_logprob_8b_judge"]),
            ("margin τ=2.0 — q8 streamed on low margins", q["margin_tau2.0"]["mean_logprob_8b_judge"]),
            ("verify-all — q8 streamed every window", q["verify-all_tau2.0"]["mean_logprob_8b_judge"])]
    W, H, ML, MR, MT = 640, 210, 26, 90, 24
    row_h = 52
    x0, x1 = -0.80, -0.40
    xs = lambda v: ML + (v - x0) / (x1 - x0) * (W - ML - MR)
    out = svg_open(W, H)
    for i, (name, v) in enumerate(rows):
        y = MT + i * row_h + 20
        color = BLUE if i else MUTED
        out.append(text(ML, y - 8, name, 12.5, INK2))
        out.append(f"<line x1='{xs(x0)}' y1='{y+9}' x2='{xs(v):.1f}' y2='{y+9}' "
                   f"stroke='{color}' stroke-width='9' stroke-linecap='round'/>")
        out.append(text(xs(v) + 10, y + 13, f"{v:.3f}", 12, INK, weight="650", mono=True))
    out.append(text(ML, H - 10, "8B-judge mean logprob (higher is better) · Qwen3-32B on a 48 GB Mac — "
                    "resident q8 is impossible", 11.5, MUTED))
    out.append("</svg>")
    (OUT / "f4_32b_outofcore.svg").write_text("\n".join(out))


if __name__ == "__main__":
    fig1_ssd(); fig2_escalation(); fig3_continuum(); fig4_scale32b()
    print("figures written to", OUT)