#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Chase Hendrick
"""inkbench: the Unroll Ink Bench as a single-file Python and CLI companion.

Compares an ink prediction for one PHerc. 139 winding with the Vesuvius Challenge
published ink layer, the same way the browser Bench does (lib/bench.js), and writes the
same bundle JSON (schema unroll-inkbench/1, see docs/BENCH_FORMAT.md).

All scores are AGREEMENT WITH ANOTHER MODEL'S OUTPUT (the published ink detection),
not agreement with ground truth.

Needs numpy and Pillow. Use as a module:
    import inkbench
    b = inkbench.compare('pred.png', segment='w046')
    print(b['metrics']['pearson'])
or as a CLI:
    inkbench compare pred.png --segment w046 -o bundle.json
    inkbench rank a.json b.json
    inkbench validate bundle.json
    inkbench adapt tif-stack stack.tif -o pred.png
"""
import argparse, datetime, hashlib, io, json, math, os, sys, urllib.request
import numpy as np

__version__ = '1.0.0'
SCHEMA = 'unroll-inkbench/1'
DEFAULT_BASE = 'https://huggingface.co/spaces/SharpMeow/unroll/resolve/main/'
DISCLAIMER = ("Scores measure agreement with another model's output (the Vesuvius Challenge "
              "published ink detection), not agreement with ground truth. Letter boxes are "
              "candidates for expert review, not readings.")
CREDITS = {
    'data': 'PHerc. 139 scan, segments and published ink: Vesuvius Challenge open data (scrollprize.org), '
            'ESRF; held by the Biblioteca Nazionale di Napoli.',
    'dataLicense': 'CC BY-NC 4.0 (https://creativecommons.org/licenses/by-nc/4.0/)',
    'code': 'Unroll Ink Bench by Chase Hendrick',
    'codeLicense': 'Apache-2.0',
}
LONG_FIT = 2          # predictions larger than 2x the reference grid are 2x2 mean pooled first
COARSE_MAX = 80       # long side of the coarse search level
SCALES = 25           # coarse scale candidates, geometric between 0.6 and 1.6
SHIFT_FRAC = 0.25     # coarse shift range, fraction of the side
MIN_VALID = 0.3       # NCC needs at least this fraction of the grid covered
TOL = 1e-9            # an alignment step must improve NCC by more than this
PR_STEPS = 50

# ------------------------------------------------------------------ numbers

def rnd(x, d):
    """Round like JavaScript Math.round(x * 10**d) / 10**d."""
    k = 10.0 ** d
    return math.floor(x * k + 0.5) / k

# ------------------------------------------------------------------ decoding

def to_gray(arr, maxval):
    a = arr.astype(np.float64)
    if a.ndim == 3:
        if a.shape[2] >= 3:
            a = (299.0 * a[:, :, 0] + 587.0 * a[:, :, 1] + 114.0 * a[:, :, 2]) / 1000.0
        else:
            a = a[:, :, 0]
    return a, maxval

def normalize(a, kind, maxval):
    """kind 'int': divide by maxval, unless the image only holds 0 and 1 (a label map).
    kind 'float': min-max over finite values. Non-finite values become 0."""
    if kind == 'float':
        f = np.isfinite(a)
        lo = float(a[f].min()) if f.any() else 0.0
        hi = float(a[f].max()) if f.any() else 1.0
        out = np.where(f, a, lo)
        if hi > lo:
            return (out - lo) / (hi - lo), 'minmax'
        return np.zeros_like(out), 'minmax'
    mx = float(a.max()) if a.size else 0.0
    if mx <= 1.0:
        return a / 1.0, 'label'
    return a / maxval, 'max%d' % int(maxval)

def decode_bytes(buf, name=''):
    """Decode PNG or TIFF bytes to a float64 [0,1] gray image. Returns (img, info)."""
    from PIL import Image
    im = Image.open(io.BytesIO(buf))
    fmt = (im.format or '').lower()
    if fmt not in ('png', 'tiff'):
        raise ValueError('Ink Bench reads PNG or TIFF (got %s)' % (im.format or 'unknown'))
    if fmt == 'png' and im.info.get('interlace'):
        raise ValueError('Interlaced PNG is not supported; save without interlacing.')
    if fmt == 'png':
        bd, ct = buf[24], buf[25]
        if bd not in (8, 16):
            raise ValueError('PNG with %d bits per sample is not supported; save as 8 or 16 bit.' % bd)
        if bd == 16 and ct != 0:
            raise ValueError('16 bit PNG must be grayscale; save color images as 8 bit.')
    if fmt == 'tiff':
        t = im.tag_v2
        bps = t.get(258, (1,)); bps = bps[0] if isinstance(bps, tuple) else bps
        spp = t.get(277, 1)
        if t.get(262, 1) in (0, 3):
            raise ValueError('Palette and WhiteIsZero TIFF are not supported.')
        if bps > 8 and spp > 1:
            raise ValueError('16 and 32 bit TIFF must have one channel.')
        if t.get(339, 1) == 2:
            raise ValueError('Signed integer TIFF is not supported; use unsigned or float.')
    mode = im.mode
    if mode == 'P':
        arr = np.asarray(im.convert('RGB')); kind, mv = 'int', 255.0
    elif mode in ('1',):
        arr = np.asarray(im).astype(np.uint8); kind, mv = 'int', 1.0
    elif mode in ('L', 'LA', 'RGB', 'RGBA'):
        arr = np.asarray(im); kind, mv = 'int', 255.0
    elif mode in ('I;16', 'I;16B', 'I;16L'):
        arr = np.asarray(im).astype(np.uint16); kind, mv = 'int', 65535.0
    elif mode == 'I':
        arr = np.asarray(im)
        kind, mv = ('int', 65535.0) if fmt == 'png' else ('int', 4294967295.0)
    elif mode == 'F':
        arr = np.asarray(im); kind, mv = 'float', 1.0
    else:
        raise ValueError('Unsupported image mode %s' % mode)
    g, mv = to_gray(arr, mv)
    if mode == 'LA':
        g = arr[:, :, 0].astype(np.float64)
    img, norm = normalize(g, kind, mv)
    return img, {'format': fmt, 'width': int(img.shape[1]), 'height': int(img.shape[0]), 'normalization': norm}

# ------------------------------------------------------------------ references

class Refs:
    """Reads data/bench/ from a local checkout (root) or from the Space (base URL)."""
    def __init__(self, root=None, base=DEFAULT_BASE):
        self.root, self.base, self.cache = root, base, {}
        self.index = json.loads(self.read('data/bench/index.json'))
    def read(self, rel):
        if rel in self.cache:
            return self.cache[rel]
        if self.root:
            b = open(os.path.join(self.root, rel), 'rb').read()
        else:
            with urllib.request.urlopen(self.base + rel, timeout=60) as r:
                b = r.read()
        self.cache[rel] = b
        return b
    def segment(self, seg):
        s = self.index['segments'].get(seg)
        if not s:
            raise KeyError('Unknown segment %s; known: %s' % (seg, ', '.join(self.index['segments'])))
        return s
    def layer(self, seg, layer):
        s = self.segment(seg)
        L = s['layers'][layer]
        b = self.read(L['file'])
        img, _ = decode_bytes(b)
        return img, L

def find_root():
    here = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')
    return here if os.path.exists(os.path.join(here, 'data/bench/index.json')) else None

# ------------------------------------------------------------------ image ops

def pool2(a):
    h, w = a.shape[0] // 2, a.shape[1] // 2
    a = a[:h * 2, :w * 2]
    return \
        (((a[0::2, 0::2] + a[0::2, 1::2]) + a[1::2, 0::2]) + a[1::2, 1::2]) * 0.25

