#!/usr/bin/env python3
"""retroDE_ps2 — Ch357: rank authentic supported draw GROUPS by clamp16 depth-rejection strength.

Codex: "Before LPDDR-Z RTL, rank authentic supported draw groups using the corrected clamp16 model and select one with
stronger depth rejection if available. Keep the current trio as regression coverage."

The reject map (z-tested owner != paint-order owner) depends ONLY on coverage + Z, not on the textures — so we rank fast
from the census verts (no per-epoch texture reconstruction). Model = clamp16 (PCSX2 SW raster: source_z=min(z,0xFFFF),
dest=stored&0xFFFF, GEQUAL). A strong acceptance group has heavy overlap where LATER-drawn draws are FARTHER (smaller Z)
so GEQUAL rejects them — paint-order and depth-order disagree.

Constraints (scheduler-supported, like gs_make_sh3_scheduler_fixture): same frame f1; >=3 DISTINCT textures AND distinct
CLUTs; PSMT8 512x512 perspective TME (fst=0); resident tex+CLUT; on-screen; union bbox that fits the 384x381 rung.

Usage: gs_sh3_z_rank.py [dump.gs.zst] [--topk N] [--max-fb 384x480]
"""
import sys, os, itertools
HERE=os.path.dirname(os.path.abspath(__file__)); ROOT=os.path.normpath(os.path.join(HERE,".."))
import gs_sh3_draw_census as C
import gs_make_sh3_multidraw_fixture as MD
edge=MD.edge

def clamp16(z): z=int(z); return 0xFFFF if z>0xFFFF else (0 if z<0 else z)

def bbox_overlap(a,b):
    ox=min(a[2],b[2])-max(a[0],b[0]); oy=min(a[3],b[3])-max(a[1],b[1])
    return max(0.0,ox)*max(0.0,oy)

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

def score_group(draws):
    """clamp16 persistent-Z reject% over a group (dump order). Returns (covered, nreject, W, H)."""
    OX=int(min(d["xmin"] for d in draws)); OY=int(min(d["ymin"] for d in draws))
    UX=max(d["xmax"] for d in draws);       UY=max(d["ymax"] for d in draws)
    W=int(UX)-OX+1; H=int(UY)-OY+1
    if W<=0 or H<=0 or W*H>700000: return (0,0,W,H)   # skip degenerate/huge
    NPX=W*H
    zbuf=[-1]*NPX; z_owner=[-1]*NPX; paint_owner=[-1]*NPX; cov=bytearray(NPX)
    for k,d in enumerate(draws):
        for (v0,v1,v2) in tris_of(d["verts"],OX,OY):
            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(W-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):
                base=py*W
                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=base+px; cov[o]=1
                    zq=clamp16(w0*v0["z"]+w1*v1["z"]+w2*v2["z"])
                    paint_owner[o]=k
                    if zq>=zbuf[o]: zbuf[o]=zq; z_owner[o]=k
    covered=sum(cov); nreject=sum(1 for o in range(NPX) if cov[o] and z_owner[o]!=paint_owner[o])
    return (covered, nreject, W, H)

def main(argv):
    a=argv[1:]; dump=a[0] if a and not a[0].startswith("--") else None
    topk=int(a[a.index("--topk")+1]) if "--topk" in a else 12
    maxfb=a[a.index("--max-fb")+1] if "--max-fb" in a else "448x480"
    MW,MH=(int(x) for x in maxfb.split("x"))
    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");
        dump=c[0]
    print(f"[Zrank] dump={os.path.basename(dump)}  clamp16 depth-rejection ranking (fits<= {MW}x{MH})")

    draws,h,vram = C.census(dump, frame_filter=1, min_prims=8)
    cand=[]
    for d in draws:
        t0=d["tex0"]
        if not (t0["psm"]==0x13 and t0.get("tw")==512 and t0.get("th")==512): continue
        if d["prim"]["tme"]!=1 or d["prim"]["fst"]!=0: continue
        if not d["tex_resident"]: continue
        if not (d["clut"] and d["clut"]["resident"]): continue
        if d["onscreen_frac"]<0.5: continue
        zc=[clamp16(v["z"]) for v in d["verts"]]
        d["_funcl"]=sum(1 for z in zc if z<0xFFFF)/len(zc); d["_zmed"]=sorted(zc)[len(zc)//2]
        d["_bbox"]=(d["xmin"],d["ymin"],d["xmax"],d["ymax"]); d["_area"]=(d["xmax"]-d["xmin"])*(d["ymax"]-d["ymin"])
        if d["_funcl"]>=0.5: cand.append(d)
    print(f"[Zrank] {len(cand)} rich-depth supported candidates; textures={sorted(set(d['tex0']['tbp'] for d in cand))}")

    # per texture, keep a manageable set of the biggest/on-screen draws to bound the combinatorics
    by_tex={}
    for d in cand: by_tex.setdefault(d["tex0"]["tbp"], []).append(d)
    for tb in by_tex: by_tex[tb]=sorted(by_tex[tb], key=lambda d:-d["_area"])[:8]
    texs=sorted(by_tex)

    # enumerate distinct-texture triples with pairwise bbox overlap + a spread of depths; score with clamp16 raster
    seen=set(); scored=[]
    for t3 in itertools.combinations(texs,3):
        for a0 in by_tex[t3[0]]:
            for a1 in by_tex[t3[1]]:
                if bbox_overlap(a0["_bbox"],a1["_bbox"])<200: continue
                for a2 in by_tex[t3[2]]:
                    if bbox_overlap(a0["_bbox"],a2["_bbox"])<200 and bbox_overlap(a1["_bbox"],a2["_bbox"])<200: continue
                    grp=sorted([a0,a1,a2], key=lambda d:d["first_idx"])   # dump order = paint order
                    # distinct CLUTs required
                    if len({g["tex0"]["cbp"] for g in grp})<3: continue
                    key=tuple(g["first_idx"] for g in grp)
                    if key in seen: continue
                    seen.add(key)
                    OX=int(min(g["xmin"] for g in grp)); OY=int(min(g["ymin"] for g in grp))
                    W=int(max(g["xmax"] for g in grp))-OX+1; Hh=int(max(g["ymax"] for g in grp))-OY+1
                    if W>MW or Hh>MH: continue                            # must fit the target rung
                    cov,nrej,Wr,Hr=score_group(grp)
                    if cov<2000: continue
                    scored.append((nrej/max(cov,1), nrej, cov, key, [g["tex0"]["tbp"] for g in grp], Wr, Hr))
    scored.sort(reverse=True)
    print(f"[Zrank] scored {len(scored)} distinct-texture overlapping triples (fit rung). TOP {topk} by reject-fraction:")
    print(f"        {'reject%':>8} {'reject':>7} {'covered':>8} {'FB':>9}  draws (dump order) / textures")
    for frac,nrej,cov,key,tbps,Wr,Hr in scored[:topk]:
        print(f"        {100*frac:7.2f}% {nrej:>7} {cov:>8} {Wr:>4}x{Hr:<4}  {list(key)}  tbp={tbps}")
    # reference: the current trio
    trio=[d for d in draws if d["first_idx"] in (11671,19562,89761)]
    if len(trio)==3:
        cov,nrej,Wr,Hr=score_group(sorted(trio,key=lambda d:d["first_idx"]))
        print(f"[Zrank] current regression trio [11671,19562,89761]: {100*nrej/max(cov,1):.2f}% ({nrej}/{cov}) FB {Wr}x{Hr}")
    return 0

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