# Before and after, in real WebGL: the glow and ring shaders exactly as they are in two versions of
# public/ideas/starlock/cutscene3d.js (and the painterly pass from the matching reveal3d.js), drawn on
# a test scene by webgl.mjs.txt, read back, and measured against the float ideal (experiment.py).
#   python3 beforeafter.py <before-starlock-dir> <after-starlock-dir> [out-dir]
# Prints the table; with out-dir, also writes crops (PNG) of each, brightened 4x so the steps show.
import sys, os, re, json, subprocess, tempfile, zlib, struct
import numpy as np
import experiment as E

def shaders(dir_):
    src = open(os.path.join(dir_, 'cutscene3d.js')).read()
    dither = re.search(r'const DITHER = `(.*?)`;', src, re.S).group(1)
    def frag(name):
        body = src[src.index(f'const {name} = {{'):]
        body = body[body.index('fragmentShader: `') + len('fragmentShader: `'):]
        return body[:body.index('`,')].replace('${DITHER}', dither)
    return frag('glowShader'), frag('ringShader')

def render_gl(dir_, scene, R, dpr, origin):
    glow, ring = shaders(dir_)
    W, H = int(origin[0] + E.FRAME), int(origin[1] + E.FRAME)
    quads = []
    for k, (kind, cx, cy, w, h, col, op, extra) in enumerate(scene):
        u = dict(uColor=list(col), uOpacity=op, uSeed=k)
        if kind == 'glow': u['uCore'] = extra
        else: u['uWidth'], u['uFill'] = extra
        quads.append(dict(frag=glow if kind == 'glow' else ring, uniforms=u, x=cx + origin[0], y=cy + origin[1], w=w, h=h))
    S = tempfile.mkdtemp()
    json.dump(dict(W=W, H=H, radius=R, stride=dpr, quads=quads), open(S + '/s.json', 'w'))
    here = os.path.dirname(os.path.abspath(__file__))
    subprocess.run(['node', '--input-type=module', '-', S + '/s.json', S + '/s'], stdin=open(os.path.join(here, 'webgl.mjs.txt')),
                   check=True, stderr=subprocess.DEVNULL, env=dict(os.environ, STARLOCK=os.path.abspath(dir_)), cwd=here)
    c = np.fromfile(S + '/s.canvas.rgba', np.uint8).reshape(H, W, 4)[..., :3].astype(float)
    return c[int(origin[1]):, int(origin[0]):]   # the frame, rows bottom-up

def png(path, rgb):
    rgb = np.flipud(np.clip(rgb, 0, 255).astype(np.uint8)); h, w, _ = rgb.shape
    raw = b''.join(b'\0' + rgb[i].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, 2, 0, 0, 0)) + chunk(b'IDAT', zlib.compress(raw, 9)) + chunk(b'IEND', b''))

if __name__ == '__main__':
    before, after = sys.argv[1], sys.argv[2]
    out = sys.argv[3] if len(sys.argv) > 3 else None
    for sname in ('stack', 'concentric', 'faint'):
        sc = E.scenes()[sname]
        for R, dpr in ((5, 2), (2, 2), (0, 2)):
            I = E.ideal(sc, R, dpr)
            row = []
            for label, d in (('before', before), ('after', after)):
                o = render_gl(d, sc, R, dpr, E.ORIGIN)
                m = E.metrics(o, I, dpr)
                row.append(f"{label} band {m['band']:.3f} grain {m['grain']:.3f} steps {m['steps']:6.0f}")
                if out and R in (5, 0):
                    png(os.path.join(out, f'{sname}-r{R}-{label}.png'), o[96:416, 320:640] * 4)
            print(f'{sname:10s} radius {R}, pixel ratio {dpr}:  ' + '  |  '.join(row))