def prepool(P, W, H):
    while P.shape[1] > LONG_FIT * W or P.shape[0] > LONG_FIT * H:
        P = pool2(P)
    return P

def axis_coords(n_out, n_in, t, s, xs=None):
    """Source coordinate for each output index along one axis."""
    x = np.arange(n_out, dtype=np.float64) if xs is None else xs
    f = n_in / n_out
    return ((((x + 0.5) - n_out / 2) - t) / s + n_out / 2) * f - 0.5

def warp(P, W, H, tx, ty, sx, sy, xs=None, ys=None, n_ref=None):
    """Bilinear sample P onto a W x H grid. Returns (values, valid)."""
    hp, wp = P.shape
    nW, nH = (W, H) if n_ref is None else n_ref
    u = axis_coords(nW, wp, tx, sx, xs)
    v = axis_coords(nH, hp, ty, sy, ys)
    vu = (u >= 0) & (u <= wp - 1)
    vv = (v >= 0) & (v <= hp - 1)
    x0 = np.clip(np.floor(u), 0, max(wp - 2, 0)).astype(np.int64)
    y0 = np.clip(np.floor(v), 0, max(hp - 2, 0)).astype(np.int64)
    fx = u - x0
    fy = v - y0
    x1 = np.minimum(x0 + 1, wp - 1)
    y1 = np.minimum(y0 + 1, hp - 1)
    p00 = P[y0[:, None], x0[None, :]]
    p10 = P[y0[:, None], x1[None, :]]
    p01 = P[y1[:, None], x0[None, :]]
    p11 = P[y1[:, None], x1[None, :]]
    FX = fx[None, :]
    top = p00 + (p10 - p00) * FX
    bot = p01 + (p11 - p01) * FX
    val = top + (bot - top) * fy[:, None]
    valid = vv[:, None] & vu[None, :]
    val = np.where(valid, val, 0.0)
    return val, valid

def ncc(a, b, m, nmin):
    n = float(m.sum())
    if n < nmin:
        return -1.0
    am = np.where(m, a, 0.0)
    Sa = float(am.sum()); Sb = float(b.sum())
    Saa = float((am * am).sum()); Sbb = float((b * b).sum()); Sab = float((am * b).sum())
    va = Saa - Sa * Sa / n
    vb = Sbb - Sb * Sb / n
    if va <= 1e-12 or vb <= 1e-12:
        return 0.0
    return (Sab - Sa * Sb / n) / math.sqrt(va * vb)

def gauss_kernel(sigma):
    r = int(math.ceil(3 * sigma))
    w = [math.exp(-(i * i) / (2 * sigma * sigma)) for i in range(-r, r + 1)]
    s = 0.0
    for x in w:
        s += x
    return [x / s for x in w], r

def blur(a, sigma):
    if sigma <= 0:
        return a
    w, r = gauss_kernel(sigma)
    H, W = a.shape
    ix = np.arange(W)
    out = np.zeros_like(a)
    for k, i in enumerate(range(-r, r + 1)):
        out = out + w[k] * a[:, np.clip(ix + i, 0, W - 1)]
    iy = np.arange(H)
    out2 = np.zeros_like(a)
    for k, i in enumerate(range(-r, r + 1)):
        out2 = out2 + w[k] * out[np.clip(iy + i, 0, H - 1), :]
    return out2

# ------------------------------------------------------------------ alignment

def pyramid(a, levels):
    out = [a]
    for _ in range(levels):
        out.append(pool2(out[-1]))
    return out

