# SPDX-License-Identifier: Apache-2.0
"""Ink Bench test suite. Run: python3 tests/test_bench.py  (needs numpy, Pillow, node 20+)

1. Alignment recovers known synthetic shifts, scales and blurs.
2. Metrics match an independent plain numpy reference.
3. The browser core (lib/bench-core.js, run in Node) and tools/inkbench.py write identical bundles,
   across PNG 8/16 bit, RGB, palette, label maps and TIFF (8 bit, 16 bit, float32, LZW).
4. Pure noise scores low.
"""
import json, os, subprocess, sys, tempfile, time
import numpy as np
from PIL import Image
HERE = os.path.dirname(os.path.abspath(__file__)); ROOT = os.path.join(HERE, '..')
sys.path.insert(0, HERE); sys.path.insert(0, os.path.join(ROOT, 'tools'))
import synth, inkbench as ib
refs = ib.Refs(root=ROOT)
results = {'recovery': [], 'reference': [], 'parity': [], 'noise': []}
fails = []
def check(cond, msg):
    if not cond:
        fails.append(msg); print('FAIL', msg)

# 1. recovery
CASES = [('w046', 23, -17, 1, 1, 2), ('w048', -12.5, 9.25, 1.12, 1.12, 1), ('w030', 0, 0, 0.85, 0.85, 3),
         ('w040', 35, 20, 1.2, 1.18, 0), ('w050', -30, -25, 0.92, 0.95, 4), ('w035', 7.3, 3.6, 1.05, 1.05, 1.5),
         ('w055', -40, 31, 0.8, 0.8, 2.5), ('w027', 12, -8, 1.3, 1.3, 1), ('w044', -5.5, 14.2, 0.97, 1.03, 3.5),
         ('w051', 18, 22, 1.0, 1.0, 6)]
for seg, tx, ty, sx, sy, sg in CASES:
    t = time.time()
    b = ib.compare(synth.make(seg, tx, ty, sx, sy, sg), seg, refs=refs)
    a, m = b['alignment'], b['metrics']
    W = b['reference']['width']
    e_t = max(abs(a['tx'] - tx), abs(a['ty'] - ty)); e_s = max(abs(a['sx'] / sx - 1), abs(a['sy'] / sy - 1))
    e_b = abs(m['blurSigmaPx'] - sg)
    r = dict(seg=seg, truth=[tx, ty, sx, sy, sg], got=[a['tx'], a['ty'], a['sx'], a['sy'], m['blurSigmaPx']],
             shift_err_px=round(e_t, 3), scale_err_pct=round(100 * e_s, 3), blur_err_px=e_b,
             pearson=m['pearson'], pearson_matched_blur=m['pearsonMatchedBlur'], secs=round(time.time() - t, 1))
    results['recovery'].append(r); print('recovery', r)
    check(e_t <= 0.25, f'{seg} shift error {e_t}')
    check(e_s <= 0.003, f'{seg} scale error {e_s}')
    check(e_b <= 0.5, f'{seg} blur error {e_b}')
    check(m['pearsonMatchedBlur'] >= 0.99, f'{seg} matched-blur pearson {m["pearsonMatchedBlur"]}')

# 2. independent numpy reference for the metric definitions
def reference_metrics(ref, val, valid, tr, tp):
    a = ref[valid]; b = val[valid]
    out = {'pearson': float(np.corrcoef(a, b)[0, 1])}
    A = a >= tr; B = b >= tp
    out['iou'] = np.logical_and(A, B).sum() / np.logical_or(A, B).sum()
    out['dice'] = 2 * np.logical_and(A, B).sum() / (A.sum() + B.sum())
    out['precision'] = np.logical_and(A, B).sum() / B.sum(); out['recall'] = np.logical_and(A, B).sum() / A.sum()
    ts = np.arange(50, -1, -1) / 50
    P = np.array([(A & (b >= t)).sum() / max((b >= t).sum(), 1) if (b >= t).sum() else 1.0 for t in ts])
    R = np.array([(A & (b >= t)).sum() / A.sum() for t in ts])
    out['ap'] = float(np.sum(np.diff(np.concatenate([[0], R])) * P))
    return out
