#!/usr/bin/env python3
"""retroDE_ps2 — Ch357 (re-scoped): PSMZ16S-accurate persistent-Z software ORACLE at the 384x381 scheduler geometry.

Codex pre-RTL gate: build a PSMZ16S-accurate oracle (authentic per-vertex Z interpolation, quantization/storage, GEQUAL,
ZTE, ZMSK), replace constant-Z flattening with authentic XYZ2 Z, and produce an expected reject/owner map quantifying how
the Z-tested framebuffer must differ from paint order.

MODEL (empirically grounded — see gs_sh3_z_census.py):
  * geometry: the EXISTING Ch356 scheduler union FB, 384x381 (OX,OY union origin), draws translated into it.
  * Z is interpolated LINEARLY in screen space (GS Z is screen-linear, NOT perspective-corrected).
  * PSMZ16S quantization: Zq = (24-bit vertex Z) >> 8  -> the high 16 bits (the whole scene uses the full 24-bit range;
    >>8 maps it onto 16 bits coherently; & 0xFFFF would wrap a surface). The signed "S" bias is order-preserving, so the
    GEQUAL reject/owner outcome is invariant to it.
  * shared PERSISTENT 16-bit Z buffer across ALL 3 epochs (cleared once to 0 = farthest). Per fragment:
      ZTE=1: pass iff Zq >= zbuf[px]  (GEQUAL, larger Z = nearer on PS2).
      on pass: write owner=epoch; ZMSK=0 -> zbuf[px]=Zq.
  * paint-order (what we ship today) = same coverage but ALWAYS overwrite (no Z test).

Outputs (to sim/data/top_psmct32_raster_demo/, LOCAL/gitignored):
  * sh3_zsched_zowner.mem  — per pixel: [31]cov [30]rejected-vs-paint [26:24]z_owner_epoch [23:8]zq16 [ ... ]
  * sh3_zsched_reject.mem  — per pixel: paint_owner vs z_owner diff mask (the far-late rejects)
  * a quantified report: covered px, pixels where Z rejects the paint-order winner, per-epoch owned counts (paint vs Z).

Usage: gs_sh3_z_oracle.py [dump.gs.zst] [--draw-list i0,i1,i2] [--emit]
"""
import sys, os
HERE=os.path.dirname(os.path.abspath(__file__)); ROOT=os.path.normpath(os.path.join(HERE,".."))
DATA=os.path.join(ROOT,"sim","data","top_psmct32_raster_demo")
import gs_sh3_recon as RC
import gs_make_sh3_multidraw_fixture as MD          # load_draws (verts incl z) + edge

DEFAULT_DRAWS=[11671, 19562, 89761]

# PSMZ16S 24/32-bit -> 16-bit reduction. NOT yet authoritatively pinned (Codex precondition): candidate hardware models
# give DIFFERENT reject maps, so the oracle is parameterized until ground-truth (PCSX2 dump replay) decides.
#   shift8 : Zq = Z >> 8    (high 16 of the scene's 24-bit range; coherent, no wrap)
#   mask16 : Zq = Z & 0xFFFF (low 16 — what PCSX2 WritePixel16 stores; WRAPS a 24-bit surface)
#   clamp16: Zq = min(Z, 0xFFFF) (near draws collapse to 0xFFFF)
def zquant(z, model):
    z=int(z)
    if model=="mask16":  return z & 0xFFFF
    if model=="clamp16": return 0xFFFF if z>0xFFFF else (0 if z<0 else z)
    q = z >> 8                                       # shift8 (default)
    return 0xFFFF if q>0xFFFF else (0 if q<0 else q)

def bbox(dr):
    xs=[v["x"] for v in dr["verts"]]; ys=[v["y"] for v in dr["verts"]]; return (min(xs),min(ys),max(xs),max(ys))