def align(ref, P, start=None):
    """Find tx, ty (reference px), sx, sy that best map P onto ref. Returns dict."""
    H, W = ref.shape
    P = prepool(P, W, H)
    K = 0
    while max(W >> K, H >> K) > COARSE_MAX:
        K += 1
    R = pyramid(ref, K)
    Q = pyramid(P, K)
    evals = 0
    if start is None:
        r = R[K]; q = Q[K]
        h, w = r.shape
        nmin = MIN_VALID * w * h
        Dx, Dy = int(math.floor(SHIFT_FRAC * w + 0.5)), int(math.floor(SHIFT_FRAC * h + 0.5))
        best = (-2.0, 1.0, 0, 0)
        for i in range(SCALES):
            s = math.exp(math.log(0.6) + (math.log(1.6) - math.log(0.6)) * i / (SCALES - 1))
            xs = np.arange(-Dx, w + Dx, dtype=np.float64)
            ys = np.arange(-Dy, h + Dy, dtype=np.float64)
            S, M = warp(q, w, h, 0.0, 0.0, s, s, xs, ys, (w, h))
            for dy in range(-Dy, Dy + 1):
                for dx in range(-Dx, Dx + 1):
                    b = S[Dy - dy:Dy - dy + h, Dx - dx:Dx - dx + w]
                    m = M[Dy - dy:Dy - dy + h, Dx - dx:Dx - dx + w]
                    c = ncc(r, b, m, nmin); evals += 1
                    if c > best[0] + TOL:
                        best = (c, s, dx, dy)
        p = [best[2] * (1 << K), best[3] * (1 << K), best[1], best[1]]
        coarse = best[0]
    else:
        p = list(start); coarse = None
    ds = 0.02
    for k in range(K, -1, -1):
        r = R[k]; q = Q[k]; h, w = r.shape; f = float(1 << k)
        nmin = MIN_VALID * w * h
        def F(pp):
            v, m = warp(q, w, h, pp[0] / f, pp[1] / f, pp[2], pp[3])
            return ncc(r, v, m, nmin)
        cur = F(p); evals += 1
        dt = f
        stop_t = 0.02 if k == 0 else f * 0.25
        n = 0
        while dt >= stop_t and n < 400:
            improved = False
            for i in range(4):
                st = dt if i < 2 else ds
                for sg in (1.0, -1.0):
                    t = list(p); t[i] = p[i] + sg * st
                    if i >= 2 and not (0.3 <= t[i] <= 3.0):
                        continue
                    c = F(t); evals += 1; n += 1
                    if c > cur + TOL:
                        p, cur, improved = t, c, True
                        break
            if not improved:
                dt *= 0.5; ds *= 0.5
        ds = max(ds * 2, 0.0005)
    out = {'tx': rnd(p[0], 3), 'ty': rnd(p[1], 3), 'sx': rnd(p[2], 5), 'sy': rnd(p[3], 5)}
    out['evals'] = evals
    out['coarseNcc'] = None if coarse is None else rnd(coarse, 6)
    return out

# ------------------------------------------------------------------ metrics

SIGMAS = [i * 0.5 for i in range(13)]

