#!/usr/bin/env python3
"""plot_results.py - charts and tables for docs/raw12-pipeline-measurements.md.

    ~/.venvs/fgc-plots/bin/python plot_results.py <out_dir> <label>=<results_dir> [...]

Each results folder is what run_raw12_camera.sh and run_raw12_encode.sh wrote
for one tower. Writes PNG charts to <out_dir> and prints Markdown tables
(stdout) for the report. Needs matplotlib and numpy.
"""
import json
import os
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

OUT = sys.argv[1]
TOWERS = dict(a.split("=", 1) for a in sys.argv[2:])
os.makedirs(OUT, exist_ok=True)
plt.rcParams.update({"figure.dpi": 110, "font.size": 9, "axes.grid": True, "grid.alpha": 0.3})
COL = {"BayerRG12p": "#d62728", "RGB8": "#1f77b4", "BayerRG8": "#2ca02c", "BayerRG12": "#9467bd",
       "BayerRG10p": "#8c564b"}
KIND_LABEL = {"rgb8": "camera RGB8 (today)", "c8_bil": "raw 8 → colour, bilinear",
              "c8_ea": "raw 8 → colour, edge-aware", "c12_bil": "raw 12 → colour, bilinear",
              "c12_ea": "raw 12 → colour, edge-aware", "raw12": "raw 12 mosaic, lossless",
              "raw12_planes": "raw 12 planes", "raw8": "raw 8 mosaic, lossless"}


def rows(path):
    if not os.path.exists(path):
        return []
    out = []
    for line in open(path):
        line = line.strip()
        if line.startswith("{"):
            try:
                out.append(json.loads(line))
            except json.JSONDecodeError:
                pass
    return out


def save(fig, name):
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, name))
    plt.close(fig)
    print(f"<!-- wrote {name} -->")


def md(header, body):
    print("| " + " | ".join(header) + " |")
    print("|" + "---|" * len(header))
    for r in body:
        print("| " + " | ".join(str(x) for x in r) + " |")
    print()


# ---- frame rate and streaming ---------------------------------------------------
print("## Frame rate and streaming at 250 MB/s\n")
body = []
for t, d in TOWERS.items():
    mf = {r["fmt"]: r for r in rows(f"{d}/maxfps.jsonl") if "max_fps" in r}
    for r in rows(f"{d}/stream.jsonl"):
        if r.get("cmd") != "stream":
            continue
        body.append([t, r["fmt"], f"{mf.get(r['fmt'], {}).get('payload', r['payload']) / 1e6:.1f} MB",
                     f"{mf.get(r['fmt'], {}).get('max_fps', r['fps_set']):.2f}", f"{r['fps']:.2f}",
                     f"{r['complete']}/{r['frames']}", r["incomplete"], r["missing_ids"],
                     f"{r['max_gap_s']:.2f} s"])
    for r in rows(f"{d}/stream_under_load.jsonl"):
        if r.get("cmd") == "stream":
            body.append([t, r["fmt"] + " **while saving 12-bit**", f"{r['payload'] / 1e6:.1f} MB",
                         f"{r['fps_set']:.2f}", f"{r['fps']:.2f}", f"{r['complete']}/{r['frames']}",
                         r["incomplete"], r["missing_ids"], f"{r['max_gap_s']:.2f} s"])
md(["tower", "format", "frame", "max fps", "measured fps", "complete", "incomplete", "lost",
    "longest gap"], body)

fig, ax = plt.subplots(figsize=(6.5, 3.2))
fmts = ["RGB8", "BayerRG12", "BayerRG12p", "BayerRG10p", "BayerRG8"]
width = 0.8 / max(1, len(TOWERS))
for i, (t, d) in enumerate(TOWERS.items()):
    mf = {r["fmt"]: r["max_fps"] for r in rows(f"{d}/maxfps.jsonl") if "max_fps" in r}
    xs = np.arange(len(fmts)) + i * width
    vals = [mf.get(f, 0) for f in fmts]
    bars = ax.bar(xs, vals, width, label=f"tower {t}")
    ax.bar_label(bars, fmt="%.1f", fontsize=8)
