#!/usr/bin/env python3
"""retroDE_ps2 — quantitative fidelity metric: our rendered frame vs the real PCSX2 frame (reference C).

Turns "does it look like SH3?" into numbers. Consumes any 640x480 RGB(A) PNG we produce (software composite B
from gs_sh3_frame_ref.py, or a framebuffer dumped from an RTL sim) and a PCSX2 screenshot (C), and reports:
  - GLOBAL: MSE, PSNR, mean-abs-error, %pixels within a per-channel tolerance.
  - PER-REGION: a GxH tile grid of per-tile MSE, the worst-N tiles ranked, and a diff heatmap PNG so the
    deficit is ATTRIBUTED (which part of the screen, how badly) instead of eyeballed.
  - Optionally restricted to PAINTED pixels only (alpha>0 in ours) so background fill doesn't dilute the score.

Reference C is the display output of the real GS; our composite is opaque-textures-only (no GS blend/fog/light),
so a perfect score is NOT expected. The value is the ATTRIBUTION: a high-error tile localized to, say, the
alpha-blended fog band tells us blending is the next rung; uniformly high error would implicate reciprocal/UV.

Alignment: C is resized to 640x480. PCSX2 NTSC output may be cropped/offset a few px vs GS screen space; use
--search P to brute-force the best integer (dx,dy) shift in [-P..P] (minimising global MSE) before scoring.

Usage:
  gs_frame_metric.py --ours frameN_composite.png --ref pcsx2_ref.png [--grid 16x12] [--worst 12]
                     [--tol 16] [--painted-only] [--search 4] [--out DIR]
"""
import sys, os

def load_rgb(path, resize=None):
    from PIL import Image
    img = Image.open(path).convert("RGBA")
    if resize and img.size != resize:
        img = img.resize(resize, Image.BILINEAR)
    return img

def main(argv):
    if len(argv) < 2 or "--ours" not in argv or "--ref" not in argv:
        print(__doc__); return 2
    def opt(n, dv=None): return argv[argv.index(n)+1] if n in argv else dv
    ours_p = opt("--ours"); ref_p = opt("--ref")
    gx, gy = (int(v) for v in opt("--grid", "16x12").lower().split("x"))
    worst_n = int(opt("--worst", "12")); tol = int(opt("--tol", "16"))
    search = int(opt("--search", "0")); painted_only = ("--painted-only" in argv)
    outdir = opt("--out", os.path.dirname(os.path.abspath(ours_p)))
    os.makedirs(outdir, exist_ok=True)

    from PIL import Image
    ours = load_rgb(ours_p)
    W, H = ours.size
    ref = load_rgb(ref_p, resize=(W, H))
    op = ours.load(); rp = ref.load()

    def scored_pixels(dx, dy):
        """Yield (x,y,(or,og,ob),(rr,rg,rb)) for pixels compared at shift (dx,dy)."""
        for y in range(H):
            ry = y + dy
            if ry < 0 or ry >= H: continue
            for x in range(W):
                rx = x + dx
                if rx < 0 or rx >= W: continue
                o = op[x, y]
                if painted_only and o[3] == 0: continue
                yield x, y, o[:3], rp[rx, ry][:3]

    def mse_at(dx, dy):
        s = n = 0
        for _, _, o, r in scored_pixels(dx, dy):
            s += (o[0]-r[0])**2 + (o[1]-r[1])**2 + (o[2]-r[2])**2; n += 1
        return (s/(3*n) if n else float("inf")), n

    # optional integer-shift alignment
    bdx = bdy = 0
    if search > 0:
        best = None
        for dy in range(-search, search+1):
            for dx in range(-search, search+1):
                m, n = mse_at(dx, dy)
                if best is None or m < best[0]: best = (m, dx, dy)
        _, bdx, bdy = best
        print(f"[metric] best alignment shift dx={bdx} dy={bdy} (searched +/-{search})")

    # global stats + per-tile accumulation at the chosen shift
    tiles_se = [[0]*gx for _ in range(gy)]
    tiles_n  = [[0]*gx for _ in range(gy)]
    tw, th = W/gx, H/gy
    S = Nabs = 0; N = 0; within = 0
    diff = Image.new("RGB", (W, H))
    dp = diff.load()
    for x, y, o, r in scored_pixels(bdx, bdy):
        e = [abs(o[i]-r[i]) for i in range(3)]
        se = e[0]**2 + e[1]**2 + e[2]**2
        S += se; Nabs += e[0]+e[1]+e[2]; N += 1
        if max(e) <= tol: within += 1
        gxi = min(gx-1, int(x/tw)); gyi = min(gy-1, int(y/th))
        tiles_se[gyi][gxi] += se; tiles_n[gyi][gxi] += 1
        # heatmap: brighter red = larger per-pixel error (scaled so 0..441 -> 0..255)
        mag = min(255, int((se**0.5)))
        dp[x, y] = (mag, 0, 128-min(128, mag//2))

    if N == 0:
        print("[metric] no comparable pixels (painted-only with empty frame?)"); return 1
    mse = S/(3*N); import math
    psnr = 10*math.log10((255.0**2)/mse) if mse > 0 else float("inf")
    print(f"[metric] compared {N} px ({'painted-only' if painted_only else 'all'}), grid {gx}x{gy}, tol +/-{tol}/ch")
    print(f"[metric] MSE={mse:.2f}  PSNR={psnr:.2f} dB  MAE={Nabs/(3*N):.2f}/ch  within-tol={100*within/N:.1f}%")

    tile_mse = []
    for gyi in range(gy):
        for gxi in range(gx):
            n = tiles_n[gyi][gxi]
            if n: tile_mse.append((tiles_se[gyi][gxi]/(3*n), gxi, gyi, n))
    tile_mse.sort(reverse=True)
    print(f"[metric] worst {min(worst_n,len(tile_mse))} tiles (col,row of {gx}x{gy} grid) by MSE:")
    for m, gxi, gyi, n in tile_mse[:worst_n]:
        px0, py0 = int(gxi*tw), int(gyi*th)
        print(f"    tile(c{gxi:2d},r{gyi:2d}) screen x[{px0}..{px0+int(tw)}] y[{py0}..{py0+int(th)}]  MSE={m:.1f}  n={n}")

    base = os.path.splitext(os.path.basename(ours_p))[0]
    hm = os.path.join(outdir, f"{base}_vs_ref_heatmap.png"); diff.save(hm)
    # triptych: ours | ref | heatmap
    trip = Image.new("RGB", (W*3+16, H), (30, 30, 30))
    trip.paste(ours.convert("RGB"), (0, 0)); trip.paste(ref, (W+8, 0)); trip.paste(diff, (W*2+16, 0))
    tp = os.path.join(outdir, f"{base}_vs_ref_triptych.png"); trip.save(tp)
    print(f"[metric] wrote {hm}\n[metric] wrote {tp} (ours | C | error-heatmap)")
    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv))