#!/usr/bin/env python3
"""debug_ident.py -- render a capture frame with each instance in a unique flat
color (no materials/textures) to identify what occupies the screen."""
import os, sys, struct
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
os.environ.setdefault('SDL_VIDEODRIVER', 'dummy')
import numpy as np
from vrboard import VirtualBoard, Assembler
from decode_anim import find_framed_start
import vrview

path = sys.argv[1] if len(sys.argv) > 1 else 'klng.raw.bin'
upto = int(sys.argv[2]) if len(sys.argv) > 2 else 900
out = sys.argv[3] if len(sys.argv) > 3 else 'ident.png'

data = open(path, 'rb').read()
board = VirtualBoard(); asm = Assembler(); asm.feed(data[find_framed_start(data):])
for m in asm:
    try: board.handle(m)
    except Exception: pass
    if board.frames >= upto: break

r = vrview.Renderer()
c = r.cache; c.maybe_rebuild(board)
V = np.linalg.inv(r.chain_matrix(board, c.cam_chain))
W, H = r.w, r.h
img = np.zeros((H, W, 3), np.float32)
zbuf = np.full((H, W), np.inf, np.float32)
vp = {'x0': -1, 'y0': -0.625, 'x1': 1, 'y1': 0.625, 'zeye': 1.0, 'hither': 1.0}
import colorsys
legend = []
for i, inst in enumerate(c.instances):
    M = r.chain_matrix(board, inst['chain']) @ V
    R, T = M[:3, :3], M[3, :3]
    col3 = np.array(colorsys.hsv_to_rgb((i * 0.37) % 1.0, 0.9, 1.0)) * 255
    drew = False
    for gh in inst['geoms']:
        mesh = c._mesh(board, gh)
        if mesh is None: continue
        pe = mesh['pos'] @ R + T
        z = -pe[:, 2]
        ok = z > vp['hither']
        if not ok.any(): continue
        zs = np.where(ok, z, 1.0)
        px = (pe[:, 0] * vp['zeye'] / zs - vp['x0']) / 2 * W
        py = (1 - (pe[:, 1] * vp['zeye'] / zs - vp['y0']) / 1.25) * H
        col = np.repeat(col3[None, :], len(pe), 0)
        r._raster(mesh['tri'], px, py, z, ok, col, None, None, None, None,
                  np.zeros(3), img, zbuf, W, H)
        drew = True
    if drew:
        legend.append((i, inst['chain'][0], [int(x) for x in col3]))

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'parser'))
import svt
rgba = np.zeros((H, W, 4), np.uint8); rgba[..., 3] = 255
rgba[..., :3] = np.clip(img, 0, 255).astype(np.uint8)
svt.write_png(out, W, H, rgba.tobytes())
print('wrote', out)

# which instances actually own visible pixels?
vis = {}
flat = rgba[..., :3].reshape(-1, 3)
for i, dcs, col in legend:
    n = int(np.all(np.abs(flat.astype(int) - col) < 12, axis=1).sum())
    if n > 200:
        vis[i] = (hex(dcs), col, n)
for i, (dcs, col, n) in sorted(vis.items(), key=lambda kv: -kv[1][2]):
    print(f'instance {i:3} dcs {dcs:>6} color={col} pixels={n}')