"""M4c-raster -- the per-draw rasterization moved onto the GPU.

Same faithful pipeline as m4b_frames.py, but the O(quads x 832 x 512) per-poly
rasterization (the CPU bottleneck, ~5.8 s/frame) runs in one GPU compute
dispatch: the shader finds each pixel's nearest poly (edge test + z-test) and
emits that poly's evaluated texz/texu/texv planes + tid. The verified M5 texel
decode (perspective divide + real-texture fetch) then runs as ONE vectorized
numpy post-pass over the whole frame (O(ntex x W x H)), not per poly.

Fidelity note: the single-pass GPU z-test picks the nearest poly regardless of
the near-black texture CUTOUT, so cut texels at a poly's edge show that poly's
hole rather than the poly behind (the CPU path defers z past the cutout). This
affects only keyed-texel boundary pixels (emblems/labels); the opaque bulk --
floor, ceiling, walls -- is identical. Documented, minor.

  python m4b_gpu.py <capture.fifodump> [max_frames]
"""
import sys, os, time, struct
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE); sys.path.insert(0, os.path.dirname(HERE))
sys.path.insert(0, r'C:\VWE\TeslaRel410\dpl3-revive\patha')
import emu860, emu_main, emu860c, igc_gpu
import numpy as np
from PIL import Image
from driver import boot, CpuShim
from vrboard import A
from texstore import build_texstore
from dpl_sampler import mode_to_texflags, wrap_index, composite
ANAME = {int(a): a.name for a in A}
emu860.Mem.log = lambda self, *a, **k: None

SRC = sys.argv[1]
MAXF = int(sys.argv[2]) if len(sys.argv) > 2 else 12
PROG_LO, PROG_HI = 0x08158000, 0x08170000
W, H = 832, 512
EDGE = {0x42, 0x0d, 0x2c}

# per-draw raster: nearest poly's texz/texu/texv planes + tid, one dispatch.
FX_SHADER = """
#version 430
#extension GL_ARB_gpu_shader_fp64 : enable
layout(local_size_x = 64, local_size_y = 1) in;
layout(std430, binding = 0) readonly buffer Prims { double prims[]; };
// 10 dvec4 / prim (40 doubles):
//   [0..5] edges 0..5: A B C used   (ALL edges, up to 6 -- matches the CPU clip)
//   [6] zdepth.A B C  nedges   (op 0x21 depth plane)
//   [7] u.A B C       hasuv     (op 0x43 ad 58)
//   [8] v.A B C       tid       (op 0x43 ad 78)
//   [9] tz.A B C      0         (op 0x43 ad 32, perspective denom)
layout(std430, binding = 1) buffer Outs { uint outbuf[]; };  // 6 words/pixel
uniform int n_prims;
uniform int width;
// double-precision plane eval, to MATCH the float64 CPU reference exactly.
void main() {
    int x = int(gl_GlobalInvocationID.x);
    int y = int(gl_GlobalInvocationID.y);
    if (x >= width) return;
    double fx = double(x), fy = double(y);
    double zbest = -1e30lf;
    double uu = 0.0lf, vv = 0.0lf, tz = 1.0lf;
    float tid = 0.0, hasuv = 0.0;
    for (int p = 0; p < n_prims; p++) {
        int b = p*40;
        bool allpos = true, allneg = true;
        for (int e = 0; e < 6; e++) {
            int eb = b + e*4;
            if (prims[eb+3] < 0.5lf) continue;
            double ve = prims[eb]*fx + prims[eb+1]*fy + prims[eb+2];
            if (ve >= 0.0lf) allneg = false; else allpos = false;
        }
        if (!(allpos || allneg)) continue;
        double z = prims[b+24]*fx + prims[b+25]*fy + prims[b+26];
        if (z <= zbest) continue;
        zbest = z;
        hasuv = float(prims[b+31]);
        uu = prims[b+28]*fx + prims[b+29]*fy + prims[b+30];
        vv = prims[b+32]*fx + prims[b+33]*fy + prims[b+34];
        tz = prims[b+36]*fx + prims[b+37]*fy + prims[b+38];
        tid = float(prims[b+35]);
    }
    int pix = y*width + x;
    // winner chosen in double (matches CPU); planes shipped as float32 for the
    // numpy float64 divide (sub-texel vs the CPU's float64 -- winner is what
    // mattered for the 18% region divergence).
    outbuf[pix*6+0] = (zbest > -1e29lf) ? 0x3f800000u : 0xff800000u; // 1.0 : -inf-ish
    outbuf[pix*6+1] = floatBitsToUint(float(uu));
    outbuf[pix*6+2] = floatBitsToUint(float(vv));
    outbuf[pix*6+3] = floatBitsToUint(float(tz));
    outbuf[pix*6+4] = floatBitsToUint(tid);
    outbuf[pix*6+5] = floatBitsToUint(hasuv);
}
"""

# --- wire -> queue + texture store ------------------------------------------
data = open(SRC, 'rb').read()
recs = []; off = 0
while True:
    i = data.find(b'VPXM', off)
    if i < 0:
        break
    ln = struct.unpack_from('<I', data, i + 4)[0]
    body = data[i + 8:i + 8 + ln]; off = i + 8 + ln
    if len(body) >= 4:
        a = struct.unpack_from('<I', body, 0)[0]
        if a < 0x100:
            recs.append((a, body[4:]))
