Made of little things.

"""Prepare named BodyParts3D soft tissues for the shared MPFB rig.
Blender --background --factory-startup --python scripts/prepare-atlas.py -- ISA_ZIP
Does not modify the released skeleton asset or the user's Blender scene.
"""
from pathlib import Path
import sys, json, struct, re, hashlib, zipfile
import numpy as np
import bpy, bmesh
sys.path.insert(0,str(Path(__file__).resolve().parent))
from anatomy_fitting import bind_region, assert_stable_topology

ROOT = Path(__file__).resolve().parents[1]
SOURCE = Path(sys.argv[sys.argv.index('--') + 1])
archive = zipfile.ZipFile(SOURCE)
raw = (ROOT/'assets/skin/body-skin.bin').read_bytes()
length = struct.unpack('<I', raw[8:12])[0]; skin = json.loads(raw[12:12+length]); start = 12+length
spec = skin['arrays']['heads4']
heads = np.frombuffer(raw, dtype='<f4', count=spec['length'], offset=start+spec['offset']).reshape(-1,3)
rig = {b['name']: i for i,b in enumerate(skin['bones'])}
parents = [b['parent'] for b in skin['bones']]
entries = {}
for path in archive.namelist():
    if not path.endswith('.obj'): continue
    text = archive.read(path).decode()
    def header(name):
        match = re.search(r'^# '+name+r' : ([^\r\n]*)', text, re.M)
        return match.group(1).strip() if match else ''
    entries[Path(path).name] = {'label':header('English name'), 'id':header('Concept ID'), 'text':text, 'source':Path(path).name}

def geometry(entry):
    points=[]; faces=[]
    for line in entry['text'].splitlines():
        if line.startswith('v '): points.append([float(v) for v in line.split()[1:4]])
        elif line.startswith('f '):
            f=[int(v.split('/')[0])-1 for v in line.split()[1:]]
            faces.extend([[f[0], f[i], f[i+1]] for i in range(1,len(f)-1)])
    p=np.array(points,dtype=float)/1000
    p=p[:,[0,2,1]]; p[:,1]+=.08; p[:,2]=-p[:,2]-.075
    return p,faces

bone_points={e['label'].lower():geometry(e)[0] for file,e in entries.items() if file[2:-4].isdigit() and 3152<=int(file[2:-4])<=3395}
def endpoint(name, top=True):
    p=bone_points[name]; lo,hi=np.quantile(p[:,1],[.025,.975]); band=min(.015,(hi-lo)*.18)
    return p[p[:,1]>=hi-band].mean(axis=0) if top else p[p[:,1]<=lo+band].mean(axis=0)

# Atlas-space joints, paired with the skin rig's reference heads. Axial structures
# share the source coordinate frame; limbs use the same landmarks as the skeleton.
source_heads=heads.copy()
for side,suffix in [('left','l'),('right','r')]:
    for bone,name,top in [('upperarm','humerus',True),('lowerarm','ulna',True),('hand','ulna',False),('thigh','femur',True),('calf','tibia',True)]:
        source_heads[rig[f'{bone}_{suffix}']]=endpoint(f'{side} {name}',top)
    source_heads[rig[f'foot_{suffix}']]=bone_points[f'{side} talus'].mean(axis=0)
    source_heads[rig[f'ball_{suffix}']]=bone_points[f'{side} first metatarsal bone'].mean(axis=0)
    for digit,label in [('thumb','thumb'),('index','index finger'),('middle','middle finger'),('ring','ring finger'),('pinky','little finger')]:
        for level in [1,2,3]:
            name=f'{"proximal" if level==1 else "middle" if level==2 else "distal"} phalanx of {side} {label}'
            if digit=='thumb': name=f'{side} first metacarpal bone' if level==1 else f'{"proximal" if level==2 else "distal"} phalanx of {side} thumb'
            source_heads[rig[f'{digit}_{level:02}_{suffix}']]=endpoint(name)

def bind(points,rigid):
    # Axial organ geometry retains its source coordinates and proportions.
    indices=np.zeros((len(points),4),dtype='<u2');indices[:,0]=rig[rigid]
    weights=np.zeros((len(points),4),dtype='<f4');weights[:,0]=1
    return points,indices,weights