ax.set_xticks(np.arange(len(fmts)) + width * (len(TOWERS) - 1) / 2, fmts)
ax.set_ylabel("max frames/s at 250 MB/s")
ax.set_title("Full-resolution frame rate per pixel format (250 MB/s link cap)")
ax.legend()
save(fig, "fps_per_format.png")


# ---- auto-exposure ---------------------------------------------------------------
def eff(e, g):
    return e * 10 ** (g / 20)


def score_ae(r):
    """Settle time and how candidate keep-rules do. "Good" is judged on the
    image: brightness within 10% of where this run finally settles (median of
    its last 2 s). Exposure x gain is not the reference: with exposure at its
    cap, brightness is what the frame shows, and clipping makes the two differ."""
    s = [x for x in r["series"] if x[3] is not None]
    if len(s) < 10:
        return None
    t_end = s[-1][0]
    final = float(np.median([x[3] for x in s if x[0] >= t_end - 2.0]))
    good = [abs(x[3] / final - 1) <= 0.10 for x in s]
    stay = next((i for i in range(len(good)) if all(good[i:])), None)
    target = (r.get("target") or 50) / 100 * 255
    tol = (r.get("tolerance") or 5) / 100 * 255
    same = lambda a, b: abs(a[1] - b[1]) <= 0.005 * a[1] and abs(a[2] - b[2]) <= 0.05
    in_band = lambda x: abs(x[3] - target) <= tol
    rules = {
        "exposure & gain unchanged 3 frames":
            lambda i: i >= 2 and same(s[i], s[i - 1]) and same(s[i - 1], s[i - 2]),
        "brightness in the camera's target band (target ± tolerance), 2 frames":
            lambda i: i >= 1 and in_band(s[i]) and in_band(s[i - 1]),
        "target band 2 frames AND exposure & gain unchanged 2 frames":
            lambda i: i >= 1 and in_band(s[i]) and in_band(s[i - 1]) and same(s[i], s[i - 1]),
        "target band 3 frames":
            lambda i: i >= 2 and in_band(s[i]) and in_band(s[i - 1]) and in_band(s[i - 2]),
        "target band AND brightness steady (≤2% change) over 3 frames":
            lambda i: i >= 2 and all(in_band(s[j]) for j in (i, i - 1, i - 2))
            and abs(s[i][3] / s[i - 1][3] - 1) <= 0.02 and abs(s[i - 1][3] / s[i - 2][3] - 1) <= 0.02,
    }
    res = {}
    for name, rule in rules.items():
        i = next((j for j in range(len(s)) if rule(j)), None)
        res[name] = (None, None) if i is None else (s[i][0], good[i])
    return dict(stay_s=s[stay][0] if stay is not None else None, stay_frame=stay, final=final,
                rules=res)


print("## Auto-exposure step response\n")
for t, d in TOWERS.items():
    ae = [r for r in rows(f"{d}/ae.jsonl") if r.get("cmd") == "ae"]
    if not ae:
        continue
    fig, axes = plt.subplots(1, 2, figsize=(9, 3.2), sharey=True)
    for ax, direction in zip(axes, ("dark", "bright")):
        for r in ae:
            if r["direction"] != direction or not r["series"]:
                continue
            xs = [x[0] for x in r["series"]]
            ys = [x[3] for x in r["series"]]
            ax.plot(xs, ys, color=COL.get(r["fmt"], "k"), alpha=0.8, lw=1.2,
                    label=f"{r['fmt']} {r['fps_set']:.1f} fps")
        if r.get("base_brightness"):
            ax.axhline(r["base_brightness"], color="grey", ls="--", lw=0.8)
        ax.set_xlim(0, 7)
        ax.set_title(f"{direction} step (×{'1/16' if direction == 'dark' else '16 max'})")
        ax.set_xlabel("seconds after auto-exposure resumes")
        h_, l_ = ax.get_legend_handles_labels()
        uniq = dict(zip(l_, h_))
        ax.legend(uniq.values(), uniq.keys(), fontsize=7)
    axes[0].set_ylabel("frame brightness (0–255)")
    fig.suptitle(f"Tower {t}: brightness recovering after a forced exposure step (dashed: settled level)")
    save(fig, f"ae_brightness_{t}.png")

    body, rule_rows = [], {}
    for r in ae:
        sc = score_ae(r)
        if not sc:
            continue
        body.append([t, r["fmt"], f"{r['fps_set']:.2f}", r["direction"], f"×{r['forced_factor']}",
                     f"{sc['stay_s']:.2f} s ({sc['stay_frame']} frames)" if sc["stay_s"] is not None else "–",
                     f"{sc['final']:.0f}"])
        for name, (ts, ok) in sc["rules"].items():
            rr = rule_rows.setdefault(name, {"fired": 0, "correct": 0, "early": 0, "never": 0, "t": []})
            if ts is None:
                rr["never"] += 1
            else:
                rr["fired"] += 1
                rr["correct" if ok else "early"] += 1
                if ok:
                    rr["t"].append(ts)
    md(["tower", "format", "fps", "step", "forced", "brightness within ±10% for good", "settled brightness"], body)
    print(f"Settle rules scored on tower {t}'s {len(ae)} steps (correct = the frame's brightness is within "
          "10% of where it finally settles when the rule fires):\n")
    md(["rule", "fired correctly", "fired too early", "never fired", "median time when correct"],
       [[k, v["correct"], v["early"], v["never"],
         f"{np.median(v['t']):.2f} s" if v["t"] else "–"] for k, v in rule_rows.items()])