print("queued %d records" % len(recs), flush=True)
store = build_texstore(recs)
texlist = list(store.items())
print("%d textures decoded" % len(texlist), flush=True)

r = boot(fw='vrend410', queue=recs)
shim = CpuShim()
RECV = emu_main.MAPS['vrend410']['receive']
g = igc_gpu.GpuTile()
ctx = g.ctx
prog_gl = ctx.compute_shader(FX_SHADER)
print("GPU:", ctx.info['GL_RENDERER'], flush=True)
# grid for the post-pass divide
yy, xx = np.mgrid[0:H, 0:W].astype(np.float64)


def f32(w):
    return struct.unpack('<f', struct.pack('<I', w & 0xffffffff))[0]


def build_prims():
    """Reconstruct per-draw primitive records from the LIVE program in C mem."""
    prog = {a: emu860c.r32(a) for a in range(PROG_LO, PROG_HI, 4) if emu860c.r32(a)}

    def rd(a):
        return prog.get(a, 0)

    setups = [a for a in sorted(prog) if rd(a) == 0x100]
    recs4 = []
    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
        zp = tzp = up = vp = None; tid = None; q = p
        for _ in range(70):
            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 == 32 and tzp is None:
                tzp = (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 op == 0xf7 and ad in (141, 142) and tid is None:
                tid = (rd(q + 4) >> 2) & 0x3f      # identical rule to render_faithful
            q += 1 * 4
        if zp is None:
            continue
        hasuv = 1.0 if (up and vp and tzp) else 0.0
        up = up or (0, 0, 0); vp = vp or (0, 0, 0); tzp = tzp or (0, 0, 1)
        eslots = [(e[0], e[1], e[2], 1.0) for e in edges[:6]]   # ALL edges (clip)
        while len(eslots) < 6:
            eslots.append((0.0, 0.0, 0.0, 0.0))
        recs4 += eslots + [(*zp, float(len(edges))), (*up, hasuv),
                           (*vp, float(tid or 0)), (*tzp, 0.0)]
    if not recs4:
        return None
    return np.array(recs4, dtype=np.float64).reshape(-1, 4)   # double planes, 10/prim


def render_gpu(nth):
    prims = build_prims()
    if prims is None:
        return False, 0
    n_prims = len(prims) // 10          # 10 vec4 (40 doubles) per prim
    b0 = ctx.buffer(prims.tobytes())
    out = ctx.buffer(reserve=W * H * 6 * 4)
    b0.bind_to_storage_buffer(0)
    out.bind_to_storage_buffer(1)
    prog_gl['n_prims'] = n_prims
    prog_gl['width'] = W
    prog_gl.run(group_x=(W + 63) // 64, group_y=H)
    raw = np.frombuffer(out.read(), dtype=np.uint32).reshape(H, W, 6)
    b0.release(); out.release()
    zb = raw[..., 0].view(np.float32)
    uu = raw[..., 1].view(np.float32).astype(np.float64)
    vv = raw[..., 2].view(np.float32).astype(np.float64)
    tz = raw[..., 3].view(np.float32).astype(np.float64)
    tid = raw[..., 4].view(np.float32).astype(np.int64)
    hasuv = raw[..., 5].view(np.float32)
    filled = zb > -1e29
    img = np.zeros((H, W, 3), np.uint8)
    img[filled & (hasuv <= 0.5)] = (60, 60, 70)          # flat polys
    tzc = np.where(np.abs(tz) < 1.0, np.sign(tz) + (tz == 0), tz)
    # VERIFIED divide, vectorized over the whole frame
    ucoord = (uu / tzc * 2048).astype(np.int64)
    vcoord = (vv / tzc * 2048).astype(np.int64)
    ntex = len(texlist)
    for slot in range(ntex):
        sel = filled & (hasuv > 0.5) & ((tid % ntex) == slot)
        if not sel.any():
            continue
        h, (tu, tv, mode, arr) = texlist[slot]
        wrap_u, wrap_v, cut_mode = mode_to_texflags(mode)
        ui = wrap_index((ucoord * tu) >> 8, tu, wrap_u)
        vi = wrap_index((vcoord * tv) >> 8, tv, wrap_v)
        samp = arr[vi, ui]
        composite(img, sel, samp, cut_mode)
    Image.fromarray(img, 'RGB').save(os.path.join(HERE, 'gpu_frame_%04d.png' % nth))
    return True, n_prims


frames = 0
prev_draw = False
t0 = time.time()
render_t = 0.0
while frames < MAXF and r.qi < len(r.queue):
    reason, _ = emu860c.run(500_000_000)
    if reason == 3:
        continue
    if reason != 0:
        print("core reason %d" % reason, flush=True); break
    pc = emu860c.getstate()['pc']
    h = r.hooks.get(pc)
    if h is None:
        print("sentinel", flush=True); break
    if pc == RECV:
        if prev_draw:
            rt0 = time.time()
            ok, np_ = render_gpu(frames)
            render_t += time.time() - rt0
            if ok:
                print("frame %d: %d prims (cmd %d)" % (frames, np_, r.qi), flush=True)
                frames += 1
            prev_draw = False
        if r.qi < len(r.queue) and ANAME.get(r.queue[r.qi][0]) == 'draw_scene':
            prev_draw = True
    if h(shim) == 'done':
        break
wall = time.time() - t0
print("done: %d GPU frames in %.1fs (%.1fs firmware, %.1fs render), cmd %d"
      % (frames, wall, wall - render_t, render_t, r.qi), flush=True)