#!/usr/bin/env python3
"""plot_common.py - the shared-frame comparison between towers.

    ~/.venvs/fgc-plots/bin/python plot_common.py <out_dir> <label>=<results-common dir> [...]

Reads what run_common_frame.sh wrote on each tower (same frame everywhere) and
writes stages_common.png, sizes_common.png and pace_common.png, and prints the
Markdown tables for the report.
"""
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:])
plt.rcParams.update({"figure.dpi": 110, "font.size": 9, "axes.grid": True, "grid.alpha": 0.3})
LABEL = {"rgb8": "camera RGB8 → JXL\n(today)", "c8_ea": "raw 8 → colour\n→ JXL",
         "c12_bil": "raw 12 → colour\n(bilinear) → JXL", "c12_ea": "raw 12 → colour\n(edge-aware) → JXL",
         "raw12": "raw 12 mosaic\nlossless", "raw12_planes": "raw 12 planes\nlossless"}
NAME = {k: v.replace("\n", " ") for k, v in LABEL.items()}
NAME["c8_bil"] = "raw 8 → colour (bilinear) → JXL"


def rows(path):
    return [json.loads(l) for l in open(path) if l.startswith("{")] if os.path.exists(path) else []


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


M = {t: [r for r in rows(f"{d}/matrix_common.jsonl") if "kind" in r] for t, d in TOWERS.items()}


def pick(t, kind, effort, threads=4, distance=None, quality=False):
    for r in M[t]:
        if r["kind"] == kind and r["effort"] == effort and r["threads"] == threads \
                and ("psnr_db" in r) == quality \
                and (distance is None or abs(r["distance"] - distance) < 1e-6):
            return r
    return None


# ---- stage breakdown, both towers side by side --------------------------------------
kinds = [("rgb8", 4, 0.8), ("c8_ea", 3, 1.0), ("c12_bil", 3, 1.0), ("c12_ea", 3, 1.0),
         ("raw12", 1, 0.0), ("raw12_planes", 1, 0.0)]
print("## Per step, shared frame, 4 threads\n")
body = []
for k, e, dist in kinds:
    for t in TOWERS:
        r = pick(t, k, e, distance=dist if k == "rgb8" else None, quality=(k == "rgb8"))
        if r is None:
            r = pick(t, k, e)
        if r:
            body.append([NAME[k], t, e, "lossless" if not r["distance"] else r["distance"],
                         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:.2f} MB"])
md(["path", "tower", "effort", "distance", "unpack ms", "colour conversion ms", "JXL encode ms",
    "write ms", "total ms", "file"], body)

fig, ax = plt.subplots(figsize=(9, 3.8))
xs, labels, stacks = [], [], []
pos = 0
for k, e, dist in kinds:
    for i, t in enumerate(TOWERS):
        r = pick(t, k, e, distance=dist if k == "rgb8" else None, quality=(k == "rgb8")) or pick(t, k, e)
        if r:
            xs.append(pos + i * 0.42)
            stacks.append([r["unpack_ms"], r["demosaic_ms"], r["encode_ms"], r["write_ms"]])
            labels.append(t)
    pos += 1
stacks = np.array(stacks)
bottom = np.zeros(len(stacks))
for j, (nm, c) in enumerate(zip(("unpack", "colour conversion / rearrange", "JPEG XL encode", "file write"),
                                ("#9edae5", "#ffbb78", "#1f77b4", "#7f7f7f"))):
    ax.bar(xs, stacks[:, j], 0.4, bottom=bottom, label=nm, color=c)
    bottom += stacks[:, j]
for x, tot, lb in zip(xs, bottom, labels):
    ax.text(x, tot, f"{tot:.0f}\n{lb}", ha="center", va="bottom", fontsize=6.5)
ax.set_xticks(np.arange(len(kinds)) + 0.21, [LABEL[k] for k, _, _ in kinds], fontsize=7)
ax.set_ylabel("ms per frame")
ax.set_ylim(0, bottom.max() * 1.18)
ax.set_title("Same frame on both computers: time per saved frame by step (4 threads; "
             "colour: distance 1.0 effort 3, today's RGB8: 0.8 / 4)")
ax.legend(fontsize=7, loc="upper right")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "stages_common.png"))
plt.close(fig)