rng = np.random.default_rng(5)
for seg, tr, tp in [('w046', 0.5, 0.5), ('w030', 0.3, 0.6), ('w050', 0.7, 0.2), ('w048', 0.5, 0.35)]:
    ref, _ = refs.layer(seg, 'ink')
    H, W = ref.shape
    val = np.clip(ref * 0.7 + rng.normal(0, 0.15, ref.shape) + 0.1, 0, 1)
    valid = np.ones_like(ref, bool); valid[:, :W // 10] = False
    val = np.where(valid, val, 0)
    m = ib.metrics(ref, val, valid, tr, tp)
    r = reference_metrics(ref, val, valid, tr, tp)
    d = {k: abs(m[k] - r[k]) for k in r}
    results['reference'].append({'seg': seg, 'tRef': tr, 'tPred': tp, 'ours': {k: m[k] for k in r}, 'max_abs_diff': max(d.values())})
    print('reference', seg, {k: (m[k], round(float(r[k]), 6)) for k in r})
    check(max(d.values()) <= 1e-6, f'{seg} reference mismatch {d}')
    # scipy-free check of the gaussian: unit sum and sigma
    w, rr = ib.gauss_kernel(2.0); x = np.arange(-rr, rr + 1)
    check(abs(sum(w) - 1) < 1e-12 and abs(np.sqrt((np.array(w) * x * x).sum()) - 2.0) < 0.02, 'gaussian kernel')

# 3. parity with the browser core
tmp = tempfile.mkdtemp()
def node(path, seg, target='ink', tr=0.5, tp=0.5, inv=False):
    o = subprocess.run(['node', os.path.join(HERE, 'node_bench.mjs'), path, seg, target, str(tr), str(tp), '1' if inv else '0'],
                       capture_output=True, text=True, check=True)
    return json.loads(o.stdout)
base = synth.make('w047', -9.5, 6.25, 1.07, 1.07, 1.5)
files = []
p = os.path.join(tmp, 'p16.png'); ib.save_png16(base, p); files.append((p, 'w047', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'p8.png'); Image.fromarray(np.round(base * 255).astype(np.uint8)).save(p); files.append((p, 'w047', 'ink', 0.4, 0.3, False))
rgb = np.stack([base, base * 0.8, base * 0.5], -1)
p = os.path.join(tmp, 'rgb.png'); Image.fromarray(np.round(rgb * 255).astype(np.uint8)).save(p); files.append((p, 'w047', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'pal.png'); Image.fromarray(np.round(base * 255).astype(np.uint8)).convert('P').save(p); files.append((p, 'w047', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'label.png'); Image.fromarray((base > 0.45).astype(np.uint8)).save(p); files.append((p, 'w047', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'f32.tif'); Image.fromarray((base * 7 - 2).astype(np.float32)).save(p); files.append((p, 'w047', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'u16_lzw.tif'); Image.fromarray(np.round(base * 65535).astype(np.uint16)).save(p, compression='tiff_lzw'); files.append((p, 'w047', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'u8_deflate.tif'); Image.fromarray(np.round(base * 255).astype(np.uint8)).save(p, compression='tiff_adobe_deflate'); files.append((p, 'w047', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'inv.png'); ib.save_png16(1 - base, p); files.append((p, 'w047', 'ink', 0.5, 0.5, True))
files.append((os.path.join(ROOT, 'data/bench/w049_gpu.png'), 'w049', 'ink', 0.5, 0.5, False))
files.append((os.path.join(ROOT, 'data/bench/w049_gpu.png'), 'w049', 'ct', 0.5, 0.5, False))
big = synth.make('w033', 4, -6, 1.0, 1.0, 1, c=3.3)
p = os.path.join(tmp, 'big.png'); ib.save_png16(big, p); files.append((p, 'w033', 'ink', 0.5, 0.5, False))
p = os.path.join(tmp, 'noise.png'); Image.fromarray((synth.noise('w046') * 255).astype(np.uint8)).save(p); files.append((p, 'w046', 'ink', 0.5, 0.5, False))
for path, seg, target, tr, tp, inv in files:
    j = node(path, seg, target, tr, tp, inv)
    py = ib.compare(path, seg, target, None, tr, tp, inv, refs=refs)
    same = ib.comparable(j) == ib.comparable(py)
    diff = [k for k in ib.comparable(py) if ib.comparable(py)[k] != ib.comparable(j).get(k)]
    results['parity'].append({'file': os.path.basename(path), 'seg': seg, 'target': target, 'identical': same, 'pearson': py['metrics']['pearson'], 'diff': diff})
    print('parity', os.path.basename(path), seg, target, same, py['metrics']['pearson'], diff)
    check(same, f'parity {path} {diff}')

# 4. noise
for seed in (1, 2, 3):
    b = ib.compare(synth.noise('w046', seed), 'w046', refs=refs)
    m = b['metrics']
    results['noise'].append({'seed': seed, 'pearson': m['pearson'], 'iou': m['iou'], 'dice': m['dice'], 'ap': m['ap'], 'alignNcc': b['alignment']['ncc']})
    print('noise', results['noise'][-1])
    check(abs(m['pearson']) < 0.1 and m['ap'] < 0.2, 'noise scored high')

json.dump(results, open(os.path.join(tempfile.gettempdir(), 'inkbench_results.json'), 'w'), indent=1)
print('FAILURES' if fails else 'ALL PASSED', len(fails))
sys.exit(1 if fails else 0)