def metrics(ref, val, valid, t_ref=0.5, t_pred=0.5, px_per_mm=None):
    H, W = ref.shape
    n = int(valid.sum())
    out = {'nValid': n, 'validFraction': rnd(n / (W * H), 6)}
    pear = ncc(ref, val, valid, 1)
    out['pearson'] = rnd(pear, 6)
    best = (-2.0, 0.0)
    for s in SIGMAS:
        c = ncc(blur(ref, s), val, valid, 1)
        if c > best[0] + TOL:
            best = (c, s)
    out['blurSigmaPx'] = best[1]
    out['blurSigmaUM'] = None if not px_per_mm else rnd(best[1] / px_per_mm * 1000, 1)
    out['pearsonMatchedBlur'] = rnd(best[0], 6)
    A = (ref >= t_ref) & valid
    B = (val >= t_pred) & valid
    inter = int((A & B).sum()); uni = int((A | B).sum())
    a, b = int(A.sum()), int(B.sum())
    out['thresholds'] = {'ref': t_ref, 'pred': t_pred}
    out['iou'] = rnd(inter / uni, 6) if uni else 0.0
    out['dice'] = rnd(2 * inter / (a + b), 6) if a + b else 0.0
    out['precision'] = rnd(inter / b, 6) if b else 1.0
    out['recall'] = rnd(inter / a, 6) if a else 0.0
    out['refPositiveFraction'] = rnd(a / n, 6) if n else 0.0
    out['predPositiveFraction'] = rnd(b / n, 6) if n else 0.0
    pr = []
    ap, prev_r, bf = 0.0, 0.0, (-1.0, 0.0)
    vv = val[valid]; aa = A[valid]
    for j in range(PR_STEPS, -1, -1):
        t = j / PR_STEPS
        pb = vv >= t
        tp = int((pb & aa).sum()); fp = int((pb & ~aa).sum())
        p = tp / (tp + fp) if tp + fp else 1.0
        r = tp / a if a else 0.0
        ap += (r - prev_r) * p; prev_r = r
        f1 = 2 * p * r / (p + r) if p + r else 0.0
        if f1 > bf[0] + TOL:
            bf = (f1, t)
        pr.append({'t': rnd(t, 6), 'precision': rnd(p, 6), 'recall': rnd(r, 6)})
    out['ap'] = rnd(ap, 6)
    out['bestF1'] = {'f1': rnd(bf[0], 6), 'threshold': rnd(bf[1], 6)}
    out['pr'] = pr
    return out

# ------------------------------------------------------------------ bundle

def sha256(b):
    return hashlib.sha256(b).hexdigest()

def compare(pred, segment, target='ink', model=None, t_ref=0.5, t_pred=0.5, invert=False,
            alignment=None, annotations=None, refs=None, name=None):
    """pred: path, bytes, or 2-D numpy array in [0,1]. Returns the bundle dict."""
    refs = refs or Refs(root=find_root())
    if isinstance(pred, np.ndarray):
        P = np.asarray(pred, dtype=np.float64)
        info = {'format': 'array', 'width': int(P.shape[1]), 'height': int(P.shape[0]), 'normalization': 'none'}
        digest = sha256(np.ascontiguousarray(P).tobytes())
    else:
        buf = pred if isinstance(pred, (bytes, bytearray)) else open(pred, 'rb').read()
        if name is None and not isinstance(pred, (bytes, bytearray)):
            name = os.path.basename(pred)
        P, info = decode_bytes(bytes(buf))
        digest = sha256(buf)
    if invert:
        P = 1.0 - P
    seg = refs.segment(segment)
    ref, Lref = refs.layer(segment, 'ink')
    tgt, Ltgt = (ref, Lref) if target == 'ink' else refs.layer(segment, target)
    H, W = ref.shape
    if alignment:
        al = {k: alignment[k] for k in ('tx', 'ty', 'sx', 'sy')}
        method = alignment.get('method', 'manual')
        al.update({'evals': 0, 'coarseNcc': None})
    else:
        al = align(tgt, P); method = 'auto'
    PP = prepool(P, W, H)
    val, valid = warp(PP, W, H, al['tx'], al['ty'], al['sx'], al['sy'])
    m = metrics(ref, val, valid, t_ref, t_pred, seg['pxPerMM'])
    alignment_out = {'method': method, 'target': target, 'tx': al['tx'], 'ty': al['ty'], 'sx': al['sx'], 'sy': al['sy'],
                     'ncc': rnd(ncc(tgt, val, valid, 1), 6), 'coarseNcc': al['coarseNcc'], 'evals': al['evals'],
                     'model': 'bilinear, scale about grid center, then shift; units are reference px'}
    return {
        'schema': SCHEMA,
        'tool': {'name': 'inkbench', 'version': __version__, 'impl': 'python'},
        'created': datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ'),
        'scroll': refs.index['scroll'],
        'segment': segment,
        'reference': {'layer': 'ink', 'label': 'Published ink (2.4 um), Vesuvius Challenge', 'file': Lref['file'],
                      'sha256': Lref['sha256'], 'width': W, 'height': H, 'pxPerMM': seg['pxPerMM'],
                      'alignTargetFile': Ltgt['file'], 'alignTargetSha256': Ltgt['sha256']},
        'prediction': {'name': name or 'prediction', 'sha256': digest, 'format': info['format'],
                       'width': info['width'], 'height': info['height'], 'normalization': info['normalization'],
                       'invert': bool(invert), 'model': model or {'name': 'unspecified'}},
        'alignment': alignment_out,
        'metrics': m,
        'annotations': annotations or [],
        'credits': CREDITS,
        'disclaimer': DISCLAIMER,
    }

