#!/usr/bin/env python3
"""
Diff demos/eyeliner-studio/strip_bake.js against the two Python tools it ports.

The studio bakes in the browser so an artist never has to leave the page, which
means the same reprojection now exists twice. That is only acceptable while the
two agree: a strip baked in the browser and one baked at the prompt have to land
the same artwork at the same place on the same face. This is what checks that.

Run from the repo root, after touching EITHER side:

    python3 tools/check_strip_bake.py

Needs numpy + Pillow (like the other tools) and node. It shells out to
tools/bake_eye_strip.py and tools/make_liner_strip.py for the reference, drives
strip_bake.js through node for the candidate, and compares the two pixel by
pixel. Exit status is non-zero if anything drifts past the tolerance.

TOLERANCE, and why it is not zero. bakeAtlasToStrip() collapses supersamples and
averages the two eyes in a different ORDER from the numpy version — numpy holds
the whole supersampled buffer and means over two axes at the end, the JS
accumulates per column to keep browser memory down. Both are the same arithmetic
in real numbers and differ only in float rounding, which shows up as at most one
8-bit level on a handful of texels (measured: 16 of 131072 at worst, 0.012%).
makeLinerStrip() has no such reordering and is expected to be EXACT.
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile

import numpy as np
from PIL import Image

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
STRIP_BAKE = os.path.join(ROOT, 'demos', 'eyeliner-studio', 'strip_bake.js')
ATLAS = os.path.join(ROOT, 'demos', 'assets', 'imgs', 'face', 'color_map_eye_liner.png')

# (label, bake_eye_strip.py flags, strip_bake.js options)
BAKE_CASES = [
    ('upper mask smooth6', '--emit mask --smooth-u 6', 'region=eyeLidUpper emit=mask smoothU=6'),
    ('upper mask raw', '--emit mask', 'region=eyeLidUpper emit=mask'),
    ('upper colour raw', '--emit color', 'region=eyeLidUpper emit=color'),
    ('lower mask', '--emit mask --region lower', 'region=eyeLidLower emit=mask'),
    ('lower mask smooth3', '--emit mask --region lower --smooth-u 3',
     'region=eyeLidLower emit=mask smoothU=3'),
    ('upper left eye only', '--emit mask --eye left', 'region=eyeLidUpper emit=mask eye=left'),
    ('upper right, stroke .2', '--emit mask --eye right --stroke 0.2',
     'region=eyeLidUpper emit=mask eye=right stroke=0.2'),
    ('upper supersample 1', '--emit mask --supersample 1',
     'region=eyeLidUpper emit=mask supersample=1'),
    ('upper ss3, wing .36', '--emit mask --supersample 3 --wing-length 0.36',
     'region=eyeLidUpper emit=mask supersample=3 wingLength=0.36'),
    ('upper 2048x256', '--emit mask --size 2048 256',
     'region=eyeLidUpper emit=mask width=2048 height=256'),
    ('upper 64 cols, cu .8', '--emit mask --columns 64 --corner-u 0.8',
     'region=eyeLidUpper emit=mask columns=64 cornerU=0.8'),
    ('lower wing .3 lift 1', '--emit mask --region lower --wing-length 0.3 --wing-lift 1.0',
     'region=eyeLidLower emit=mask wingLength=0.3 wingLift=1.0'),
]

LINER_CASES = [
    ('lower default', '--region lower', 'region=eyeLidLower'),
    ('upper default', '--region upper', 'region=eyeLidUpper'),
    ('upper cat-eye',
     '--region upper --thick-inner 0.10 --thick-outer 1.0 --start-u 0.06 --end-u 0.80 '
     '--wing-taper 0.90 --edge 0.10',
     'region=eyeLidUpper thickInner=0.10 thickOuter=1.0 startU=0.06 endU=0.80 '
     'wingTaper=0.90 edge=0.10'),
    ('lower smudge',
     '--region lower --thick-inner 0.25 --thick-outer 0.75 --start-u 0.10 --edge 0.55',
     'region=eyeLidLower thickInner=0.25 thickOuter=0.75 startU=0.10 edge=0.55'),
    ('upper 2048x256', '--region upper --size 2048 256',
     'region=eyeLidUpper width=2048 height=256'),
    ('lower corner-u .8', '--region lower --corner-u 0.8', 'region=eyeLidLower cornerU=0.8'),
]

# node driver, written to the temp dir so the import path is absolute and node
# does not need the repo's package.json to resolve anything
DRIVER = '''
import {readFileSync, writeFileSync} from 'fs';
import {bakeAtlasToStrip, makeLinerStrip, parseGeometry, CentripetalCatmullRom, RINGS}
  from %(module)s;

const [kind, outBin, ...kv] = process.argv.slice(2);
const opts = {};
for (const p of kv) {
  const [k, v] = p.split('=');
  opts[k] = isNaN(Number(v)) ? v : Number(v);
}

if (kind === 'curve') {
  // the one thing this file re-implements rather than imports: three's
  // CatmullRomCurve3 with curveType 'centripetal'
  const THREE = await import(%(three)s);
  const geo = parseGeometry(readFileSync(%(geometry)s, 'utf8'));
  const V = geo.vertices;
  let worst = 0;
  for (const region of ['eyeLidUpper', 'eyeLidLower']) {
    for (let e = 0; e < 2; e++) {
      const pts = RINGS[region][e].lash.map(i => [V[i*3], V[i*3+1], V[i*3+2]]);
      const mine = new CentripetalCatmullRom(pts);
      const theirs = new THREE.CatmullRomCurve3(
        pts.map(p => new THREE.Vector3(...p)), false, 'centripetal');
      for (let k = 0; k <= 2000; k++) {
        const t = k / 2000;
        for (const fn of ['getPoint', 'getPointAt']) {
          const a = mine[fn](t), b = theirs[fn](t, new THREE.Vector3());
          worst = Math.max(worst, Math.abs(a[0]-b.x), Math.abs(a[1]-b.y), Math.abs(a[2]-b.z));
        }
      }
    }
  }
  console.log(JSON.stringify({worst}));
} else if (kind === 'liner') {
  const out = makeLinerStrip(opts);
  writeFileSync(outBin, Buffer.from(out.data.buffer, out.data.byteOffset, out.data.length));
} else {
  const raw = readFileSync(opts.src);
  const source = {data: new Uint8ClampedArray(raw), width: opts.sw, height: opts.sh};
  delete opts.src; delete opts.sw; delete opts.sh;
  const geometry = parseGeometry(readFileSync(%(geometry)s, 'utf8'));
  const out = bakeAtlasToStrip({source, geometry, ...opts});
  writeFileSync(outBin, Buffer.from(out.data.buffer, out.data.byteOffset, out.data.length));
}
'''


def rgba(path):
    return np.asarray(Image.open(path).convert('RGBA'), dtype=np.int32)


def compare(label, ref_png, cand_bin):
    ref = rgba(ref_png)
    h, w = ref.shape[:2]
    cand = np.frombuffer(open(cand_bin, 'rb').read(), dtype=np.uint8)
    if cand.size != h * w * 4:
        return label, None, f'size mismatch: {cand.size} bytes for {w}x{h}'
    cand = cand.astype(np.int32).reshape(h, w, 4)
    d = np.abs(ref - cand)
    n = int((d.max(axis=2) > 0).sum())
    return label, int(d.max()), f'{w}x{h}  {n} of {w * h} texels differ ({100.0 * n / (w * h):.4f}%)'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--tolerance', type=int, default=1,
                    help='largest 8-bit difference accepted on the bake (default 1); '
                         'the parametric generator is always required to be exact')
    args = ap.parse_args()

    tmp = tempfile.mkdtemp(prefix='strip_bake_check.')
    driver = os.path.join(tmp, 'driver.mjs')
    open(driver, 'w').write(DRIVER % {
        'module': json.dumps(STRIP_BAKE),
        'geometry': json.dumps(os.path.join(ROOT, 'src', 'geometry.ts')),
        'three': json.dumps(os.path.join(ROOT, 'node_modules', 'three', 'build', 'three.module.js')),
    })

    src_bin = os.path.join(tmp, 'src.bin')
    atlas = rgba(ATLAS)
    open(src_bin, 'wb').write(atlas.astype(np.uint8).tobytes())
    sh, sw = atlas.shape[:2]

    def node(*argv):
        return subprocess.run(['node', driver, *argv], capture_output=True, text=True)

    failures = []

    # --- the ported curve, against three's own -----------------------------
    print('CentripetalCatmullRom vs three.CatmullRomCurve3(centripetal):')
    r = node('curve', '-')
    if r.returncode != 0:
        print('  SKIPPED — ' + (r.stderr.strip().splitlines() or ['node failed'])[-1][:160])
    else:
        worst = json.loads(r.stdout)['worst']
        ok = worst < 1e-12
        print(f'  max |difference| over 4 rings x 2001 samples x 2 methods: {worst:.3e}'
              f'  {"OK" if ok else "DRIFTED"}')
        if not ok:
            failures.append('curve')

    ref = os.path.join(tmp, 'ref.png')
    cand = os.path.join(tmp, 'cand.bin')

    print(f'\nbakeAtlasToStrip vs tools/bake_eye_strip.py  (tolerance {args.tolerance})')
    for label, pyflags, jsopts in BAKE_CASES:
        p = subprocess.run(
            ['python3', os.path.join(ROOT, 'tools', 'bake_eye_strip.py'), ATLAS,
             *pyflags.split(), '--out', ref],
            capture_output=True, text=True, cwd=ROOT)
        if p.returncode != 0:
            print(f'  {label:24s} PYTHON FAILED: {p.stderr.strip()[-160:]}')
            failures.append(label)
            continue
        r = node('bake', cand, f'src={src_bin}', f'sw={sw}', f'sh={sh}', *jsopts.split())
        if r.returncode != 0:
            print(f'  {label:24s} NODE FAILED: {r.stderr.strip()[-160:]}')
            failures.append(label)
            continue
        _, mx, detail = compare(label, ref, cand)
        bad = mx is None or mx > args.tolerance
        print(f'  {label:24s} max|diff| {"?" if mx is None else mx:>3}  {detail}'
              f'{"   DRIFTED" if bad else ""}')
        if bad:
            failures.append(label)

    print('\nmakeLinerStrip vs tools/make_liner_strip.py  (must be exact)')
    for label, pyflags, jsopts in LINER_CASES:
        p = subprocess.run(
            ['python3', os.path.join(ROOT, 'tools', 'make_liner_strip.py'),
             *pyflags.split(), '--out', ref],
            capture_output=True, text=True, cwd=ROOT)
        if p.returncode != 0:
            print(f'  {label:24s} PYTHON FAILED: {p.stderr.strip()[-160:]}')
            failures.append(label)
            continue
        r = node('liner', cand, *jsopts.split())
        if r.returncode != 0:
            print(f'  {label:24s} NODE FAILED: {r.stderr.strip()[-160:]}')
            failures.append(label)
            continue
        _, mx, detail = compare(label, ref, cand)
        bad = mx is None or mx > 0
        print(f'  {label:24s} max|diff| {"?" if mx is None else mx:>3}  {detail}'
              f'{"   DRIFTED" if bad else ""}')
        if bad:
            failures.append(label)

    print()
    if failures:
        print(f'{len(failures)} case(s) drifted: ' + ', '.join(failures))
        print('The browser bake and the shell bake would place artwork differently. '
              'Fix before shipping either.')
        return 1
    print('strip_bake.js agrees with both Python tools on every case.')
    return 0


if __name__ == '__main__':
    sys.exit(main())