# E-OS v4 — evaluation: full-res (app) PSNR on fresh renders + preview sheet.
import os, sys, math, json, random
import numpy as np
import torch
sys.path.insert(0, "/home/z/my-project/scripts/os4")
from os4_models import FrameGenerator
import os4_spec as S
import os4_draw as D
from gen_os4_dataset import build_vec
from PIL import Image

CKPT = "/home/z/my-project/checkpoints4/gen_ckpt_final.pt"
OUT = "/home/z/my-project/download/e-os-v4"
os.makedirs(OUT, exist_ok=True)

def main():
    m = FrameGenerator()
    m.load_state_dict(torch.load(CKPT, map_location="cpu", weights_only=False)["m"])
    m.eval()
    rng = random.Random(4242)
    ps, previews = [], []
    N = 40
    with torch.no_grad():
        for i in range(N):
            st = __import__("gen_os4_dataset").gen_scene(900000 + i)
            hint = np.asarray(D.draw_scene(st, "hint"), np.float32) / 255.0
            gt = np.asarray(D.draw_scene(st, "styled"), np.float32) / 255.0
            h = torch.from_numpy(hint).permute(2, 0, 1).unsqueeze(0)
            s = torch.from_numpy(build_vec(st)).unsqueeze(0)
            pred = m(h, s)[0].permute(1, 2, 0).numpy()
            mse = ((pred - gt) ** 2).mean()
            ps.append(10 * math.log10(1.0 / max(1e-9, mse)))
            if i % 8 == 0:
                h2 = np.asarray(Image.fromarray((hint * 255).astype(np.uint8))
                                .resize((S.FULL_W, S.FULL_H)), np.float32) / 255.0
                row = np.concatenate([h2, gt, np.clip(pred, 0, 1)], axis=1)
                previews.append((row * 255).astype(np.uint8))
    ps = np.array(ps)
    print("app-res PSNR: mean %.2f dB | median %.2f | p10 %.2f | min %.2f (N=%d)" %
          (ps.mean(), np.median(ps), np.percentile(ps, 10), ps.min(), N))
    sheet = np.concatenate(previews, axis=0)
    Image.fromarray(sheet).save(os.path.join(OUT, "eval_preview_grid.png"))
    json.dump({"psnr_mean": float(ps.mean()), "psnr_median": float(np.median(ps)),
               "psnr_p10": float(np.percentile(ps, 10)), "psnr_min": float(ps.min()),
               "n": N}, open(os.path.join(OUT, "eval_metrics.json"), "w"), indent=1)
    print("saved", os.path.join(OUT, "eval_preview_grid.png"))

if __name__ == "__main__":
    main()