Made of little things.

"""Offline, region-constrained atlas fitting. Units are meters, +X patient left.
No topology operations: existing triangle/barycentric addresses remain valid.
"""
import numpy as np

def rotation(a,b):
    a=a/np.linalg.norm(a);b=b/np.linalg.norm(b);v=np.cross(a,b);c=a@b
    k=np.array([[0,-v[2],v[1]],[v[2],0,-v[0]],[-v[1],v[0],0]])
    return np.eye(3)+k+k@k/max(1+c,1e-9)

def segment_matrix(a,b,c,d):
    old=b-a;new=d-c;axis=old/np.linalg.norm(old)
    matrix=rotation(old,new)@(np.eye(3)+np.outer(axis,axis)*(np.linalg.norm(new)/np.linalg.norm(old)-1))
    return matrix,c-matrix@a

def palm_matrix(source,heads,rig,suffix):
    def frame(points):
        wrist=points[rig[f'hand_{suffix}']];middle=points[rig[f'middle_01_{suffix}']]-wrist
        across=points[rig[f'index_01_{suffix}']]-points[rig[f'pinky_01_{suffix}']]
        y=middle/np.linalg.norm(middle);x=across-y*(across@y);width=np.linalg.norm(x);x/=width
        return wrist,np.column_stack([x,y,np.cross(x,y)]),np.array([width,np.linalg.norm(middle),width])
    a,old,oldsize=frame(source);b,new,newsize=frame(heads)
    matrix=new@np.diag(newsize/oldsize)@old.T
    return matrix,b-matrix@a

def candidates(label,points,rig):
    name=label.lower();suffix='l' if 'left' in name else 'r' if 'right' in name else None
    axial=[1,2,3,4,43,44]
    if not suffix:return axial
    side=lambda names:[rig[f'{n}_{suffix}'] for n in names]
    if any(s in name for s in ['oblique','intercostal','transversus thoracis','serratus posterior','spinal','longissimus','iliocostalis','rotator','intertransvers','colli','capitis','cervic','hyoid','platysma','digastric','scalen','sternocleidomastoid','levatores costarum']):
        # The hand/foot adductor oblique heads are handled by their location.
        if 'pollicis' not in name and 'hallucis' not in name:return axial
    if any(s in name for s in ['pubococcygeus','puborectalis','iliococcygeus','coccygeus']):return [1]
    if any(s in name for s in ['pectoralis','deltoid','subscapularis','supraspinatus','infraspinatus','teres major','teres minor','trapezius','rhomboid','subclavius','serratus anterior','levator scapulae']):return [3,4,43]+side(['clavicle','upperarm'])
    if any(s in name for s in ['gluteus','iliacus','psoas','obturator','gemellus','quadratus femoris','piriformis']):return [1,2]+side(['thigh'])
    if any(s in name for s in ['sartorius','semitendinosus','semimembranosus','tensor fascia','pectineus','adductor brevis','adductor longus','adductor magnus','adductor minimus','rectus femoris','vastus','gracilis','biceps femoris','iliotibial']):return [1]+side(['thigh','calf'])
    if 'foot' in name or 'hallucis brevis' in name or 'adductor hallucis' in name or 'abductor hallucis' in name or 'flexor accessorius' in name or 'flexor digitorum brevis' in name:return side(['foot','ball'])
    if any(s in name for s in ['tibialis','fibularis','soleus','plantaris','popliteus','gastrocnemius','calcaneal','hallucis longus','digitorum longus','hallucis brevis']):return side(['thigh','calf','foot','ball'])
    # Remaining named upper-limb structures cannot bind to torso or legs even
    # when a finger lies next to the thigh in the source's neutral pose.
    if points[:,1].mean()>1.05:return side(['clavicle','upperarm','lowerarm','hand'])
    return [i for i,b in enumerate(rig) if b.endswith('_'+suffix) and not any(s in b for s in ['thigh','calf','foot','ball','clavicle'])]

