# E-OS v6 — eval: PSNR on fresh val scenes @ app-res + preview grids.
import os, sys, json, math
import numpy as np
import torch
sys.path.insert(0, "/home/z/my-project/scripts/os6")
from os6_models import FrameGenerator
from os6_draw import render_frame, rasterize_hint, build_state
from os6_spec import lay_home, lay_random, CARD, CIRCLE
from PIL import Image

CKPT = sys.argv[1] if len(sys.argv) > 1 else \
    "/home/z/my-project/checkpoints6/gen_ckpt_final.pt"
OUT = "/home/z/my-project/download/e-os-v6/eval"
os.makedirs(OUT, exist_ok=True)

m = FrameGenerator()
ck = torch.load(CKPT, map_location="cpu", weights_only=False)
m.load_state_dict(ck["m"])
m.eval()
print("ckpt step %s best %s" % (ck.get("step"), ck.get("best")))

rng = np.random.default_rng(999)


def run(shapes, theme, day, shimmer, ripples, pressed=0):
    hint = rasterize_hint(shapes, ripples, 256, 544)
    st = build_state(theme, shimmer, day, pressed)
    with torch.no_grad():
        p = m(torch.from_numpy(hint).unsqueeze(0),
              torch.from_numpy(st).unsqueeze(0))
    return p[0].permute(1, 2, 0).numpy()


def gt(shapes, theme, day, shimmer, ripples, pressed=0):
    return render_frame(shapes, theme, day, shimmer, ripples, pressed) / 255.0


# ---- 40 fresh val scenes ----
ps = []
samples = []
for k in range(40):
    shapes = (lay_home if k % 3 == 0 else
              lay_random if k % 3 == 1 else
              (lay_home if k % 2 else lay_random))(rng)
    shapes.sort(key=lambda s: s.get("z", 0))
    theme = int(rng.integers(0, 4))
    day = float(rng.uniform(0.05, 0.95))
    shimmer = float(rng.uniform(0, 1))
    rip = (float(rng.uniform(0.15, 0.85)), float(rng.uniform(0.15, 0.85)),
           float(rng.uniform(0.05, 0.9))) if k % 4 != 3 else None
    args = (shapes, theme, day, shimmer, [rip] if rip else [])
    g = gt(*args)
    p = np.clip(run(*args), 0, 1)
    mse = ((g - p) ** 2).mean()
    ps.append(10 * math.log10(1.0 / max(1e-8, mse)))
    if k < 6:
        samples.append((shapes, theme, rip, g, p))
ps = np.array(ps)
print("PSNR mean %.2f median %.2f (40 scenes @ 512x1088)" %
      (ps.mean(), np.median(ps)))

# ---- preview grid: GT | NPU-Output (6 rows) ----
rows = []
for (shapes, theme, rip, g, p) in samples:
    rows.append(np.concatenate([g, p], axis=1))
grid = (np.concatenate(rows, axis=0) * 255).astype(np.uint8)
Image.fromarray(grid).save(os.path.join(OUT, "eval_gt_vs_npu.png"))

# ---- GENERALIZATION PROOF: edited layout the net never saw ----
sh = lay_home(rng)
sh[0] = dict(sh[0], x0=100, x1=430, y0=90, y1=330)       # card moved/resized
sh.append(dict(t=CIRCLE, x0=60, y0=430, x1=150, y1=520, z=9))  # extra circle
sh.append(dict(t=CARD, x0=210, y0=560, x1=470, y1=700, z=8))   # extra card
sh.sort(key=lambda s: s.get("z", 0))
theme = 1
rip = (0.4, 0.5, 0.35)
ge = gt(sh, theme, 0.5, 0.4, [rip])
pe = np.clip(run(sh, theme, 0.5, 0.4, [rip]), 0, 1)
mse = ((ge - pe) ** 2).mean()
psnr_edit = 10 * math.log10(1.0 / max(1e-8, mse))
hint_e = rasterize_hint(sh, [rip], 256, 544)
wire = np.stack([hint_e[0]] * 3, -1)
proof = (np.concatenate([ge, pe], axis=1) * 255).astype(np.uint8)
Image.fromarray(proof).save(os.path.join(OUT, "eval_generalization.png"))
print("EDITED-LAYOUT PSNR %.2f dB (proof of generalization)" % psnr_edit)

json.dump(dict(psnr_mean=float(ps.mean()), psnr_median=float(np.median(ps)),
               psnr_edited_layout=float(psnr_edit), n=40, app_res=[512, 1088],
               ckpt_step=ck.get("step")),
          open(os.path.join(OUT, "eval_metrics.json"), "w"), indent=1)
print("EVAL DONE")