# SPDX-License-Identifier: Apache-2.0
"""Example Ink Bench bundles for the leaderboard: the public ink_9um runs on every winding and two
baselines any ink model should beat, X-ray intensity and uniform noise (seeded), on the same grid."""
import io, json, os, sys
import numpy as np
from PIL import Image
sys.path.insert(0, os.path.dirname(__file__))
import inkbench as ib
root = os.path.join(os.path.dirname(__file__), '..'); out = os.path.join(root, 'docs/examples/bench')
refs = ib.Refs(root=root)
ident = {'tx': 0.0, 'ty': 0.0, 'sx': 1.0, 'sy': 1.0, 'method': 'manual'}
for w in [k for k, v in refs.index['segments'].items() if 'gpu' in v['layers']]:
    s = refs.segment(w)
    b = ib.compare(os.path.join(root, s['layers']['gpu']['file']), w, refs=refs,
                   model={'name': 'scrollprize/ink_9um (seed 42, step 75000) on the 9.362 um scan', 'url': 'https://huggingface.co/scrollprize/ink_9um'})
    json.dump(b, open(os.path.join(out, w + '-ink9um.json'), 'w'), indent=1)
    b = ib.compare(os.path.join(root, s['layers']['ct']['file']), w, refs=refs, alignment=ident,
                   model={'name': 'baseline: X-ray intensity (same grid, no alignment needed)'})
    json.dump(b, open(os.path.join(out, w + '-ct-baseline.json'), 'w'), indent=1)
    rng = np.random.default_rng(int(w[1:]))
    buf = io.BytesIO(); Image.fromarray((rng.random((s['height'], s['width'])) * 255).astype(np.uint8)).save(buf, 'PNG')
    b = ib.compare(buf.getvalue(), w, refs=refs, alignment=ident, name='noise-seed%d.png' % int(w[1:]),
                   model={'name': 'baseline: uniform noise', 'notes': 'numpy default_rng(%d).random, 8 bit' % int(w[1:])})
    json.dump(b, open(os.path.join(out, w + '-noise-baseline.json'), 'w'), indent=1)
    print(w, 'done')