"""UV stretch audit of a kit.glb (contract v2: world-scale UVs, u,v = metres / tile size). For every triangle, compare
UV area with 3D area / tile^2; a triangle whose ratio is under 1/8 or over 8 (or whose UV collapses to a line) is stretched
and shows as streaks in the game. usage: kit_uv_audit.py <kit dir> [--json]. Prints atoms ranked by stretched 3D area."""
import json, struct, sys
import numpy as np

def load(path):
    b = open(path, 'rb').read(); n = struct.unpack_from('<I', b, 12)[0]
    js = json.loads(b[20:20 + n]); off = 20 + n; ln = struct.unpack_from('<I', b, off)[0]; binb = b[off + 8: off + 8 + ln]
    return js, binb

CT = {5126: ('f4', 4), 5123: ('u2', 2), 5125: ('u4', 4), 5121: ('u1', 1)}
NC = {'SCALAR': 1, 'VEC2': 2, 'VEC3': 3, 'VEC4': 4}
def acc(js, binb, i):
    a = js['accessors'][i]; bv = js['bufferViews'][a['bufferView']]; dt, sz = CT[a['componentType']]; nc = NC[a['type']]
    o = bv.get('byteOffset', 0) + a.get('byteOffset', 0); stride = bv.get('byteStride', sz * nc)
    if stride == sz * nc: return np.frombuffer(binb, dtype=dt, count=a['count'] * nc, offset=o).reshape(a['count'], nc).astype('f8')
    return np.array([np.frombuffer(binb, dtype=dt, count=nc, offset=o + k * stride) for k in range(a['count'])], dtype='f8')

def main(kdir, as_json=False):
    js, binb = load(f'{kdir}/kit.glb')
    shelf = json.load(open(f'{kdir}/materials.json')); mats = shelf.get('materials', shelf)
    tile = {}
    for m in (mats if isinstance(mats, list) else [dict(name=k, **v) for k, v in mats.items()]):
        t = m.get('tile_size_metres') or m.get('tile') or 1; tile[m.get('name')] = (t[0] if isinstance(t, list) else t)
    out = []
    for node in js['nodes']:
        if 'mesh' not in node: continue
        bad = tot = 0.0; worst = []
        for prim in js['meshes'][node['mesh']]['primitives']:
            at = prim['attributes']
            if 'TEXCOORD_0' not in at or 'indices' not in prim: continue
            P, U, I = acc(js, binb, at['POSITION']), acc(js, binb, at['TEXCOORD_0']), acc(js, binb, prim['indices']).astype(int).reshape(-1, 3)
            mname = js['materials'][prim['material']]['name'] if 'material' in prim else '?'
            if any(k in mname.lower() for k in ('glass', 'iron', 'metal', 'lead')): continue
            t = tile.get(mname, 1.0)
            a3 = 0.5 * np.linalg.norm(np.cross(P[I[:, 1]] - P[I[:, 0]], P[I[:, 2]] - P[I[:, 0]]), axis=1)
            e1, e2 = U[I[:, 1]] - U[I[:, 0]], U[I[:, 2]] - U[I[:, 0]]
            a2 = 0.5 * np.abs(e1[:, 0] * e2[:, 1] - e1[:, 1] * e2[:, 0])
            big = a3 > 1e-4; ratio = np.where(big, a2 * t * t / np.maximum(a3, 1e-12), 1)
            s = big & ((ratio < 1 / 8) | (ratio > 8))
            bad += a3[s].sum(); tot += a3[big].sum()
            if s.any(): worst.append(mname)
        if tot: out.append({'atom': node.get('name'), 'stretched_m2': round(bad, 3), 'share': round(bad / tot, 3), 'materials': sorted(set(worst))})
    out.sort(key=lambda x: -x['stretched_m2'])
    if as_json: print(json.dumps(out, indent=1)); return
    for o in out:
        if o['share'] > 0.01: print(f"{o['atom']:28s} {o['stretched_m2']:8.2f} m2  {o['share']*100:5.1f} %  {', '.join(o['materials'])}")
    print(f"{sum(1 for o in out if o['share'] > 0.01)} of {len(out)} atoms have stretched UVs on more than 1 % of their area")

if __name__ == '__main__': main(sys.argv[1], '--json' in sys.argv)