# Skull shell + hair-fit numbers for any 65-joint Quaternius GLB/glTF.
#
# The game rebinds Quaternius hair (100% rigid to the single `Head` bone) onto whatever body
# loaded, correcting it with CharacterBuildData.HeadFit = a diagonal scale about the Head rest
# pivot plus a bind translation. Those constants are only valid against the skull they were
# tuned on, and they go stale silently every time a body is rebaked. This measures the skull
# they SHOULD be tuned against, and predicts where a given hair style's crown actually lands.
#
#   python tools/measure_head_shell.py <body.glb> [<body.glb> ...]
#   python tools/measure_head_shell.py --hair <hair.gltf> --on <body.glb> --scale X,Y,Z --offset X,Y,Z
#
# Two shell measures are printed. The FULL Head-weighted shell is what Ozan's MakoHeadFit
# comment quotes; it is inflated by jaw/ear/neck weight bleed and is NOT the surface a hair cap
# sits on. The CRANIUM band (Head-weighted verts within 8 cm of the skull top) is the honest
# one — use it for any scale ratio.
#
# See docs/characters/head-features-contract.md.
import json, struct, sys, os, array

CT = {5120: ('b', 1), 5121: ('B', 1), 5122: ('h', 2), 5123: ('H', 2), 5125: ('I', 4), 5126: ('f', 4)}
NW = {'SCALAR': 1, 'VEC2': 2, 'VEC3': 3, 'VEC4': 4, 'MAT4': 16}


def load(path):
    if path.endswith('.glb'):
        d = open(path, 'rb').read()
        n = struct.unpack_from('<I', d, 12)[0]
        j = json.loads(d[20:20 + n])
        off = 20 + n
        blen, _ = struct.unpack_from('<I4s', d, off)
        return j, [d[off + 8:off + 8 + blen]]
    j = json.load(open(path))
    base = os.path.dirname(path)
    return j, [open(os.path.join(base, b['uri']), 'rb').read() for b in j['buffers']]


def acc(j, bufs, i):
    a = j['accessors'][i]
    bv = j['bufferViews'][a['bufferView']]
    off = bv.get('byteOffset', 0) + a.get('byteOffset', 0)
    n = NW[a['type']]
    fmt, sz = CT[a['componentType']]
    stride = bv.get('byteStride')
    buf = bufs[bv['buffer']]
    if stride and stride != n * sz:
        out = []
        for k in range(a['count']):
            ar = array.array(fmt)
            ar.frombytes(buf[off + k * stride:off + k * stride + n * sz])
            out.append(tuple(ar))
        return out
    ar = array.array(fmt)
    ar.frombytes(buf[off:off + a['count'] * n * sz])
    return [tuple(ar[k * n:(k + 1) * n]) for k in range(a['count'])]


def mat_mul(a, b):
    """Column-major 4x4 (glTF order) product a*b."""
    out = [0.0] * 16
    for c in range(4):
        for r in range(4):
            out[c * 4 + r] = sum(a[k * 4 + r] * b[c * 4 + k] for k in range(4))
    return out


def node_local(nd):
    if 'matrix' in nd:
        return list(nd['matrix'])
    t = nd.get('translation', [0, 0, 0])
    q = nd.get('rotation', [0, 0, 0, 1])
    s = nd.get('scale', [1, 1, 1])
    x, y, z, w = q
    r = [
        1 - 2 * (y * y + z * z), 2 * (x * y + z * w), 2 * (x * z - y * w),
        2 * (x * y - z * w), 1 - 2 * (x * x + z * z), 2 * (y * z + x * w),
        2 * (x * z + y * w), 2 * (y * z - x * w), 1 - 2 * (x * x + y * y),
    ]
    return [
        r[0] * s[0], r[1] * s[0], r[2] * s[0], 0.0,
        r[3] * s[1], r[4] * s[1], r[5] * s[1], 0.0,
        r[6] * s[2], r[7] * s[2], r[8] * s[2], 0.0,
        t[0], t[1], t[2], 1.0,
    ]


def node_world(j):
    """Rest-pose world matrix per node (full TRS chain)."""
    parent = {}
    for i, nd in enumerate(j['nodes']):
        for c in nd.get('children', []):
            parent[c] = i
    cache = {}

    def m(i):
        if i in cache:
            return cache[i]
        local = node_local(j['nodes'][i])
        cache[i] = mat_mul(m(parent[i]), local) if i in parent else local
        return cache[i]

    def w(i):
        mm = m(i)
        return (mm[12], mm[13], mm[14])
    return w