# ---- colour conversion alone ------------------------------------------------------------
print("## Colour conversion and unpack alone (median over all runs), shared frame\n")
body = []
for t in TOWERS:
    for k in ("c8_bil", "c8_ea", "c12_bil", "c12_ea"):
        rs = [r for r in M[t] if r["kind"] == k]
        if rs:
            body.append([t, NAME[k], f"{np.median([r['unpack_ms'] for r in rs]):.0f} ms",
                         f"{np.median([r['demosaic_ms'] for r in rs]):.0f} ms"])
md(["tower", "path", "unpack", "colour conversion"], body)

# ---- size and time vs distance ---------------------------------------------------------
print("## Quality setting, shared frame (effort 3, 4 threads)\n")
fig, axes = plt.subplots(1, 2, figsize=(9, 3.2))
body = []
for t in TOWERS:
    for k in ("rgb8", "c8_ea", "c12_ea", "raw12_planes"):
        pts = sorted([r for r in M[t] if r["kind"] == k and "psnr_db" in r and r["effort"] == 3],
                     key=lambda r: r["distance"])
        lossy = [r for r in pts if r["distance"] > 0]
        if not lossy:
            continue
        ls = "-" if t == list(TOWERS)[0] else "--"
        axes[0].plot([r["distance"] for r in lossy], [r["bytes"] / 1e6 for r in lossy], ls, marker="o", ms=3,
                     label=f"{t}: {NAME[k]}")
        axes[1].plot([r["distance"] for r in lossy], [r["total_ms"] for r in lossy], ls, marker="o", ms=3,
                     label=f"{t}: {NAME[k]}")
        for r in pts:
            body.append([t, NAME[k], "lossless" if r["distance"] == 0 else r["distance"],
                         f"{r['bytes'] / 1e6:.2f} MB", f"{r['total_ms']:.0f} ms"])
axes[0].set_ylabel("file size (MB)")
axes[1].set_ylabel("time per frame (ms)")
for ax in axes:
    ax.set_xlabel("distance (lower = higher quality)")
axes[0].legend(fontsize=6)
fig.suptitle("Same frame: file size and save time by quality setting")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "sizes_common.png"))
plt.close(fig)
md(["tower", "path", "distance", "file", "time"], body)

# ---- pace -------------------------------------------------------------------------------
print("## Save pace, shared frame\n")
body = []
for t, d in TOWERS.items():
    for r in [x for x in rows(f"{d}/throughput.jsonl") + rows(f"{d}/pileup.jsonl") if x.get("summary") == "pipeline"]:
        per = f"{r['period_ms'] / 1000:.1f} s" if r["period_ms"] else "back to back"
        body.append([t, NAME.get(r["kind"], r["kind"]), r["effort"], r["distance"] or "lossless",
                     f"{r['workers']} × {r['threads']}", per, f"{r['frames_per_s']:.2f}",
                     f"{1 / r['frames_per_s']:.2f} s" if r["frames_per_s"] else "–",
                     f"{r['latency_ms_p50'] / 1000:.2f} / {r['latency_ms_max'] / 1000:.2f} s",
                     r["backlog_at_end"], f"{r['peak_rss_mb']:.0f} MB"])
md(["tower", "path", "effort", "distance", "workers × threads", "one frame every", "saves/s",
    "one save per", "saved after (median / worst)", "behind at end", "peak memory"], body)

fig, ax = plt.subplots(figsize=(7, 3.2))
cfgs = [("rgb8", 1), ("rgb8", 2), ("c8_ea", 1), ("c8_ea", 2), ("c12_ea", 1), ("c12_ea", 2)]
w = 0.8 / len(TOWERS)
for i, (t, d) in enumerate(TOWERS.items()):
    thr = {(r["kind"], r["workers"]): r["frames_per_s"] for r in rows(f"{d}/throughput.jsonl")
           if r.get("summary") == "pipeline"}
    vals = [thr.get(c, 0) for c in cfgs]
    b = ax.bar(np.arange(len(cfgs)) + i * w, vals, w, label=f"tower {t}")
    ax.bar_label(b, fmt="%.2f", fontsize=7)
ax.set_xticks(np.arange(len(cfgs)) + w * (len(TOWERS) - 1) / 2,
              [f"{NAME[k].split(' →')[0]}\n{'1 × 4' if n == 1 else '2 × 2'}" for k, n in cfgs], fontsize=7)
ax.set_ylabel("saves per second (back to back)")
ax.set_title("Same frame: how many frames each computer can save per second")
ax.legend(fontsize=7)
fig.tight_layout()
fig.savefig(os.path.join(OUT, "pace_common.png"))
plt.close(fig)