#!/usr/bin/env python3
"""retroDE_ps2 — Ch357 follow-up: Z-usage CENSUS (read-only; no RTL, no board).

Codex: "Census Z usage afterward; no speculative Z work." Measures whether the scheduler's draws actually rely on depth
testing, and whether cross-draw Z persistence would change the composite vs the paint-order we already ship.

Reports, from the AUTHENTIC dump (NOT the flattened-Z fixtures):
  * per target epoch: TEST.ZTE / TEST.ZTST, ZBUF.ZMSK / ZBP / PSM, and the real per-vertex Z range.
  * scene-wide tally: for every textured draw kick, the active (ZTE, ZTST, ZMSK) — how much of the scene uses real Z.
  * if any target epoch does REAL depth testing (ZTE=1, ZTST in {GEQUAL,GREATER}, ZMSK=0): a pairwise screen-overlap +
    depth-order analysis — at pixels where a later draw overlaps an earlier one, does Z ever REJECT the later fragment
    (i.e. would Z reorder the composite away from paint-order)?
  * verdict: is cross-draw Z NEEDED for these draws, or is paint-order sufficient?

Usage: gs_sh3_z_census.py [dump.gs.zst] [--draw-list i0,i1,i2]
"""
import sys, os
HERE=os.path.dirname(os.path.abspath(__file__)); ROOT=os.path.normpath(os.path.join(HERE,".."))
import gs_texture_residency as R          # R.collect walks the GS event stream (same as the fixture tooling)
import gs_make_sh3_multidraw_fixture as MD

DEFAULT_DRAWS=[11671, 19562, 89761]
ZTST_NAME={0:"NEVER",1:"ALWAYS",2:"GEQUAL",3:"GREATER"}
ZPSM_NAME={0x00:"PSMZ32",0x01:"PSMZ24",0x02:"PSMZ16",0x0A:"PSMZ16S"}

def dec_test(t): return dict(ate=t&1, atst=(t>>1)&7, zte=(t>>16)&1, ztst=(t>>17)&3)
def dec_zbuf(z): return dict(zbp=z&0x1FF, psm=(z>>24)&0xF, zmsk=(z>>32)&1)

