# Data for the page's lattice explorer: Starlock's pass per tap (taps.json, for computing H and the
# error live in the browser), and two small greyscale PNGs over the torus of (a, b) in [0, 1)^2:
# hmap.png, |H| (dark where the filter lets little through), and jmap.png, each lattice's RMS error
# (dark is good), 251 x 251 (a prime, so every point is a lattice with 251 levels), both centred on zero:
# pixel (i, j) is the frequency or lattice ((i - 125) / 251, (125 - j) / 251). python3 maps.py <out-dir>
import sys, json, zlib, struct
import numpy as np
import filters as F, lattice as L

def png(path, g):  # 8-bit greyscale, rows top to bottom
    h, w = g.shape; raw = b''.join(b'\0' + g[i].astype(np.uint8).tobytes() for i in range(h))
    chunk = lambda t, d: struct.pack('>I', len(d)) + t + d + struct.pack('>I', zlib.crc32(t + d) & 0xffffffff)
    open(path, 'wb').write(b'\x89PNG\r\n\x1a\n' + chunk(b'IHDR', struct.pack('>IIBBBBB', w, h, 8, 0, 0, 0, 0)) + chunk(b'IDAT', zlib.compress(raw, 9)) + chunk(b'IEND', b''))

out = sys.argv[1] if len(sys.argv) > 1 else '.'
t = F.kuwahara(5); n = 251
json.dump([[dx, dy, round(w, 8)] for dx, dy, w in t], open(f'{out}/taps.json', 'w'))
centre = lambda a: np.flipud(np.roll(np.roll(a, n // 2, 0), n // 2, 1))   # zero in the middle, b up
H = np.sqrt(F.grid_response(t, n))                       # [fy, fx]
png(f'{out}/hmap.png', centre(np.clip((H - 0.12) / (0.6 - 0.12), 0, 1) ** 0.5 * 255))
Jm = np.sqrt(L.search(F.grid_response(t, n), n))          # [B, A]; only (0, 0) is not a lattice
Jm[~np.isfinite(Jm)] = 0.4
png(f'{out}/jmap.png', centre(np.clip((Jm - 0.058) / (0.11 - 0.058), 0, 1) ** 0.7 * 255))
print('ok')