def comparable(b):
    """The part of a bundle that must be identical between implementations."""
    b = json.loads(json.dumps(b))
    for k in ('tool', 'created'):
        b.pop(k, None)
    return b

# ------------------------------------------------------------------ validation and ranking

REQUIRED = {
    '': ['schema', 'scroll', 'segment', 'reference', 'prediction', 'alignment', 'metrics', 'annotations', 'disclaimer'],
    'reference': ['layer', 'sha256', 'width', 'height'],
    'prediction': ['sha256', 'model'],
    'alignment': ['method', 'tx', 'ty', 'sx', 'sy'],
    'metrics': ['pearson', 'iou', 'dice', 'thresholds', 'pr'],
}

def validate(b):
    errs = []
    if not isinstance(b, dict):
        return ['bundle is not a JSON object']
    if not str(b.get('schema', '')).startswith('unroll-inkbench/1'):
        errs.append('schema must be unroll-inkbench/1')
    for sect, keys in REQUIRED.items():
        o = b if not sect else b.get(sect)
        if not isinstance(o, dict):
            errs.append('missing object %s' % sect); continue
        for k in keys:
            if k not in o:
                errs.append('missing %s%s' % (sect + '.' if sect else '', k))
    for i, a in enumerate(b.get('annotations') or []):
        for k in ('x', 'y', 'w', 'h', 'letter', 'confidence', 'status'):
            if k not in a:
                errs.append('annotations[%d] missing %s' % (i, k))
        if a.get('status') != 'candidate for expert review':
            errs.append('annotations[%d].status must be "candidate for expert review"' % i)
    try:
        import jsonschema
        here = find_root()
        if here:
            sch = json.load(open(os.path.join(here, 'docs/bench-bundle.schema.json')))
            for e in jsonschema.Draft202012Validator(sch).iter_errors(b):
                errs.append('schema: %s' % e.message)
    except ImportError:
        pass
    return errs

def rank(bundles, metric='pearson'):
    rows = []
    for b in bundles:
        m = b['metrics']
        v = m['bestF1']['f1'] if metric == 'f1' else m[metric]
        rows.append({'segment': b['segment'], 'model': b['prediction']['model'].get('name', '?'),
                     'value': v, 'pearson': m['pearson'], 'iou': m['iou'], 'dice': m['dice'], 'ap': m.get('ap')})
    out = {}
    for r in rows:
        out.setdefault(r['segment'], []).append(r)
    for s in out:
        out[s].sort(key=lambda r: -r['value'])
        for i, r in enumerate(out[s]):
            r['rank'] = i + 1
    return out

# ------------------------------------------------------------------ adapters

def save_png16(a, path):
    """Save a float [0,1] array as a 16-bit gray PNG (what the browser Bench reads exactly)."""
    from PIL import Image
    a = np.clip(np.nan_to_num(np.asarray(a, dtype=np.float64)), 0, 1)
    Image.fromarray(np.round(a * 65535).astype(np.uint16)).save(path)

def _minmax(a):
    a = a.astype(np.float64); lo, hi = np.nanmin(a), np.nanmax(a)
    return (a - lo) / (hi - lo) if hi > lo else np.zeros_like(a)