def real_z(test,zbuf):
    """A draw does REAL depth rejection iff Z-test is on, the method can actually reject (GEQUAL/GREATER), and Z is
    being written (ZMSK=0). ZTE=0 or ZTST=ALWAYS/NEVER, or ZMSK=1 -> Z cannot reorder overlapping draws."""
    te=dec_test(test); zb=dec_zbuf(zbuf)
    return te["zte"]==1 and te["ztst"] in (2,3) and zb["zmsk"]==0

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)
    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"[Zcensus] dump={os.path.basename(dump)}  target epochs={idxs}")

    # ---- scene-wide tally: active (ZTE,ZTST,ZMSK) at every textured draw kick ----
    d,h,events,uploads,runs,vram = R.collect(dump, 0)
    st=dict(PRIM=None, TEST_1=None, TEST_2=None, ZBUF_1=None, ZBUF_2=None)
    prim=dict(tme=0,ctxt=0); cur_key=None
    from collections import Counter
    tally=Counter(); tally_f1=Counter(); nkick=0; nkick_f1=0; frame=0
    def ctx(): return 1 if prim["ctxt"]==0 else 2
    for e in events:
        if e.kind=="FRAME_BOUNDARY": frame=e.frame+1; continue
        if e.kind!="GSREG": continue
        frame=e.frame; r,v=e.reg,e.value
        if   r=="PRIM": prim=dict(tme=(v>>4)&1, ctxt=(v>>9)&1); cur_key=None
        elif r in ("TEST_1","TEST_2","ZBUF_1","ZBUF_2"): st[r]=v
        elif r in ("XYZF2","XYZ2","XYZF3","XYZ3"):
            if not prim["tme"]: continue
            c=ctx(); test=st["TEST_%d"%c]; zbuf=st["ZBUF_%d"%c]
            if test is None or zbuf is None: continue
            # count ONE record per draw run (state latched at first kick of the run)
            k=(id(e),)  # unused; we key on run transitions via cur_key below
            if cur_key is None:
                te=dec_test(test); zb=dec_zbuf(zbuf)
                key=(te["zte"], te["ztst"], zb["zmsk"])
                tally[key]+=1; nkick+=1
                if frame==1: tally_f1[key]+=1; nkick_f1+=1
                cur_key=key
    def fmt_tally(t,n):
        out=[]
        for (zte,ztst,zmsk),c in sorted(t.items(), key=lambda kv:-kv[1]):
            lbl = "Z-OFF" if zte==0 else f"ZTE zt={ZTST_NAME[ztst]} zmsk={zmsk}"
            real = "REAL-DEPTH" if (zte==1 and ztst in (2,3) and zmsk==0) else "no-reorder"
            out.append(f"    {c:5d} ({100.0*c/max(n,1):4.1f}%)  {lbl:24s} -> {real}")
        return "\n".join(out)
    print(f"\n[Zcensus] SCENE-WIDE draw-run tally (all frames, {nkick} runs):")
    print(fmt_tally(tally,nkick))
    print(f"\n[Zcensus] frame f1 only ({nkick_f1} runs):")
    print(fmt_tally(tally_f1,nkick_f1))

    # ---- per-target-epoch detail ----
    got,_=MD.load_draws(dump, idxs)
    for i in idxs:
        if i not in got: sys.exit(f"[Zcensus] FAIL: idx{i} not found as a textured draw")
    eps=[got[i] for i in idxs]
    print(f"\n[Zcensus] TARGET epochs (authentic depth state + per-vertex Z):")
    print(f"        {'idx':>7} {'ZTE':>3} {'ZTST':>8} {'ZMSK':>4} {'ZBP':>4} {'ZPSM':>8} {'z_min':>10} {'z_max':>10} {'z_span':>8}  real-depth?")
    any_real=False
    for e in eps:
        te=dec_test(e["state"]["test"]); zb=dec_zbuf(e["state"]["zbuf"])
        zs=[v["z"] for v in e["verts"]]; zmin=min(zs); zmax=max(zs)
        rz = real_z(e["state"]["test"], e["state"]["zbuf"]); any_real|=rz
        print(f"        {e['first_idx']:>7} {te['zte']:>3} {ZTST_NAME[te['ztst']]:>8} {zb['zmsk']:>4} {zb['zbp']:>4} "
              f"{ZPSM_NAME.get(zb['psm'],hex(zb['psm'])):>8} {zmin:>10} {zmax:>10} {zmax-zmin:>8}  {'YES' if rz else 'no'}")

    # ---- verdict ----
    print()
    if not any_real:
        # explain WHY paint-order is exactly the hardware behaviour for these draws
        reasons=set()
        for e in eps:
            te=dec_test(e["state"]["test"]); zb=dec_zbuf(e["state"]["zbuf"])
            if te["zte"]==0: reasons.add("ZTE=0 (Z-test disabled)")
            elif te["ztst"]==1: reasons.add("ZTST=ALWAYS (test never rejects)")
            elif te["ztst"]==0: reasons.add("ZTST=NEVER")
            if zb["zmsk"]==1: reasons.add("ZMSK=1 (Z buffer read-only, no depth written)")
        print(f"[Zcensus] VERDICT: cross-draw Z is NOT needed for these draws — {', '.join(sorted(reasons))}.")
        print(f"[Zcensus] The GS applies no depth rejection here, so paint-order (dump order) IS the hardware result.")
        print(f"[Zcensus] Recommendation (matches Codex 'no speculative Z'): do NOT add Z for this scheduler scene.")
        return 0

    # real depth present -> pairwise screen-overlap + depth-order analysis (does Z reorder vs paint-order?)
    print(f"[Zcensus] REAL depth testing present -> checking whether Z would REORDER the composite vs paint-order...")
    edge=MD.edge
    def raster_z(e):
        """per-pixel {(x,y): z_interp} for a draw, using its authentic vertex Z (barycentric)."""
        cov={}
        fv=e["verts"]
        for i in range(2,len(fv)):
            v0,v1,v2=fv[i-2],fv[i-1],fv[i]
            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=int(min(v0["x"],v1["x"],v2["x"])); maxx=int(max(v0["x"],v1["x"],v2["x"]))+1
            miny=int(min(v0["y"],v1["y"],v2["y"])); maxy=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
                    z=w0*v0["z"]+w1*v1["z"]+w2*v2["z"]
                    cov[(px,py)]=z    # last tri wins within a draw (fine for a coverage/Z sample)
        return cov
    zmaps=[raster_z(e) for e in eps]
    reorder_total=0
    for bi in range(len(eps)):
        for ai in range(bi):   # A earlier (ai<bi) in dump order, B later
            A=zmaps[ai]; B=zmaps[bi]; diff=0; ov=0
            for p,zb in B.items():
                if p in A:
                    ov+=1
                    # GS: larger Z = nearer. GEQUAL/GREATER: later B passes iff zB >= (or >) zA. If it would be REJECTED
                    # (zB < zA) then Z keeps A -> differs from paint-order (which always takes B).
                    if zb < A[p]: diff+=1
            reorder_total+=diff
            print(f"[Zcensus]   pair A=idx{eps[ai]['first_idx']} over-drawn by B=idx{eps[bi]['first_idx']}: "
                  f"overlap {ov}px, Z-would-reject-B {diff}px ({100.0*diff/max(ov,1):.2f}% of overlap)")
    print()
    if reorder_total==0:
        print(f"[Zcensus] VERDICT: real Z state present, but Z would reject the later draw at 0 overlapping pixels -> the "
              f"depth order MATCHES paint-order for these draws. Cross-draw Z would NOT change the composite.")
        print(f"[Zcensus] Recommendation: paint-order is sufficient for THIS scene; no Z needed (revisit only if a chosen "
              f"scene shows reorder>0).")
    else:
        print(f"[Zcensus] VERDICT: Z WOULD reorder the composite at ~{reorder_total} pixels -> cross-draw Z IS needed to "
              f"match the hardware for these draws. This justifies (non-speculative) Z work as the next rung.")
        print(f"[Zcensus] CAVEAT: the exact px count uses a RAW vertex-Z compare; the precise value depends on the "
              f"PSMZ16S quantization this census does not model. The ORDERING (near vs far cluster) is robust to that.")
    return 0

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