def main(argv):
    a=argv[1:]; dump=a[0] if a and not a[0].startswith("--") else None
    idxs=[int(x) for x in a[a.index("--draw-list")+1].split(",")] if "--draw-list" in a else list(DEFAULT_DRAWS)
    ZMODEL = a[a.index("--zmodel")+1] if "--zmodel" in a else "clamp16"
    # KNOWN VECTORS pinned against the PCSX2 SOFTWARE RASTER path (authoritative; NOT the bare memory write):
    #   GSLocalMemory.cpp:199   PSMZ16S -> PSM_FMT_16
    #   GSRendererSW.cpp:1439    z_max = 0xffffffff >> (fmt*8) = 0xFFFF
    #   GSDrawScanline…:1129     source_z = CLAMP(interpolated_z, z_max)   ("Clamp Z to ZPSM_FMT_MAX")
    #   GSDrawScanline…:1157     dest_z   = stored_z & 0xFFFF ; pass = source_z >= dest_z (GEQUAL)
    # => the model is CLAMP16 (min with 0xFFFF), NOT mask16. WritePixel16((u16)c) is only the final store, after the clamp.
    # The "S" in PSMZ16S = swizzled STORAGE layout, not signed depth (Codex) — address swizzle only, not the value reduction.
    for zin,exp in ((0x00FED407,0xFFFF),(0x00015534,0xFFFF),(0x0000B2C2,0xB2C2),(0xFFFFFFFF,0xFFFF),(0x00008000,0x8000),(0,0)):
        assert zquant(zin,"clamp16")==exp, f"clamp16 vector fail 0x{zin:x}->0x{zquant(zin,'clamp16'):x} exp 0x{exp:x}"
    print("[Zoracle] PSMZ16S known-vector self-test PASS (clamp16 == PCSX2 SW raster source-clamp to 0xFFFF)")
    if dump is None:
        import glob; c=glob.glob(os.path.join(ROOT,"captures","gs","silenthill3","*224139*.gs.zst"))
        if not c: sys.exit("no SH3 dump found; pass the .gs.zst path")
        dump=c[0]
    print(f"[Zoracle] dump={os.path.basename(dump)}  epochs={idxs}  (PSMZ16S zmodel={ZMODEL}, GEQUAL, persistent)")

    got,_=MD.load_draws(dump, idxs)
    for i in idxs:
        if i not in got: sys.exit(f"[Zoracle] FAIL: idx{i} not found as a textured draw")
    eps=[got[i] for i in idxs]

    # depth state gate: all real GEQUAL, ZMSK=0, shared ZBP (the census precondition)
    for e in eps:
        t=e["state"]["test"]; z=e["state"]["zbuf"]
        zte=(t>>16)&1; ztst=(t>>17)&3; zmsk=(z>>32)&1
        if not (zte==1 and ztst==2 and zmsk==0):
            sys.exit(f"[Zoracle] FAIL: idx{e['first_idx']} not real GEQUAL depth (zte={zte} ztst={ztst} zmsk={zmsk})")
    zbps={ (e['state']['zbuf']>>0)&0x1FF for e in eps }
    print(f"[Zoracle] all epochs GEQUAL/ZMSK=0; shared ZBP={zbps} (single Z buffer)")

    # union geometry = the EXISTING Ch356 scheduler FB (384x381)
    OX=int(min(bbox(e)[0] for e in eps)); OY=int(min(bbox(e)[1] for e in eps))
    UX=max(bbox(e)[2] for e in eps);      UY=max(bbox(e)[3] for e in eps)
    W=int(UX)-OX+1; H=int(UY)-OY+1; FBW=(W+63)//64; FBPXW=FBW*64
    print(f"[Zoracle] union origin=({OX},{OY}) -> FB {FBPXW}x{H} (FBW={FBW})  [existing Ch356 scheduler geometry]")
    NPX=FBPXW*H; edge=MD.edge

    # per-epoch texture idx + palette (authentic), for color/owner readout
    TEX=512*512; epdata=[]
    for e in eps:
        t0=e["state"]["tex0"]; mem,*_=RC.build_localmem_to(dump, e["first_idx"])
        idx=mem.read_psmt8(t0['tbp'], t0['tbw'], 512, 512)
        pal=RC.read_clut32(mem, t0['cbp'], order="grid")
        epdata.append(dict(idx=idx, pal=pal))

    def tris(e):
        fv=[dict(x=v["x"]-OX,y=v["y"]-OY,z=v["z"],s=v["s"],t=v["t"],q=v["q"]) for v in e["verts"]]
        return [(fv[i-2],fv[i-1],fv[i]) for i in range(2,len(fv))]

    # buffers
    zbuf=[-1]*NPX                       # stored Zq (16-bit); -1 = uninitialised (cleared "farther than any" => first frag passes)
    z_owner=[-1]*NPX; z_zq=[0]*NPX      # Z-tested owner epoch + its Zq
    paint_owner=[-1]*NPX                # paint-order owner (last covering, no Z)
    cov=[0]*NPX

    for k,e in enumerate(eps):
        for (v0,v1,v2) in tris(e):
            ar=edge(v0["x"],v0["y"],v1["x"],v1["y"],v2["x"],v2["y"])
            if abs(ar)<1e-9: continue
            inv=1.0/ar
            minx=max(0,int(min(v0["x"],v1["x"],v2["x"]))); maxx=min(FBPXW-1,int(max(v0["x"],v1["x"],v2["x"]))+1)
            miny=max(0,int(min(v0["y"],v1["y"],v2["y"]))); maxy=min(H-1,int(max(v0["y"],v1["y"],v2["y"]))+1)
            for py in range(miny,maxy+1):
                for px in range(minx,maxx+1):
                    cx,cy=px+0.5,py+0.5
                    w0=edge(v1["x"],v1["y"],v2["x"],v2["y"],cx,cy)*inv
                    w1=edge(v2["x"],v2["y"],v0["x"],v0["y"],cx,cy)*inv
                    w2=1.0-w0-w1
                    if w0<-0.001 or w1<-0.001 or w2<-0.001: continue
                    o=py*FBPXW+px; cov[o]=1
                    zf=w0*v0["z"]+w1*v1["z"]+w2*v2["z"]          # screen-linear Z interp (24-bit)
                    zq=zquant(zf, ZMODEL)                         # PSMZ16S reduction (model-parameterized, see zquant)
                    paint_owner[o]=k                             # paint-order always overwrites
                    if zq >= zbuf[o]:                            # GEQUAL (>= handles uninit -1 too)
                        zbuf[o]=zq; z_owner[o]=k; z_zq[o]=zq     # ZTE pass + ZMSK=0 write

    covered=sum(cov)
    # reject vs paint-order: covered pixels where the Z-tested owner != the paint-order owner (far-late fragment rejected)
    reject=[1 if (cov[o] and z_owner[o]!=paint_owner[o]) else 0 for o in range(NPX)]
    nreject=sum(reject)
    paint_ct=[0]*len(eps); z_ct=[0]*len(eps)
    for o in range(NPX):
        if cov[o]:
            if paint_owner[o]>=0: paint_ct[paint_owner[o]]+=1
            if z_owner[o]>=0:     z_ct[z_owner[o]]+=1

    print(f"\n[Zoracle] covered px={covered}")
    print(f"[Zoracle] per-epoch OWNED pixels (final, after all 3 epochs):")
    for k,e in enumerate(eps):
        print(f"        epoch {k} idx{e['first_idx']:>6}: paint-order={paint_ct[k]:>7}   Z-tested={z_ct[k]:>7}   (delta {z_ct[k]-paint_ct[k]:+d})")
    print(f"[Zoracle] pixels where Z REJECTS the paint-order winner (final FB MUST differ here): {nreject} "
          f"({100.0*nreject/max(covered,1):.2f}% of covered)")
    if nreject==0:
        print("[Zoracle] (no difference — paint order would already match; Z not needed. Unexpected for this scene.)")
    else:
        print(f"[Zoracle] => the far-late draw is depth-rejected in the overlaps; paint-order output demonstrably differs.")

    if "--emit" in a:
        def wmem(name, words, banner):
            with open(os.path.join(DATA,name),"w") as f:
                f.write(f"// {banner}\n")
                for x in words: f.write(f"{x&0xFFFFFFFF:08x}\n")
        # z-owner map: [31]cov [30]reject-vs-paint [27:24]z_owner(+1,0=none) [23:8]zq16 [3:0]paint_owner(+1)
        owner=[]
        for o in range(NPX):
            zo=(z_owner[o]+1)&0xF; po=(paint_owner[o]+1)&0xF
            owner.append((cov[o]<<31)|(reject[o]<<30)|(zo<<24)|((z_zq[o]&0xFFFF)<<8)|po)
        wmem("sh3_zsched_zowner.mem", owner, "Ch357 LOCAL persistent-Z oracle: cov|reject|z_owner|zq16|paint_owner. gitignored.")
        wmem("sh3_zsched_reject.mem", reject, "Ch357 LOCAL per-pixel reject-vs-paint mask (far-late Z rejects). gitignored.")
        print(f"[Zoracle] emitted sh3_zsched_zowner.mem + sh3_zsched_reject.mem -> {DATA}")
    return 0

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