# ---- auto white balance -------------------------------------------------------------
print("## Auto white balance step response\n")
for t, d in TOWERS.items():
    awb = [r for r in rows(f"{d}/awb.jsonl") if r.get("cmd") == "awb"]
    if not awb:
        continue
    fig, axes = plt.subplots(1, 2, figsize=(9, 3.2), sharey=True)
    for ax, k in zip(axes, (0.5, 2.0)):
        for r in awb:
            if r["factor"] != k:
                continue
            xs = [x[0] for x in r["series"]]
            ax.plot(xs, [x[1] / r["ref_red"] for x in r["series"]], color=COL.get(r["fmt"], "k"),
                    lw=1.2, label=f"{r['fmt']} {r['fps_set']:.1f} fps, red")
            ax.plot(xs, [x[2] / r["ref_blue"] for x in r["series"]], color=COL.get(r["fmt"], "k"),
                    lw=1.0, ls=":", label=f"{r['fmt']} blue")
        ax.axhspan(0.95, 1.05, color="green", alpha=0.08)
        ax.set_xlim(0, 3)
        ax.set_title(f"ratios forced to ×{k}")
        ax.set_xlabel("seconds after auto white balance resumes")
        h_, l_ = ax.get_legend_handles_labels()
        uniq = dict(zip(l_, h_))
        ax.legend(uniq.values(), uniq.keys(), fontsize=6)
    axes[0].set_ylabel("ratio / settled ratio (band: ±5%)")
    fig.suptitle(f"Tower {t}: auto white balance recovering")
    save(fig, f"awb_{t}.png")
    md(["tower", "format", "fps", "forced", "within ±5% (and stays)", "settled red / blue",
        "spread on a still scene"],
       [[t, r["fmt"], f"{r['fps_set']:.2f}", f"×{r['factor']}",
         f"{r['stays5pct_s']:.2f} s ({r['stays5pct_frame']} frames)" if r["stays5pct_s"] is not None else "–",
         f"{r['ref_red']:.2f} / {r['ref_blue']:.2f}",
         f"{100 * r['ref_spread_red']:.1f}% / {100 * r['ref_spread_blue']:.1f}%"] for r in awb])


# ---- processing ------------------------------------------------------------------------
def matrix(d, which):
    return [r for r in rows(f"{d}/matrix_{which}.jsonl") if "kind" in r]