CRANIUM_BAND_M = 0.08   # verts within this of the skull top = the surface a hair cap sits on


def head_points(path, weight_min=0.5):
    """(Head pivot, [Head-weighted verts], total vert count) for a skinned body or hair file."""
    j, bufs = load(path)
    skin = j['skins'][0]
    names = [j['nodes'][n].get('name') for n in skin['joints']]
    if 'Head' not in names:
        raise SystemExit(f"{os.path.basename(path)}: no Head joint")
    hidx = names.index('Head')
    pivot = node_world(j)(skin['joints'][hidx])
    pts, total = [], 0
    for m in j['meshes']:
        for pr in m['primitives']:
            at = pr['attributes']
            if 'JOINTS_0' not in at:
                continue
            P = acc(j, bufs, at['POSITION'])
            J = acc(j, bufs, at['JOINTS_0'])
            W = acc(j, bufs, at['WEIGHTS_0'])
            total += len(P)
            for p, jj, ww in zip(P, J, W):
                if sum(ww[k] for k in range(4) if jj[k] == hidx) >= weight_min:
                    pts.append(p)
    if not pts:
        raise SystemExit(f"{os.path.basename(path)}: no Head-weighted verts")
    return pivot, pts, total


def _pct(v, p):
    v = sorted(v)
    k = (len(v) - 1) * p
    f = int(k)
    return v[f] + (v[min(f + 1, len(v) - 1)] - v[f]) * (k - f)


def landmarks(path, slab=0.01):
    """Two robust cranium landmarks in Head-pivot-relative metres, plus spans at the wider one.

    `top`  = skull crown.
    `wide` = the height where the cranium is widest (temple/ear line). Two landmarks is exactly
             what a scale-about-pivot + translate has degrees of freedom for, per axis.
    """
    pivot, pts, _ = head_points(path)
    ys = [p[1] for p in pts]
    top = _pct(ys, 0.999)
    best = None
    y = top
    while y > top - 0.16:                      # search the top 16 cm — cranium, never the jaw
        slabpts = [p for p in pts if y - slab <= p[1] < y]
        if len(slabpts) >= 8:
            xs = [p[0] for p in slabpts]
            span = _pct(xs, 0.98) - _pct(xs, 0.02)
            if best is None or span > best[1]:
                zs = [p[2] for p in slabpts]
                best = (y - slab / 2, span, _pct(zs, 0.98) - _pct(zs, 0.02),
                        (_pct(zs, 0.98) + _pct(zs, 0.02)) / 2)
        y -= slab
    if best is None:
        raise SystemExit(f"{os.path.basename(path)}: cranium slabs too sparse")
    wide_y, wide_w, wide_d, wide_zc = best
    return {
        'name': os.path.basename(path), 'pivot_y': pivot[1],
        'top': top - pivot[1], 'wide': wide_y - pivot[1],
        'wide_w': wide_w, 'wide_d': wide_d, 'wide_zc': wide_zc,
    }


def fit(stock_path, body_path):
    """Derive HeadScale / HairOffset that carry the stock cranium onto this body's cranium.

    Runtime maps an authored vert at Head-local h to  shipPivot + h*scale + offset, so two
    matched landmarks per axis determine both constants exactly — no bed guessing.
    """
    s, t = landmarks(stock_path), landmarks(body_path)
    sy = (t['top'] - t['wide']) / (s['top'] - s['wide'])
    oy = t['top'] - s['top'] * sy
    sx = t['wide_w'] / s['wide_w']
    sz = t['wide_d'] / s['wide_d']
    oz = t['wide_zc'] - s['wide_zc'] * sz
    for d in (s, t):
        print(f"   {d['name']:42s} top {d['top']*100:6.2f}  wide-line {d['wide']*100:6.2f}"
              f"  w {d['wide_w']*100:6.2f}  d {d['wide_d']*100:6.2f}   (cm above pivot)")
    print(f"   HeadScale  = new Vector3({sx:.3f}f, {sy:.3f}f, {sz:.3f}f)")
    print(f"   HairOffset = new Vector3(0f, {oy:.4f}f, {oz:.4f}f)")
    return (sx, sy, sz), (0.0, oy, oz)


