"""Port the app's Z-Anatomy mesh without changing its geometry. CC BY-SA 4.0.
Usage: python3 tools/export-anatomy.py ../WorkoutNotes/AnatomyMesh.json dist/model
"""
import array, base64, gzip, hashlib, json, sys
from pathlib import Path

# Same ordered muscle rules as AnatomicalScene.swift (specific deltoids first).
RULES = [
 ('Chest',['pectoralis major']), ('Lats',['latissimus','teres major']),
 ('Upper back',['trapezius','rhomboid']), ('Biceps',['biceps brachii','brachialis']),
 ('Triceps',['triceps']), ('Front shoulders',['anterior part of deltoid','clavicular part of deltoid']),
 ('Rear shoulders',['posterior part of deltoid','spinal part of deltoid']), ('Side shoulders',['deltoid']),
 ('Core',['rectus abdominis','oblique','transversus abdominis']), ('Glutes',['gluteus']),
 ('Quads',['vastus','rectus femoris']), ('Hamstrings',['biceps femoris','semitendinosus','semimembranosus']),
 ('Calves',['gastrocnemius','soleus'])]

def export(source, output):
    assert sys.byteorder == 'little'
    parts = json.loads(Path(source).read_text())
    merged = {}
    for part in parts:
        group = next((g for g, words in RULES if any(w in part['name'].lower() for w in words)), '')
        key = (group, 'bone' if part['kind'] == 'bone' else 'muscle')
        mesh = merged.setdefault(key, {'p':array.array('f'), 'n':array.array('f'), 'i':array.array('I'), 'parts':0})
        p = array.array('f',base64.b64decode(part['p']))
        n = array.array('f',base64.b64decode(part['n']))
        i = array.array('I',base64.b64decode(part['i']))
        assert len(p)==len(n) and len(p)%3==0 and len(i)%3==0
        assert not i or max(i)<len(p)//3
        start=len(mesh['p'])//3
        mesh['p'].extend(p); mesh['n'].extend(n); mesh['i'].extend(v+start for v in i); mesh['parts']+=1
    binary=bytearray(); meshes=[]
    for (group,kind), mesh in merged.items():
        record={'group':group,'kind':kind,'sourceParts':mesh['parts']}
        for key in ['p','n','i']:
            record[key]={'offset':len(binary),'count':len(mesh[key])}
            binary.extend(mesh[key].tobytes())
        meshes.append(record)
    out=Path(output);out.mkdir(parents=True,exist_ok=True)
    compressed=gzip.compress(bytes(binary),compresslevel=9,mtime=0)
    (out/'anatomy.bin.gz').write_bytes(compressed)
    manifest={'version':1,'sourceParts':len(parts),'byteLength':len(binary),'sha256':hashlib.sha256(binary).hexdigest(),'groups':[x[0] for x in RULES],'meshes':meshes}
    (out/'anatomy.json').write_text(json.dumps(manifest,separators=(',',':')))
    print(f'{len(parts)} source parts → {len(meshes)} draw groups; {len(binary)} bytes, {len(compressed)} bytes compressed')
    return manifest

if __name__=='__main__': export(sys.argv[1],sys.argv[2])