"""Honest battle geometry: 0x2c/0x42/0x0d edges as proper half-planes with
polygon-level winding (inside = all>=0 OR all<0), real z-buffer. The scene is
TEXTURED (texz/texu/texv planes, no Gouraud colour), so with no TXDN texels we
show STRUCTURE: flat distinct colour per polygon + depth + u/v. numpy, offline."""
import pickle, struct, os
import numpy as np
from PIL import Image
HERE = os.path.dirname(os.path.abspath(__file__))
d = pickle.load(open(os.path.join(HERE, 'battle_prog.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)
W, H = 832, 512
setups = [a for a in sorted(prog) if rd(a) == 0x100]
EDGE = {0x42, 0x0d, 0x2c}

yy, xx = np.mgrid[0:H, 0:W].astype(np.float64)
zbuf = np.full((H, W), -1e30)
flat = np.zeros((H, W, 3), np.uint8)
uv = np.zeros((H, W, 3))
PAL = [(200,80,80),(80,200,120),(90,140,220),(220,200,90),(200,120,220),
       (120,220,220),(240,150,80),(150,120,230),(90,210,160),(230,110,150)]
drawn = 0
for a in setups:
    p = a + 4
    edges = []
    while ((rd(p) >> 8) & 0xff) in EDGE and len(edges) < 6:
        edges.append((f32(rd(p+4)), f32(rd(p+8)), f32(rd(p+12))))
        p += 16
    if len(edges) < 3:
        continue
    vsum_pos = np.ones((H, W), bool)
    vsum_neg = np.ones((H, W), bool)
    for A, B, C in edges:
        v = A * xx + B * yy + C
        vsum_pos &= v >= 0
        vsum_neg &= v < 0
    en = vsum_pos | vsum_neg
    if not en.any():
        continue
    # z = the op-0x21 within the next ~10 words
    zp = None; zq = p
    for _ in range(12):
        if ((rd(zq) >> 8) & 0xff) == 0x21:
            zp = (f32(rd(zq+4)), f32(rd(zq+8)), f32(rd(zq+12))); break
        zq += 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]
    flat[en] = PAL[drawn % len(PAL)]
    # texu/texv planes (addr 58 / 78 via op 0x43) for the uv view
    up = vp = None; q = zq
    for _ in range(50):
        w = rd(q)
        if ((w >> 8) & 0xff) == 0x43 and (w & 0xff) == 58 and up is None:
            up = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
        if ((w >> 8) & 0xff) == 0x43 and (w & 0xff) == 78 and vp is None:
            vp = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
        if up and vp: break
        q += 4
    if up and vp:
        u = up[0]*xx + up[1]*yy + up[2]
        v = vp[0]*xx + vp[1]*yy + vp[2]
        uv[..., 0][en] = u[en]; uv[..., 1][en] = v[en]
    drawn += 1

print("drawn %d polygons of %d setups; covered %d px" %
      (drawn, len(setups), int((zbuf > -1e29).sum())))
Image.fromarray(flat, 'RGB').save(os.path.join(HERE, 'battle_flat.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, np.ptp(zv[m]))
    zi[m] = (40 + 215*zn[m]).astype(np.uint8)
Image.fromarray(zi, 'L').save(os.path.join(HERE, 'battle_depth.png'))
# uv normalized globally (spatially coherent, unlike per-prim)
um, vm = uv[..., 0], uv[..., 1]
rgb = np.zeros((H, W, 3), np.uint8)
if m.any():
    un = (um - um[m].min())/max(1e-9, np.ptp(um[m]))
    vn = (vm - vm[m].min())/max(1e-9, np.ptp(vm[m]))
    rgb[..., 0][m] = (40 + 200*un[m]).astype(np.uint8)
    rgb[..., 1][m] = (40 + 200*vn[m]).astype(np.uint8)
    rgb[..., 2][m] = 120
Image.fromarray(rgb, 'RGB').save(os.path.join(HERE, 'battle_uv.png'))
print("wrote battle_flat.png + battle_depth.png + battle_uv.png")