def measure(path):
    pivot, pts, total = head_points(path)
    ys = [p[1] for p in pts]
    top, chin = max(ys), min(ys)
    band = [p for p in pts if p[1] >= top - CRANIUM_BAND_M]
    def span(v):
        return (max(v) - min(v)) * 100
    print(f"{os.path.basename(path):46s} headverts {len(pts):6d}/{total}")
    print(f"   Head pivot Y {pivot[1]:.4f}   skull top Y {top:.4f}   chin Y {chin:.4f}")
    print(f"   full shell     w {span([p[0] for p in pts]):6.2f}  d {span([p[2] for p in pts]):6.2f}   (cm)")
    print(f"   cranium band   w {span([p[0] for p in band]):6.2f}  d {span([p[2] for p in band]):6.2f}"
          f"   n={len(band)}   <- use THIS for scale ratios")
    print(f"   skull top above pivot {(top - pivot[1]) * 100:6.2f} cm"
          f"   pivot above chin {(pivot[1] - chin) * 100:6.2f} cm")


def predict(hair_path, body_path, stock_path, scale, offset, terse=False):
    """Does this style land on this body the way it was authored to land on the stock head?

    Runtime (HeadAccessoryBindMath.RestToWorldStepwise) on a vert rigid to Head:
        world = bodyHeadRest + ScaleAbout(authoredLocal, 0, headScale) + offset

    The crown's ABSOLUTE height over the skull is not the test — a mohawk stands proud and a
    beard hangs at the chin by design. The test is whether the style keeps the clearance it
    was drawn with on its own authoring rig.
    """
    hpivot, hpts, _ = head_points(hair_path)
    bpivot, bpts, _ = head_points(body_path)
    spivot, spts, _ = head_points(stock_path)
    crown_local = max(p[1] for p in hpts) - hpivot[1]
    authored = crown_local - (_pct([p[1] for p in spts], 0.999) - spivot[1])
    skull_top = _pct([p[1] for p in bpts], 0.999)
    landed = bpivot[1] + crown_local * scale[1] + offset[1]
    got = landed - skull_top
    err = (got - authored) * 100
    verdict = "ok" if abs(err) <= 1.0 else ("HIGH" if err > 0 else "LOW")
    if terse:
        print(f"   {os.path.basename(hair_path):26s} authored {authored*100:+6.2f}  "
              f"got {got*100:+6.2f}  err {err:+6.2f} cm  {verdict}")
        return err
    print(f"{os.path.basename(hair_path)} on {os.path.basename(body_path)}")
    print(f"   authored crown {crown_local*100:6.2f} cm above its pivot"
          f"  = {authored*100:+.2f} cm clearance over the stock skull")
    print(f"   scale {scale}  offset {offset}")
    print(f"   lands Y {landed:.4f}  skull top Y {skull_top:.4f}"
          f"  = {got*100:+.2f} cm clearance")
    print(f"   error vs authored  {err:+6.2f} cm   {verdict}")
    return err


def sweep(stock_path, body_path, hair_dir, scale, offset):
    print(f"   {'style':26s} {'authored':>8s} {'got':>8s} {'err':>10s}")
    errs = []
    for f in sorted(os.listdir(hair_dir)):
        if not f.startswith('Hair_') or not f.endswith('.gltf') or 'Teen' in f:
            continue
        errs.append(abs(predict(os.path.join(hair_dir, f), body_path, stock_path,
                                scale, offset, terse=True)))
    print(f"   -> worst {max(errs):.2f} cm, mean {sum(errs)/len(errs):.2f} cm over {len(errs)} styles")


def _vec(s):
    return tuple(float(v) for v in s.split(','))


if __name__ == '__main__':
    a = sys.argv[1:]
    sc = _vec(a[a.index('--scale') + 1]) if '--scale' in a else (1, 1, 1)
    off = _vec(a[a.index('--offset') + 1]) if '--offset' in a else (0, 0, 0)
    if '--fit' in a:
        fit(a[a.index('--fit') + 1], a[a.index('--on') + 1])
    elif '--sweep' in a:
        sweep(a[a.index('--stock') + 1], a[a.index('--on') + 1],
              a[a.index('--sweep') + 1], sc, off)
    elif '--hair' in a:
        predict(a[a.index('--hair') + 1], a[a.index('--on') + 1],
                a[a.index('--stock') + 1], sc, off)
    else:
        for p in a:
            measure(p)