"""Render fxtest's per-draw effect primitives from fx_program.pkl: per pair
{19w setup: edges+z, 80w ladder: texz/u/v planes}, rasterize with z-buffer and
visualize (u,v) as colour + z as depth. Structure render (no texel sampling)."""
import pickle, os, struct
import numpy as np
from PIL import Image
S = r'C:\Users\cyd\AppData\Local\Temp\claude\c--VWE-TeslaRel410\4e848c76-6e89-4034-8047-d8d491cb32d8\scratchpad'
d = pickle.load(open(os.path.join(S, 'fx_program.pkl'), 'rb'))
prog = d['prog']
def f32(w): return struct.unpack('<f', struct.pack('<I', w & 0xffffffff))[0]
def rd(a): return prog.get(a, 0)

# find setup blocks: SETENABS followed by 0x..2c00-family edge headers
setups = []
addrs = sorted(prog)
for a in addrs:
    if rd(a) == 0x100 and ((rd(a + 4) >> 8) & 0xff) == 0x2c:
        setups.append(a)
print("%d effect primitives" % len(setups))

W, H = 832, 512
yy, xx = np.mgrid[0:H, 0:W].astype(np.float64)
zbuf = np.full((H, W), -1e30)
rgb = np.zeros((H, W, 3))
drawn = 0
for a in setups:
    # edges: 4-word {hdr(0x2c), A, B, C} groups until a non-0x2c header
    en_p = np.ones((H, W), bool); en_n = np.ones((H, W), bool)
    p = a + 4
    nedges = 0
    while ((rd(p) >> 8) & 0xff) == 0x2c and nedges < 6:
        A, B, C = f32(rd(p+4)), f32(rd(p+8)), f32(rd(p+12))
        v = A * xx + B * yy + C
        en_p &= v >= 0; en_n &= v < 0
        p += 16; nedges += 1
    en = en_p | en_n
    if not en.any():
        continue
    # scan forward for the z instr (op 0x21) then the ladder (first texu/texv L3)
    zp = None; up = None; vp = None
    q = p
    for _ in range(40):
        w = rd(q); op = (w >> 8) & 0xff; ad = w & 0xff
        if op == 0x21 and zp is None:
            zp = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
        if op == 0x43 and ad == 58 and up is None:
            up = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
        if op == 0x43 and ad == 78 and vp is None:
            vp = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
        if up and vp and zp:
            break
        q += 4
    if zp is None:
        continue
    z = zp[0] * xx + zp[1] * yy + zp[2]
    en &= z > zbuf
    if not en.any():
        continue
    zbuf[en] = z[en]
    if up and vp:
        u = up[0] * xx + up[1] * yy + up[2]
        v = vp[0] * xx + vp[1] * yy + vp[2]
        # u/v as colour: hue-ish mapping, normalized per prim
        un = np.clip((u - u[en].min()) / max(1e-9, np.ptp(u[en])), 0, 1)
        vn = np.clip((v - v[en].min()) / max(1e-9, np.ptp(v[en])), 0, 1)
        rgb[..., 0][en] = 60 + 180 * un[en]
        rgb[..., 1][en] = 40 + 160 * vn[en]
        rgb[..., 2][en] = 220 - 140 * un[en]
    else:
        rgb[..., 0][en] = 200; rgb[..., 1][en] = 200; rgb[..., 2][en] = 200
    drawn += 1
print("drawn %d" % drawn)
img = np.clip(rgb, 0, 255).astype(np.uint8)
Image.fromarray(img, 'RGB').save(os.path.join(S, 'fx_render.png'))
zv = zbuf.copy(); m = zv > -1e29
zi = np.zeros((H, W), np.uint8)
if m.any():
    zn = (zv - zv[m].min()) / max(1e-9, (zv[m].max() - zv[m].min()))
    zi[m] = (40 + 215 * zn[m]).astype(np.uint8)
Image.fromarray(zi, 'L').save(os.path.join(S, 'fx_depth.png'))
print("wrote fx_render.png + fx_depth.png")