def adapt_tif_stack(path, reduce='mean', page=None):
    """Multi-page TIFF (a layer stack or per-layer ink probabilities) to one 2-D map."""
    from PIL import Image, ImageSequence
    im = Image.open(path)
    pages = [np.asarray(p, dtype=np.float64) for p in ImageSequence.Iterator(im)]
    if page is not None:
        return _minmax(pages[page])
    st = np.stack(pages)
    return _minmax(st.max(0) if reduce == 'max' else st.mean(0))

def adapt_villa_folder(folder):
    """A vesuvius or villa style segment folder: uses the ink prediction PNG or TIFF found in
    the folder (names containing 'prediction' or 'ink', not 'inklabels'), masked by mask.png."""
    from PIL import Image
    cands = []
    for root, _, files in os.walk(folder):
        for f in sorted(files):
            fl = f.lower()
            if fl.endswith(('.png', '.tif', '.tiff')) and ('pred' in fl or 'ink' in fl) and 'label' not in fl:
                cands.append(os.path.join(root, f))
    if not cands:
        raise FileNotFoundError('no *prediction* or *ink* image in %s' % folder)
    img, _ = decode_bytes(open(cands[0], 'rb').read())
    mpath = next((os.path.join(folder, m) for m in ('mask.png', os.path.basename(folder.rstrip('/')) + '_mask.png')
                  if os.path.exists(os.path.join(folder, m))), None)
    if mpath:
        m = np.asarray(Image.open(mpath).convert('L').resize((img.shape[1], img.shape[0]), Image.NEAREST)) > 0
        img = np.where(m, img, 0.0)
    return img, cands[0], mpath

def adapt_nnunet(path, channel=1):
    """nnU-Net output: a label PNG (0/1, read as is) or a .npz with 'probabilities' (C, [Z,] H, W)."""
    if path.endswith('.npz'):
        p = np.load(path)['probabilities'][channel]
        while p.ndim > 2:
            p = p.max(0)
        return np.clip(p.astype(np.float64), 0, 1)
    img, _ = decode_bytes(open(path, 'rb').read())
    return img

def adapt_timesformer(path):
    """TimeSformer style ink inference output: a gray PNG of probabilities (0 to 255)."""
    img, _ = decode_bytes(open(path, 'rb').read())
    return img

# ------------------------------------------------------------------ CLI

def bibtex():
    return ('@software{hendrick_unroll_inkbench_2026,\n'
            '  author = {Hendrick, Chase},\n'
            '  title = {Unroll Ink Bench: model-agnostic agreement metrics for Herculaneum ink detection},\n'
            '  year = {2026},\n  version = {%s},\n'
            '  url = {https://huggingface.co/spaces/SharpMeow/unroll},\n'
            '  license = {Apache-2.0},\n'
            '  note = {Data: Vesuvius Challenge open data, CC BY-NC 4.0}\n}\n' % __version__)