def bind_region(points,label,source,heads,rig,parents):
    tails=[next((j for j,p in enumerate(parents) if p==i),-1) for i in range(len(heads))]
    for suffix in ['l','r']:tails[rig[f'hand_{suffix}']]=rig[f'middle_01_{suffix}']
    end=lambda i: source[tails[i]] if tails[i]>=0 else source[i]+(source[i]-source[parents[i]])*.72
    allowed=candidates(label,points,rig)
    distances=[]
    for i in allowed:
        a=source[i];v=end(i)-a;t=np.clip((points-a)@v/max(v@v,1e-8),0,1)
        distances.append(np.sum((points-a-t[:,None]*v)**2,axis=1))
    d=np.array(distances).T;nearest=np.argsort(d,axis=1)[:,:4];chosen=np.take_along_axis(d,nearest,axis=1)
    # Continuous positive weights: a hard distance cutoff made sharp creases at
    # the influence boundary. Restriction is per anatomical structure instead.
    w=1/np.maximum(chosen,.0004)**2;w/=w.sum(1)[:,None]
    indices=np.zeros((len(points),4),dtype='<u2');weights=np.zeros((len(points),4),dtype='<f4')
    indices[:,:len(nearest[0])]=np.array(allowed)[nearest];weights[:,:len(nearest[0])]=w
    transforms=[]
    for i in range(len(heads)):
        if i in [0,1,2,3,4,43,44]:transforms.append((np.eye(3),heads[i]-source[i]));continue
        tail=tails[i];newend=heads[tail] if tail>=0 else heads[i]+(heads[i]-heads[parents[i]])*.72
        transforms.append(segment_matrix(source[i],end(i),heads[i],newend))
    # The same palm frame as the skeleton avoids an independently twisted hand.
    for suffix in ['l','r']:transforms[rig[f'hand_{suffix}']]=palm_matrix(source,heads,rig,suffix)
    fitted=np.zeros_like(points)
    for i,(A,b) in enumerate(transforms):
        contribution=np.sum(np.where(indices==i,weights,0),axis=1)
        if np.any(contribution):
            transformed=points@A.T+b
            if i in [rig['lowerarm_l'],rig['lowerarm_r']]:
                suffix='l' if i==rig['lowerarm_l'] else 'r';j=rig[f'hand_{suffix}']
                transformed=forearm_points(points,source[i],source[j],heads[i],heads[j],palm_matrix(source,heads,rig,suffix))
            fitted+=transformed*contribution[:,None]
    return fitted,indices,weights

def forearm_points(points,a,b,target_a,target_b,palm):
    """Blend the forearm's proximal frame into the palm frame at the wrist.
    The atlas palm faces forward while the display rig palm faces inward. Both
    radius/ulna use the same transition, preserving their shared wrist contact.
    This is an illustrative rest fit, not simulated pronation mechanics.
    """
    A,shift=segment_matrix(a,b,target_a,target_b);palmA,palmb=palm
    t=np.clip((points-a)@(b-a)/np.dot(b-a,b-a),0,1)
    t=np.clip((t-.45)/.55,0,1);t=t*t*(3-2*t)
    return (points@A.T+shift)*(1-t[:,None])+(points@palmA.T+palmb)*t[:,None]

def correct_skeleton_part(points,name,part,bone,tail,heads,source,rig,parents,endpoint):
    side='left' if part.startswith('left') else 'right';suffix='l' if side=='left' else 'r'
    if part.endswith('-forearm'):
        a=source[rig[f'lowerarm_{suffix}']];b=source[rig[f'hand_{suffix}']];ta=heads[bone];tb=heads[tail]
        A,shift=segment_matrix(a,b,ta,tb);original=(points+ta-shift)@np.linalg.inv(A).T
        return forearm_points(original,a,b,ta,tb,palm_matrix(source,heads,rig,suffix))-ta,tail
    if not part.endswith('-hand'):return points,tail
    ta=heads[bone];tb=heads[tail] if tail>=0 else ta+(ta-heads[parents[bone]])*.72
    finger=any(w in name for w in ['finger','thumb','metacarpal'])
    if finger:a,b=endpoint(name),endpoint(name,False)
    else:a=endpoint(f'{side} ulna',False);b=a+np.array([0,-.08,0]);tb=heads[rig[f'middle_01_{suffix}']]
    A,shift=segment_matrix(a,b,ta,tb);original=(points+ta-shift)@np.linalg.inv(A).T
    palmA,palmb=palm_matrix(source,heads,rig,suffix);fitted=original@palmA.T+palmb
    if finger and 'metacarpal' not in name:
        A,shift=segment_matrix(palmA@a+palmb,palmA@b+palmb,ta,tb);fitted=fitted@A.T+shift
    elif 'metacarpal' in name and 'first' not in name:tail=rig[f'middle_01_{suffix}']
    return fitted-ta,tail

def assert_stable_topology(path,parts,blob):
    """Reject a re-export that would relocate existing saved triangle addresses."""
    import json,struct
    if not path.exists():return
    raw=path.read_bytes();start=12+struct.unpack('<I',raw[8:12])[0];old=json.loads(raw[12:start])
    byid={p['id']:p for p in old['parts']}
    for p in parts:
        if p['id'] not in byid:continue
        previous=byid[p['id']]
        assert (p['vertices'],p['triangles'])==(previous['vertices'],previous['triangles']),f"Topology changed: {p['label']}; version anchors before exporting"
        original=np.frombuffer(raw,'<u2',previous['triangles']*3,start+previous['indices'])
        fresh=np.frombuffer(blob,'<u2',p['triangles']*3,p['indices'])
        assert np.array_equal(original,fresh),f"Triangle addresses changed: {p['label']}; version anchors before exporting"