print("## Stage times per frame (4 threads)\n")
for which in ("real", "synth"):
    fig, ax = plt.subplots(figsize=(8, 3.6))
    labels, stacks = [], []
    for t, d in TOWERS.items():
        m = matrix(d, which)
        for k in ("rgb8", "c8_ea", "c12_bil", "c12_ea", "raw12", "raw12_planes"):
            e = 3 if k.startswith(("rgb", "c")) else 1
            r = next((x for x in m if x["kind"] == k and x["threads"] == 4 and x["effort"] == e
                      and "psnr_db" not in x), None)
            if r:
                labels.append(f"{t}\n{k}")
                stacks.append([r["unpack_ms"], r["demosaic_ms"], r["encode_ms"], r["write_ms"]])
    if not stacks:
        plt.close(fig)
        continue
    stacks = np.array(stacks)
    bottom = np.zeros(len(stacks))
    for j, (name, c) in enumerate(zip(("unpack", "de-Bayer / rearrange", "JXL encode", "file write"),
                                      ("#9edae5", "#ffbb78", "#1f77b4", "#7f7f7f"))):
        ax.bar(labels, stacks[:, j], bottom=bottom, label=name, color=c)
        bottom += stacks[:, j]
    for i, tot in enumerate(bottom):
        ax.text(i, tot, f"{tot:.0f}", ha="center", va="bottom", fontsize=7)
    ax.set_ylabel("ms per frame")
    ax.set_title(f"Time per saved frame by stage, effort 3 (raw: 1), 4 threads — "
                 f"{'real plant frame' if which == 'real' else 'synthetic detailed frame'}")
    ax.legend(fontsize=7)
    plt.xticks(fontsize=7)
    save(fig, f"stages_{which}.png")

    body = []
    for t, d in TOWERS.items():
        for r in matrix(d, which):
            if r["threads"] == 4 and "psnr_db" not in r and r["effort"] in (1, 3):
                body.append([t, KIND_LABEL.get(r["kind"], r["kind"]), r["effort"],
                             r["distance"] or "lossless", f"{r['unpack_ms']:.0f}",
                             f"{r['demosaic_ms']:.0f}", f"{r['encode_ms']:.0f}",
                             f"{r['write_ms']:.0f}", f"**{r['total_ms']:.0f}**",
                             f"{r['bytes'] / 1e6:.1f} MB"])
    print(f"### {'Real plant frame' if which == 'real' else 'Synthetic detailed frame'}\n")
    md(["tower", "kind", "effort", "distance", "unpack ms", "de-Bayer ms", "encode ms", "write ms",
        "total ms", "file"], body)

for t, d in TOWERS.items():
    m = matrix(d, "real")
    kinds = [k for k in ("rgb8", "c12_ea") if any(x["kind"] == k for x in m)]
    if not kinds:
        continue
    fig, axes = plt.subplots(1, len(kinds), figsize=(4 * len(kinds), 3))
    axes = np.atleast_1d(axes)
    for ax, k in zip(axes, kinds):
        grid = np.full((4, 4), np.nan)
        for x in m:
            if x["kind"] == k and "psnr_db" not in x:
                grid[x["effort"] - 1, x["threads"] - 1] = x["total_ms"]
        im = ax.imshow(grid, cmap="viridis_r")
        for i in range(4):
            for j in range(4):
                if not np.isnan(grid[i, j]):
                    ax.text(j, i, f"{grid[i, j]:.0f}", ha="center", va="center", color="w", fontsize=8)
        ax.set_xticks(range(4), [1, 2, 3, 4])
        ax.set_yticks(range(4), [1, 2, 3, 4])
        ax.set_xlabel("threads")
        ax.set_ylabel("effort")
        ax.set_title(f"{KIND_LABEL[k]} (ms)")
        ax.grid(False)
        fig.colorbar(im, ax=ax, shrink=0.8)
    fig.suptitle(f"Tower {t}: time per frame by effort and threads (real frame, distance 1.0)")
    save(fig, f"effort_threads_{t}.png")