def main(argv=None):
    ap = argparse.ArgumentParser(prog='inkbench', description=__doc__.split('\n')[0])
    sub = ap.add_subparsers(dest='cmd', required=True)
    c = sub.add_parser('compare', help='score a prediction against the published ink')
    c.add_argument('pred'); c.add_argument('--segment', required=True)
    c.add_argument('--target', default='ink', choices=['ink', 'ct'])
    c.add_argument('--model', default=None, help='model name for the bundle')
    c.add_argument('--model-url', default=None)
    c.add_argument('--t-ref', type=float, default=0.5); c.add_argument('--t-pred', type=float, default=0.5)
    c.add_argument('--invert', action='store_true')
    c.add_argument('--alignment', default=None, help='JSON file or bundle whose alignment to reuse')
    c.add_argument('--annotations', default=None, help='JSON file or bundle whose annotations to copy')
    c.add_argument('--root', default=None, help='local checkout of the Space (default: this repo, else the Space URL)')
    c.add_argument('--base', default=DEFAULT_BASE)
    c.add_argument('-o', '--out', default=None)
    c.add_argument('--min-pearson', type=float, default=None, help='exit 1 if pearson is below this (for CI)')
    r = sub.add_parser('rank', help='rank bundles per segment'); r.add_argument('bundles', nargs='+')
    r.add_argument('--metric', default='pearson', choices=['pearson', 'iou', 'dice', 'ap', 'f1'])
    v = sub.add_parser('validate', help='check a bundle against the format'); v.add_argument('bundles', nargs='+')
    sub.add_parser('cite', help='print a BibTeX entry')
    a = sub.add_parser('adapt', help='convert pipeline outputs to a Bench PNG')
    a.add_argument('kind', choices=['tif-stack', 'villa-folder', 'nnunet', 'timesformer'])
    a.add_argument('path'); a.add_argument('-o', '--out', required=True)
    a.add_argument('--reduce', default='mean', choices=['mean', 'max']); a.add_argument('--page', type=int, default=None)
    a.add_argument('--channel', type=int, default=1)
    sub.add_parser('segments', help='list segments with reference layers')
    args = ap.parse_args(argv)
    if args.cmd == 'compare':
        root = args.root or find_root()
        refs = Refs(root=root, base=args.base)
        def load(p):
            return json.load(open(p)) if p else None
        al = load(args.alignment); an = load(args.annotations)
        if al and 'alignment' in al: al = dict(al['alignment'], method='reused')
        if isinstance(an, dict): an = an.get('annotations', [])
        model = {'name': args.model or 'unspecified'}
        if args.model_url: model['url'] = args.model_url
        b = compare(args.pred, args.segment, args.target, model, args.t_ref, args.t_pred, args.invert, al, an, refs)
        b['tool']['impl'] = 'python-cli'
        s = json.dumps(b, indent=1, ensure_ascii=False)
        if args.out:
            open(args.out, 'w').write(s + '\n')
        else:
            print(s)
        m = b['metrics']
        print('pearson %.4f  iou %.4f  dice %.4f  ap %.4f  (agreement with published ink, not ground truth)'
              % (m['pearson'], m['iou'], m['dice'], m['ap']), file=sys.stderr)
        if args.min_pearson is not None and m['pearson'] < args.min_pearson:
            return 1
    elif args.cmd == 'rank':
        bs = [json.load(open(p)) for p in args.bundles]
        for seg, rows in sorted(rank(bs, args.metric).items()):
            print(seg)
            for r_ in rows:
                print('  %d. %-30s %s=%.4f  pearson=%.4f iou=%.4f dice=%.4f' % (r_['rank'], r_['model'][:30], args.metric, r_['value'], r_['pearson'], r_['iou'], r_['dice']))
    elif args.cmd == 'validate':
        bad = 0
        for p in args.bundles:
            e = validate(json.load(open(p)))
            print(p, 'OK' if not e else 'INVALID'); [print('  ' + x) for x in e]; bad += bool(e)
        return 1 if bad else 0
    elif args.cmd == 'cite':
        print(bibtex(), end='')
    elif args.cmd == 'segments':
        refs = Refs(root=find_root())
        for k, s in refs.index['segments'].items():
            print(k, '%dx%d' % (s['width'], s['height']), ' '.join(s['layers']))
    elif args.cmd == 'adapt':
        if args.kind == 'tif-stack':
            img = adapt_tif_stack(args.path, args.reduce, args.page)
        elif args.kind == 'villa-folder':
            img, src, mask = adapt_villa_folder(args.path)
            print('prediction', src, 'mask', mask, file=sys.stderr)
        elif args.kind == 'nnunet':
            img = adapt_nnunet(args.path, args.channel)
        else:
            img = adapt_timesformer(args.path)
        save_png16(img, args.out)
        print('wrote', args.out, img.shape[1], 'x', img.shape[0], file=sys.stderr)
    return 0

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