"""Extract Kaleidostrobe's deterministic tables by running the original code.

- 22 colour maps (makeMap000-021): mapRGB[0], 256 RGBA entries
- base wave images (initWaveShape000-022), *before* their random transform()
  (transform is stubbed out here; the port applies it with its own rand()),
  plus the gWaveIncScale each one sets
Every call is checked to touch nothing but its output.
"""
import struct, sys
import numpy as np
sys.path.insert(0, __file__.rsplit('/', 1)[0])
import macho
from plugin import Plugin
from unicorn import UC_HOOK_CODE

MAPRGB = 0x23fd4
TEXIDX = 0x457f8          # textureImageIndexArrays
CURTEX = 0x13dbc
CURFRAME = 0x13dc0
WAVEINC = 0x257f0         # gWaveIncScale (2 floats)

p = Plugin(seed=1)
k = p.k
tr = macho.func_size['_transform'][0]
k.u.hook_add(UC_HOOK_CODE, lambda u, a, s, d: k._ret(), begin=tr, end=tr + 3)
k.u.ctl_remove_cache(0, 0x80000)

def snapshot():
    return bytearray(k.mem(0x13000, 0x759f8 - 0x13000))

def changed(before, after):
    b = np.frombuffer(bytes(before), np.uint8); a = np.frombuffer(bytes(after), np.uint8)
    idx = np.nonzero(a != b)[0] + 0x13000
    return idx

def regions(idx):
    out = set()
    for a in idx[::max(1, len(idx)//200)]:
        out.add(macho.symfor(int(a)).split('+')[0])
    return sorted(out)

k.w32(CURTEX, 0)
k.w32(CURFRAME, 0); k.w32(CURFRAME + 4, 0)

maps = np.zeros((22, 256, 4), np.uint8)
for m in range(22):
    k.u.mem_write(MAPRGB, b'\x55' * 1024)
    before = snapshot()
    k.call(f'_makeMap{m:03d}')
    after = snapshot()
    idx = changed(before, after)
    stray = idx[(idx < MAPRGB) | (idx >= MAPRGB + 1024)]
    data = np.frombuffer(k.mem(MAPRGB, 1024), np.uint8).reshape(256, 4)
    untouched = int((np.frombuffer(k.mem(MAPRGB, 1024), np.uint8) == 0x55).sum())
    maps[m] = data
    print(f'map {m:03d}: stray={regions(stray) if len(stray) else "-"} 0x55-left={untouched}')

WAVES = list(range(23)) + [31, 32, 33]
waves = {}
incs = {}
for w in WAVES:
    if w >= 31:
        continue
    for fill in (0x00, 0xff):
        k.u.mem_write(TEXIDX, bytes([fill]) * 0x8000)
        k.wf(WAVEINC, -999.0)
        before = snapshot()
        k.call(f'_initWaveShape{w:03d}')
        after = snapshot()
        img = np.frombuffer(k.mem(TEXIDX, 0x4000), np.uint8).reshape(128, 128).copy()
        if fill == 0:
            img0 = img
        else:
            img1 = img
        idx = changed(before, after)
        stray = idx[((idx < TEXIDX) | (idx >= TEXIDX + 0x4000)) & ((idx < WAVEINC) | (idx >= WAVEINC + 4))]
    partial = int((img0 != img1).sum())
    waves[w] = img0
    incs[w] = k.rf(WAVEINC)
    print(f'wave {w:03d}: inc={incs[w]!r} partial-pixels={partial} stray={regions(stray) if len(stray) else "-"}')

np.savez_compressed(sys.argv[1] if len(sys.argv) > 1 else 'tables.npz', maps=maps,
                    waves=np.stack([waves[w] for w in range(23)]),
                    wave_inc=np.array([incs[w] for w in range(23)], np.float32))
print('saved')