print("## Quality setting: size, time and quality (effort 3, 4 threads)\n")
for which in ("real", "synth"):
    fig, axes = plt.subplots(1, 3, figsize=(11, 3.2))
    body = []
    for t, d in TOWERS.items():
        q = [x for x in matrix(d, which) if "psnr_db" in x and x["effort"] == 3]
        for k in ("rgb8", "c8_ea", "c12_bil", "c12_ea", "raw12_planes"):
            pts = sorted([x for x in q if x["kind"] == k and x["distance"] > 0], key=lambda x: x["distance"])
            if not pts:
                continue
            ls = "-" if t == list(TOWERS)[0] else "--"
            ds = [x["distance"] for x in pts]
            axes[0].plot(ds, [x["bytes"] / 1e6 for x in pts], ls, marker="o", ms=3, label=f"{t} {k}")
            axes[1].plot(ds, [x["psnr_db"] for x in pts], ls, marker="o", ms=3, label=f"{t} {k}")
            axes[2].plot(ds, [x["detail"] for x in pts], ls, marker="o", ms=3, label=f"{t} {k}")
            for x in pts:
                body.append([t, KIND_LABEL[k], x["distance"], f"{x['bytes'] / 1e6:.2f} MB",
                             f"{x['total_ms']:.0f} ms", f"{x['psnr_db']:.1f} dB", f"{x['detail']:.3f}"])
        for x in q:
            if x["distance"] == 0:
                body.append([t, KIND_LABEL[x["kind"]], "lossless", f"{x['bytes'] / 1e6:.2f} MB",
                             f"{x['total_ms']:.0f} ms", "exact" if x["differing"] == 0 else f"{x['psnr_db']:.1f} dB",
                             f"{x['detail']:.3f}"])
    if not body:
        plt.close(fig)
        continue
    for ax, yl in zip(axes, ("file size (MB)", "PSNR vs uncompressed (dB)", "fine detail kept (1 = all)")):
        ax.set_xlabel("distance (lower = better quality)")
        ax.set_ylabel(yl)
    axes[0].legend(fontsize=6)
    fig.suptitle(f"Quality setting on the {'real plant' if which == 'real' else 'synthetic detailed'} frame")
    save(fig, f"distance_{which}.png")
    print(f"### {'Real plant frame' if which == 'real' else 'Synthetic detailed frame'}\n")
    md(["tower", "kind", "distance", "file", "time", "PSNR", "detail kept"], body)

# ---- throughput, pile-up, heat ---------------------------------------------------------------
def summ(path):
    return [r for r in rows(path) if r.get("summary") == "pipeline"]


print("## Throughput and pile-up (synthetic frame)\n")
body = []
for t, d in TOWERS.items():
    for r in summ(f"{d}/throughput.jsonl") + summ(f"{d}/pileup.jsonl") + summ(f"{d}/pileup_real.jsonl"):
        per = f"{r['period_ms'] / 1000:.1f} s" if r["period_ms"] else "back to back"
        body.append([t, KIND_LABEL.get(r["kind"], r["kind"]), r["effort"], r["distance"] or "lossless",
                     f"{r['workers']} × {r['threads']}", per, f"{r['frames_per_s']:.2f}",
                     f"{r['latency_ms_p50'] / 1000:.2f} / {r['latency_ms_max'] / 1000:.2f} s",
                     r["backlog_at_end"], f"{r['cpu_s_per_frame']:.2f}", f"{r['peak_rss_mb']:.0f} MB",
                     f"{r['temp_c_mean']:.0f} / {r['temp_c_max']:.0f} °C", f"{r['watts_mean']:.1f} W"])
md(["tower", "kind", "effort", "distance", "workers × threads", "one frame every", "saves/s",
    "saved after (median / worst)", "behind at end", "CPU-s/frame", "peak memory", "temp mean / max",
    "power"], body)

fig, ax = plt.subplots(figsize=(7, 3.4))
for t, d in TOWERS.items():
    cfgs = {}
    for r in summ(f"{d}/pileup.jsonl"):
        cfgs.setdefault((r["kind"], r["effort"], r["distance"], r["workers"], r["threads"]), []).append(r)
    for (k, e, dist, w, th), rs in cfgs.items():
        rs.sort(key=lambda r: -r["period_ms"])
        ax.plot([r["period_ms"] / 1000 for r in rs], [r["latency_ms_max"] / 1000 for r in rs],
                marker="o", ms=4, ls="-" if t == list(TOWERS)[0] else "--",
                label=f"{t}: {k} d{dist or 'LL'} {w}×{th}")
ax.set_xlabel("one frame every … s (scan pace)")
ax.set_ylabel("worst time from capture to saved file (s)")
ax.set_yscale("log")
ax.invert_xaxis()
ax.set_title("When saving falls behind the scan (log scale; flat = keeping up)")
ax.legend(fontsize=6)
save(fig, "pileup.png")