def region(indices,weights):
    score=np.bincount(indices.reshape(-1),weights=weights.reshape(-1),minlength=len(heads))
    name=skin['bones'][int(score.argmax())]['name']
    for suffix,side in [('l','left'),('r','right')]:
        if not name.endswith('_'+suffix):continue
        for stem,part in [('upperarm','upperarm'),('lowerarm','forearm'),('thigh','thigh'),('calf','calf'),('foot','foot'),('ball','foot'),('clavicle','shoulder')]:
            if name.startswith(stem):return f'{side}-{part}'
        return f'{side}-hand'
    return 'head' if name=='head' else 'neck' if name=='neck_01' else 'pelvis' if name=='pelvis' else 'torso'

def build(layer,selected):
    grouped={}
    for e,rigid,color in selected:
        # Multiple source pieces may share a concept; join them under its one ID.
        key=e['id'];assert key.startswith('FMA') and e['label']
        if key not in grouped: grouped[key]={'label':e['label'],'entries':[],'rigid':rigid,'color':color}
        grouped[key]['entries'].append(e)
    parts=[];catalog=[];blob=bytearray();files={}
    def write(array,dtype):
        while len(blob)%4:blob.append(0)
        offset=len(blob);blob.extend(np.array(array,dtype=dtype).tobytes());return offset
    for id,g in sorted(grouped.items()):
        points=[];faces=[]
        for e in g['entries']:
            p,f=geometry(e);offset=len(points);points.extend(p.tolist());faces.extend([[v+offset for v in face] for face in f])
            files[e['source']]=hashlib.sha256(e['text'].encode()).hexdigest()
        geometry_data=bpy.data.meshes.new(id);geometry_data.from_pydata(points,[],faces);geometry_data.update()
        welded=bmesh.new();welded.from_mesh(geometry_data)
        bmesh.ops.remove_doubles(welded,verts=list(welded.verts),dist=.000001)
        welded.to_mesh(geometry_data);welded.free();geometry_data.update()
        obj=bpy.data.objects.new(id,geometry_data);bpy.context.scene.collection.objects.link(obj);bpy.context.view_layer.objects.active=obj;obj.select_set(True)
        budget=650 if layer=='muscles' else 1400
        if len(geometry_data.polygons)>budget:
            mod=obj.modifiers.new('Web reduction','DECIMATE');mod.ratio=budget/len(geometry_data.polygons);bpy.ops.object.modifier_apply(modifier=mod.name)
        geometry_data.calc_loop_triangles()
        p=np.array([v.co[:] for v in geometry_data.vertices]);f=np.array([t.vertices[:] for t in geometry_data.loop_triangles])
        # Decimation can leave many loose source vertices. They are neither
        # drawable nor paintable and should not carry weight/storage overhead.
        used,inverse=np.unique(f,return_inverse=True);p=p[used];f=inverse.reshape(-1,3)
        p,indices,weights=bind(p,g['rigid']) if g['rigid'] else bind_region(p,g['label'],source_heads,heads,rig,parents)
        assert np.isfinite(p).all() and len(p)<65536 and np.allclose(weights.sum(axis=1),1)
        info={'id':id,'label':g['label'],'part':region(indices,weights),'color':g['color'],'vertices':len(p),'triangles':len(f),'positions':write(p,'<f4'),'indices':write(f,'<u2'),'skinIndices':write(indices,'<u2'),'skinWeights':write(weights,'<f4')}
        parts.append(info);catalog.append({'id':id,'label':g['label'],'triangles':len(f),'part':info['part'],'sources':[e['source'] for e in g['entries']]})
        bpy.data.objects.remove(obj,do_unlink=True);bpy.data.meshes.remove(geometry_data)
    id=f'bp3d-{layer}-1';meta=json.dumps({'id':id,'parts':parts},separators=(',',':')).encode();meta+=b' '*((-len(meta))%4)
    output=b'BJATLS01'+struct.pack('<I',len(meta))+meta+blob
    directory=ROOT/'assets/anatomy';directory.mkdir(exist_ok=True)
    assert_stable_topology(directory/f'{layer}.bin',parts,blob)
    (directory/f'{layer}.bin').write_bytes(output)
    provenance={'id':id,'source':'BodyParts3D 4.0 IS-A tree, 99% reduction','url':'https://dbarchive.biosciencedbc.jp/data/bodyparts3d/LATEST/isa_BP3D_4.0_obj_99.zip','archiveSha256':hashlib.sha256(SOURCE.read_bytes()).hexdigest(),'attribution':'BodyParts3D, © The Database Center for Life Science licensed under CC Attribution 4.0 International','license':'CC-BY-4.0','licenseSource':'https://dbarchive.biosciencedbc.jp/en/bodyparts3d/lic.html','licenseUpdated':'2025-02-27','preparationBlenderVersion':bpy.app.version_string,'fittingRevision':3,'modifications':'Region-constrained continuous muscle weights and shared wrist frame; Axis/unit conversion; coincident-vertex welding; decimation and unused-vertex removal; joining pieces with the same FMA concept; approximate landmark fit and blended rig weights for muscles, axial binding for organs; binary packaging. No clinical or biomechanical validation.','structures':len(parts),'triangles':sum(p['triangles'] for p in parts),'bytes':len(output),'sha256':hashlib.sha256(output).hexdigest(),'sourceFiles':files}
    (directory/f'{layer}-provenance.json').write_text(json.dumps(provenance,indent=2)+'\n')
    print(layer,{k:provenance[k] for k in ['structures','triangles','bytes','sha256']},flush=True)
    return catalog

