"""Plot the patch-number sweep: 4 small-multiple panels (one metric each, since
they move in opposite directions -> never a dual axis), P on a log2 x-axis with
real P ticks, sweep scatter + quadratic fit, and trained-reference markers
(single-state normalized @P=1, patch @P=8) for cross-check.
"""
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

J = json.load(open("workspace/markdown/patch_number_sweep.json"))
ps = J["ps"]; sweep = {int(k): v for k, v in J["sweep"].items()}
refs = J["refs"]; fits = J["fits"]
x = np.log2(np.array(ps, float))

# CVD-safe: blue (sweep) + orange (trained refs); recessive gray grid/ink.
BLUE, ORANGE, INK, MUTED, GRID = "#2F6FBF", "#E8710A", "#1c1c1c", "#6b6b6b", "#e6e6e6"
TITLES = {
    "clip_text_align": "CLIP text-align (prompt fidelity) ↑",
    "clip_temporal": "CLIP temporal coherence ↑",
    "flow_smoothness": "Optical-flow smoothness ↑",
    "flow_magnitude": "Flow magnitude (dynamic degree)",
}
metrics = J["metrics"]
ref_pt = {"normalized": 1, "patch": 8}  # model -> its trained P

fig, axes = plt.subplots(2, 2, figsize=(11, 7.6), dpi=130)
fig.suptitle("Quality vs patch number P  —  arl2_normalized_patch (P=8 trained ckpt, inference-P swept)\n"
             "480p / 21f / seed 0 / 8 prompts · CAVEAT: P=8 in-distribution, other P inference-OOD",
             fontsize=11, y=0.99)

for ax, k in zip(axes.flat, metrics):
    y = np.array([sweep[p][k] for p in ps])
    # quadratic fit in log2(P)
    xs = np.array(fits[k]["curve_x_log2P"]); ys = np.array(fits[k]["curve_y"])
    ax.plot(xs, ys, "-", color=BLUE, lw=2, alpha=0.55, zorder=1)
    ax.scatter(x, y, s=46, color=BLUE, zorder=3, edgecolor="white", linewidth=1.2)
    # trained-reference cross-check markers
    for m, P in ref_pt.items():
        if m in refs and k in refs[m]:
            ax.scatter([np.log2(P)], [refs[m][k]], marker="*", s=200, color=ORANGE,
                       zorder=4, edgecolor="white", linewidth=1.0,
                       label=f"trained {m} (P={P})")
    ax.set_title(TITLES.get(k, k), fontsize=10.5, color=INK, loc="left")
    ax.set_xticks(x); ax.set_xticklabels([str(p) for p in ps], fontsize=9)
    ax.set_xlabel("patch number P (log₂ spacing)", fontsize=9, color=MUTED)
    ax.tick_params(colors=MUTED, labelsize=9)
    ax.grid(True, color=GRID, lw=0.8, zorder=0)
    for s in ax.spines.values():
        s.set_color(GRID)
    r = fits[k]["pearson_r_vs_log2P"]
    ax.text(0.03, 0.06, f"Pearson r vs log₂P = {r:+.2f}", transform=ax.transAxes,
            fontsize=9, color=INK, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec=GRID, alpha=0.85))
    if k == "clip_text_align":
        ax.legend(loc="upper left", fontsize=8, frameon=False)

fig.tight_layout(rect=[0, 0, 1, 0.95])
out = "workspace/markdown/patch_number_sweep_curve.png"
fig.savefig(out, bbox_inches="tight")
print("wrote", out)