for name, title in (("heat_max", "flat out (2 × 2)"), ("heat_paced", "one frame per 2 s (1 × 4)")):
    fig, ax1 = plt.subplots(figsize=(7, 3))
    ax2 = ax1.twinx()
    any_ = False
    for t, d in TOWERS.items():
        tr = [r for r in rows(f"{d}/{name}.jsonl") if "t" in r and "summary" not in r]
        if not tr:
            continue
        any_ = True
        ls = "-" if t == list(TOWERS)[0] else "--"
        ax1.plot([r["t"] / 60 for r in tr], [r["temp_c"] for r in tr], ls, color="#d62728", lw=0.8,
                 label=f"{t} temperature")
        ax2.plot([r["t"] / 60 for r in tr], [r["watts"] for r in tr], ls, color="#1f77b4", lw=0.8,
                 label=f"{t} power")
    if not any_:
        plt.close(fig)
        continue
    ax1.set_xlabel("minutes")
    ax1.set_ylabel("CPU package °C", color="#d62728")
    ax2.set_ylabel("package W", color="#1f77b4")
    ax2.grid(False)
    ax1.set_title(f"12-bit saving, {title}")
    l1, n1 = ax1.get_legend_handles_labels()
    l2, n2 = ax2.get_legend_handles_labels()
    ax1.legend(l1 + l2, n1 + n2, fontsize=7, loc="lower right")
    save(fig, f"{name}.png")

print("## Heat\n")
body = []
for t, d in TOWERS.items():
    idle = [r for r in rows(f"{d}/idle_start.jsonl") if r.get("summary") == "idle"]
    if idle:
        i = idle[0]
        body.append([t, "idle", f"{i['temp_c_mean']:.0f} / {i['temp_c_max']:.0f} °C",
                     f"{i['mhz_mean'] / 1000:.2f} GHz", f"{i['cpu_pct_mean']:.0f}%", f"{i['watts_mean']:.1f} W", "–"])
    for name in ("heat_paced", "heat_max"):
        for r in summ(f"{d}/{name}.jsonl"):
            tr = [x for x in rows(f"{d}/{name}.jsonl") if "done" in x]
            first = tr[59]["done"] if len(tr) > 60 else None
            last = tr[-1]["done"] - tr[-61]["done"] if len(tr) > 61 else None
            drift = f"{first} → {last} per min" if first is not None and last is not None else "–"
            body.append([t, f"{r['kind']} {r['workers']}×{r['threads']}, "
                            f"{'flat out' if not r['period_ms'] else 'one per %.1f s' % (r['period_ms'] / 1000)}",
                         f"{r['temp_c_mean']:.0f} / {r['temp_c_max']:.0f} °C", f"{r['mhz_mean'] / 1000:.2f} GHz",
                         f"{r['cpu_pct_mean']:.0f}%", f"{r['watts_mean']:.1f} W", drift])
md(["tower", "load (10 min)", "temp mean / max", "clock", "CPU busy", "power", "saves first → last minute"], body)

# ---- storage --------------------------------------------------------------------------------
print("## Storage per day (real plant frame sizes)\n")
fig, ax = plt.subplots(figsize=(7, 3.2))
body = []
periods = (2, 5, 10)
for t, d in TOWERS.items():
    q = {(x["kind"], x["distance"]): x["bytes"] for x in matrix(d, "real") if "psnr_db" in x and x["effort"] == 3}
    cands = [("rgb8", 0.8, "camera RGB8 d0.8 (today's setting)"), ("c12_ea", 1.0, "12-bit colour d1.0"),
             ("c12_ea", 0.5, "12-bit colour d0.5"), ("c12_ea", 0.0, "12-bit colour lossless"),
             ("raw12_planes", 0.0, "12-bit raw planes lossless")]
    for k, dist, label in cands:
        b = q.get((k, dist))
        if b is None:
            continue
        per_day = [b * 86400 / p / 1e9 for p in periods]
        body.append([t, label, f"{b / 1e6:.2f} MB"] + [f"{v:.0f} GB" for v in per_day])
        ax.barh(f"{t}: {label}", per_day[0], color="#1f77b4")
ax.set_xlabel("GB per day at one frame every 2 s")
ax.set_title("Storage, 24 h of scanning")
save(fig, "storage.png")
md(["tower", "saved as", "per frame"] + [f"per day, one frame / {p} s" for p in periods], body)