# Explicit source ranges contain musculoskeletal structures. Exclude unnamed and
# non-muscle connective tissues rather than inventing labels for blank headers.
muscles=[]
for file,e in entries.items():
    number=int(re.search(r'\d+',file).group())
    if not 1383<=number<=1601 or not e['label'] or not e['id']:continue
    if any(w in e['label'].lower() for w in ['ligament','membrane','retinaculum','linea alba','tendinous arch','intermediate tendon']):continue
    color='#a66569' if 'tendon' not in e['label'].lower() and 'tract' not in e['label'].lower() else '#c6b6ad'
    muscles.append((e,None,color))

# Selected major visceral structures. Source names are retained, including lobes
# and walls instead of mislabeling incomplete geometry as an entire organ.
organ_files={
    'FJ2561':'spine_02','FJ2563':'spine_03','FJ2564':'spine_02','FJ2565':'pelvis',
    'FJ2566':'spine_01','FJ2567':'spine_01','FJ2571':'pelvis','FJ2572':'spine_02','FJ2573':'spine_02',
    'FJ1895':'spine_02','FJ2817':'spine_02','FJ3144':'spine_01','FJ3145':'spine_02',
    'FJ3146':'spine_01','FJ3147':'spine_02','FJ3149':'pelvis','FJ2541':'spine_03',
    'FJ2428':'spine_03','FJ2438':'spine_03','FJ2439':'spine_03',
    'FJ3129':'spine_02','FJ3130':'spine_02',
}
organs=[]
for file,e in entries.items():
    number=int(re.search(r'\d+',file).group());stem=file[:-4];name=e['label'].lower()
    rigid=organ_files.get(stem)
    if 2816<=number<=2824:rigid='spine_02'
    if 2574<=number<=2628:rigid='spine_01'
    if any(w in name for w in ['gyrus','lobule','occipital lobe','cerebellum']) and 1732<=number<=1842:rigid='head'
    if rigid and e['label'] and e['id']:
        color='#c89487' if rigid=='head' else '#bca5a1' if 'bronch' in name or 'lung' in name or 'pulmonary' in name else '#ad7274' if any(w in name for w in ['atrium','ventricle','kidney']) else '#a97961' if 2816<=number<=2824 else '#bbaa88'
        organs.append((e,rigid,color))
muscle_catalog=build('muscles',muscles);organ_catalog=build('organs',organs)
(ROOT/'src/atlas-catalog.ts').write_text('// Generated by scripts/prepare-atlas.py. Bump model IDs if topology changes.\nexport const muscleCatalog = '+json.dumps(muscle_catalog,separators=(',',':'))+';\nexport const organCatalog = '+json.dumps(organ_catalog,separators=(',',':'))+';\n')