#!/usr/bin/env python3
"""Rebuild the Adom pins board lead GLB from the REAL Fusion board (Draco-decoded),
with: transparent FR4, real pins grouped per-node, 300um solder-paste balls on each pad,
baked silkscreen text geometry, animated contact shadows, and a SUBTLE insertion animation.
Fully renderable/verifiable offline (uncompressed)."""
import DracoPy, pygltflib, numpy as np, trimesh, sys
from matplotlib.textpath import TextPath
from matplotlib.font_manager import FontProperties, findfont
from shapely.geometry import Polygon, MultiPolygon
from shapely.ops import unary_union
FONT_PATH=findfont(FontProperties(family="DejaVu Sans", weight="bold"))

def text_mesh(s, height_mm, thick_mm=0.12):
    """Extruded silk text with REAL cut-out counters (holes in D,R,0,e...), normal +Z,
    reads correctly viewed from +Z (the wiki view). Centered on origin."""
    tp=TextPath((0,0), s, size=height_mm, prop=FontProperties(family="DejaVu Sans", weight="bold"))
    raw=[Polygon(pp).buffer(0) for pp in tp.to_polygons(closed_only=True) if len(pp)>=3]
    # even-odd fill: a contour is a HOLE if it sits inside an odd number of others.
    # Robust: union all solids, union all holes, subtract -> correct counters, nothing dropped.
    solids=[]; holes=[]
    for i,s_ in enumerate(raw):
        # depth = how many OTHER contours fully contain this one (full-polygon containment;
        # a centroid point would fall in the counter for O/D/e and misclassify the outline).
        depth=sum(1 for j,o in enumerate(raw) if j!=i and o.contains(s_))
        (holes if depth%2 else solids).append(s_)
    geom=unary_union(solids)
    if holes: geom=geom.difference(unary_union(holes))
    parts=list(geom.geoms) if isinstance(geom,MultiPolygon) else [geom]
    meshes=[trimesh.creation.extrude_polygon(p, height=thick_mm) for p in parts if p.area>1e-6]
    m=trimesh.util.concatenate(meshes)
    m.apply_translation([-(m.bounds[0][0]+m.bounds[1][0])/2, -(m.bounds[0][1]+m.bounds[1][1])/2, 0])
    return m

SRC = "/tmp/AdomBoardFusion.glb"          # raw real Fusion board (Draco)
OUT = sys.argv[1] if len(sys.argv)>1 else "/tmp/AdomBoard-lead2.glb"
LIFT_MM = float(sys.argv[2]) if len(sys.argv)>2 else 2.6   # gentle press-in, not a launch

g = pygltflib.GLTF2().load(SRC); blob = g.binary_blob()
mats = g.materials or []
def matcol(idx):
    if idx is None or idx>=len(mats): return [0.7,0.7,0.7,1.0]
    pbr = mats[idx].pbrMetallicRoughness
    return list(pbr.baseColorFactor) if pbr and pbr.baseColorFactor else [0.7,0.7,0.7,1.0]

# world transforms per mesh index (+ node name)
def world_mats(g):
    W={}
    def rec(ni,M):
        n=g.nodes[ni]; L=np.eye(4)
        if n.matrix: L=np.array(n.matrix,float).reshape(4,4).T
        elif n.translation: L[:3,3]=n.translation
        Mn=M@L
        if n.mesh is not None: W.setdefault(n.mesh,[]).append((Mn,n.name or ""))
        for c in (n.children or []): rec(c,Mn)
    for r in g.scenes[g.scene or 0].nodes: rec(r,np.eye(4))
    return W
W = world_mats(g)

# decode + split into board verts and per-pin clusters (meters)
pin_pts=[]; board_meshes=[]
PIN_CENTERS=np.array([[0.010,0.010],[0.040,0.010],[0.010,0.040],[0.040,0.040]])
pin_clusters={i:[] for i in range(4)}
for mi,mesh in enumerate(g.meshes):
    for p in mesh.primitives:
        ext=p.extensions["KHR_draco_mesh_compression"]; bv=g.bufferViews[ext["bufferView"]]
        off=bv.byteOffset or 0; dm=DracoPy.decode(blob[off:off+bv.byteLength])
        v=np.array(dm.points).reshape(-1,3); f=np.array(dm.faces).reshape(-1,3)
        col=matcol(p.material)
        for (M,nm) in W.get(mi,[(np.eye(4),"")]):
            vv=(M[:3,:3]@v.T).T+M[:3,3]
            is_pin = ("MpinContact" in nm) or nm.startswith("=>")
            if is_pin:
                c=vv[:,:2].mean(0); k=int(np.argmin(((PIN_CENTERS-c)**2).sum(1)))
                pin_clusters[k].append((vv,f,col))
            else:
                board_meshes.append((vv,f,col))

# board top surface z (max z of board/pad geometry)
board_top = max(vv[:,2].max() for vv,_,_ in board_meshes)
print("board top z (mm):", round(board_top*1000,3))

scene = trimesh.Scene()

ROT_CENTER=None   # once set to (cx,cy) below, every mesh is rotated 180 about the board-centre Z axis
def _rot(vv):
    if ROT_CENTER is None: return vv
    cx,cy=ROT_CENTER; out=np.asarray(vv,dtype=float).copy()
    out[:,0]=2*cx-out[:,0]; out[:,1]=2*cy-out[:,1]; return out
def add_mesh(vv,f,rgba,name,metal=False,double=True):
    vv=_rot(vv)
    m=trimesh.Trimesh(vertices=vv,faces=f,process=False)
    r,gr,b,a=rgba
    m.visual=trimesh.visual.TextureVisuals(material=trimesh.visual.material.PBRMaterial(
        baseColorFactor=[int(r*255),int(gr*255),int(b*255),int(a*255)],
        metallicFactor=0.9 if metal else 0.2, roughnessFactor=0.35 if metal else 0.55,
        alphaMode="BLEND" if a<0.99 else "OPAQUE", doubleSided=double))
    scene.add_geometry(m,node_name=name,geom_name=name)

# REPLACE the thin (0.3mm) Fusion FR4 with a proper 1.6mm slab; drop the flat Fusion pads
# (we rebuild copper as real 1oz layers below).
board_xy=np.vstack([vv[:,:2] for vv,f,col in board_meshes])
bx0,by0=board_xy.min(0); bx1,by1=board_xy.max(0)
ROT_CENTER=None   # 3D text reads correctly at the default camera WITHOUT rotating the assembly

FR4_T=0.0016   # 1.6mm standard FR4
board_bottom=board_top-FR4_T
slab=trimesh.creation.box(extents=[bx1-bx0, by1-by0, FR4_T])
slab.apply_translation([(bx0+bx1)/2,(by0+by1)/2, board_top-FR4_T/2])
cutters=[]
DRILL={0:0.00055,1:0.00055,2:0.001725,3:0.001725}   # med Ø1.1, lrg Ø3.45
PAD  ={0:0.00080,1:0.00080,2:0.00260,3:0.00260}      # med Ø1.6, lrg Ø5.2
for i,(cx,cy) in enumerate(PIN_CENTERS):
    cyl=trimesh.creation.cylinder(radius=DRILL[i], height=FR4_T*2.5)
    cyl.apply_translation([cx,cy,board_top-FR4_T/2]); cutters.append(cyl)
slab=slab.difference(trimesh.util.concatenate(cutters), engine='manifold')  # REAL through-holes
slab.fix_normals()   # outward normals so single-sided transparency renders the clean outer shell
# NOTE: the board is added at the very END with a silkscreen TEXTURE (flat print) instead of a
# solid colour, so the text/callouts are surface print (zero geometry) and cast NO shadow.

# REAL copper: 1oz (~35um) annular ring PROUD on the FR4 top AND bottom, + a plated via barrel
# lining the drilled hole (sticks out a hair top/bottom). Not buried in the FR4.
COPPER=[0.82,0.53,0.24,1.0]; CU_T=0.000035
for i,(cx,cy) in enumerate(PIN_CENTERS):
    dr,pr=DRILL[i],PAD[i]
    # top+bottom pad rings inner-radius = dr-CU_T so they OVERLAP the barrel wall (dr-CU_T..dr):
    # the plated hole is one continuous copper piece (pad flares into the barrel), not a pad merely
    # touching a separate tube.
    top=trimesh.creation.annulus(r_min=dr-CU_T, r_max=pr, height=CU_T)
    top.apply_translation([cx,cy,board_top+CU_T/2])            # sits ON the top surface
    add_mesh(top.vertices,top.faces,COPPER,f"cu_top_{i}",metal=True)
    bot=trimesh.creation.annulus(r_min=dr-CU_T, r_max=pr, height=CU_T)
    bot.apply_translation([cx,cy,board_bottom-CU_T/2])         # sticks out the bottom
    add_mesh(bot.vertices,bot.faces,COPPER,f"cu_bot_{i}",metal=True)
    barrel=trimesh.creation.annulus(r_min=dr-CU_T, r_max=dr, height=FR4_T+2*CU_T)
    barrel.apply_translation([cx,cy,(board_top+board_bottom)/2])  # plated hole wall, overlaps both rings
    add_mesh(barrel.vertices,barrel.faces,COPPER,f"cu_via_{i}",metal=True)
print(f"board: 1.6mm FR4 + 1oz copper (proud top/bottom rings + plated via barrels) x4")

# pins: merge each cluster into ONE node so it animates as a unit; detect barrel radius
pin_info={}
for k in range(4):
    parts=pin_clusters[k]
    if not parts: continue
    V=[]; F=[]; off=0
    for vv,f,col in parts:
        V.append(vv); F.append(f+off); off+=len(vv)
    V=np.vstack(V); F=np.vstack(F)
    c=V[:,:2].mean(0)
    barrel_r=float(np.hypot(V[:,0]-c[0], V[:,1]-c[1]).max())   # widest = socket barrel
    # SEAT the pin's collar ON the copper top: the collar (the ring near the board too wide to enter the
    # hole) must rest ON the proud copper pad, not sink to the FR4 level. Shift the pin up so the collar
    # bottom == copper top. The tail still passes through and out the bottom.
    rr=np.hypot(V[:,0]-c[0],V[:,1]-c[1])
    thr=0.0036 if barrel_r>0.0015 else 0.0014
    near=(V[:,2]>board_top-0.001)&(V[:,2]<board_top+0.003)&(2*rr>thr)
    if near.any():
        V=V.copy(); V[:,2]+=(board_top+CU_T)-V[near,2].min()
    pin_info[k]=(c, barrel_r)
    add_mesh(V,F,[0.90,0.66,0.30,1.0],f"PIN_{k}",metal=True)
    print(f"PIN_{k} center=({c[0]*1000:.1f},{c[1]*1000:.1f})mm barrel OD={barrel_r*2000:.2f}mm")

# 300um Essemtec jet solder-paste balls, ON the annular ring (between drill edge and pad edge),
# at ~0.45mm pitch (paste.rs standard). Real spec size = 0.30mm dia; do NOT oversize.
BALL_D=0.00030; BALL_R=BALL_D/2.0; PITCH=0.00045
# per-pin footprint geometry (meters): medium drill1.1/pad1.6 ; large drill3.45/pad5.2
FP={0:(0.00055,0.00080), 1:(0.00055,0.00080), 2:(0.001725,0.00260), 3:(0.001725,0.00260)}
for k,(c,barrel_r) in pin_info.items():
    drill_r,pad_r=FP[k]
    ring_r=(drill_r+pad_r)/2.0                 # centerline of the annular ring
    n=max(6, int(round(2*np.pi*ring_r/PITCH)))
    for j in range(n):
        th=2*np.pi*j/n
        s=trimesh.creation.icosphere(subdivisions=2,radius=BALL_R)
        s=s.slice_plane([0,0,0],[0,0,1],cap=True)   # HEMISPHERE: flat bottom sits on the pad
        s.apply_translation([c[0]+ring_r*np.cos(th), c[1]+ring_r*np.sin(th), board_top+CU_T])   # flat bottom ON the copper top surface, not sunk to FR4
        add_mesh(s.vertices,s.faces,[0.62,0.63,0.66,1.0],f"ball_{k}_{j}",metal=True)
    print(f"  balls PIN_{k}: {n} hemispheres x {BALL_D*1000:.2f}mm on ring r={ring_r*1000:.2f}mm")

# --- SILKSCREEN + CALLOUTS as 3D extruded text on the board top. These go through add_mesh, so
#     they get the SAME 180 rotation as the pins -> labels stay matched to pins AND read left-to-right
#     at the default camera. (Solid geometry, unlike the flat-texture attempt whose UV kept desyncing.)
# alpha 0.98 -> alphaMode BLEND -> viewer skips it for shadows (transparencyShadow off) => text casts NO shadow, still visually solid
WHITE=[0.97,0.97,0.97,0.98]; AMBER=[0.98,0.74,0.10,0.98]; STEEL=[0.60,0.80,0.97,0.98]
def place_text(txt, cx_mm, cy_mm, h_mm, color, name):
    t=text_mesh(txt, h_mm); t.apply_scale(0.001)
    t.apply_translation([cx_mm/1000.0, cy_mm/1000.0, board_top+0.00006])
    add_mesh(t.vertices, t.faces, color, name)
def arrow(sx,sy,tx,ty,color,name):
    z=board_top+0.00006
    lead=trimesh.creation.cylinder(radius=0.00016,
        segment=[[sx/1000.0,sy/1000.0,z],[tx/1000.0,ty/1000.0,z]])
    add_mesh(lead.vertices,lead.faces,color,name+"_lead")
    d=np.array([tx-sx,ty-sy],float); d/=np.linalg.norm(d); HH=0.0012
    head=trimesh.creation.cone(radius=0.00048,height=HH,sections=16)   # apex at +Z*HH
    zc=np.array([0,0,1.]); axd=np.array([d[0],d[1],0.]); axd/=np.linalg.norm(axd)
    v=np.cross(zc,axd); s_=np.linalg.norm(v); cth=np.dot(zc,axd)
    if s_>1e-6:
        vx=np.array([[0,-v[2],v[1]],[v[2],0,-v[0]],[-v[1],v[0],0]])
        head.apply_transform(np.block([[np.eye(3)+vx+vx@vx*((1-cth)/(s_*s_)),np.zeros((3,1))],[np.zeros(3),1]]))
    # place the APEX exactly on the tip (base sits back toward the callout)
    head.apply_translation([tx/1000.0-HH*d[0], ty/1000.0-HH*d[1], z])
    add_mesh(head.vertices,head.faces,color,name+"_arrow")

LABELS={0:("MP-MED-STD","hole Ø1.1  pin Ø1.2"), 1:("MP-MED-SHORT","hole Ø1.1  pin Ø1.2"),
        2:("MP-LRG-STD","hole Ø3.45  pin Ø3.6"), 3:("MP-LRG-SHORT","hole Ø3.45  pin Ø3.6")}
for k,(c,barrel_r) in pin_info.items():
    nm,sp=LABELS[k]; cx,cy=c[0]*1000,c[1]*1000; y0=cy-barrel_r*1000-2.8
    place_text(nm, cx, y0,     1.35, WHITE, f"silk_{k}_0")
    place_text(sp, cx, y0-1.9, 1.15, WHITE, f"silk_{k}_1")

def callout(title, ax, ay, spec, why, tx, ty, sx, sy, prefix, color):
    place_text(title, ax, ay+2.2, 2.3, color, prefix+"_t")
    place_text(spec,  ax, ay,     1.5, color, prefix+"_s")
    place_text(why,   ax, ay-2.0, 1.4, color, prefix+"_w")
    arrow(sx,sy,tx,ty,color,prefix)
# PRESS-FIT: medium -> pin @ (40,10); large -> pin @ (10,40)
callout("PRESS-FIT",21.5,18.0,"hole Ø1.1   pin Ø1.2","pin 0.1mm wider = interference fit",
        39.2,10.7, 30.0,14.0, "cmed", AMBER)
callout("PRESS-FIT",27.0,30.0,"hole Ø3.45   pin Ø3.6","pin 0.15mm wider = interference fit",
        12.6,38.4, 19.5,34.0, "clrg", AMBER)
# SOLDER PASTE -> nearest ball center on the large pad @ (40,40)
_pc=np.array([40.0,40.0]); _ca=np.array([24.0,45.5]); _u=(_ca-_pc)/np.linalg.norm(_ca-_pc)
_ringmm=(FP[3][0]+FP[3][1])/2.0*1000; _n=max(6,int(round(2*np.pi*(FP[3][0]+FP[3][1])/2.0/PITCH)))
_ang=round(np.arctan2(_u[1],_u[0])/(2*np.pi/_n))*(2*np.pi/_n)
_ball=_pc+_ringmm*np.array([np.cos(_ang),np.sin(_ang)])
callout("SOLDER PASTE",24.0,45.5,"300µm (0.3mm)","jet-dispensed paste balls",
        _ball[0],_ball[1], 30.0,45.0, "csol", STEEL)
print(f"  3D silk: 4 labels + 3 callouts; solder arrow -> ball ({_ball[0]:.1f},{_ball[1]:.1f})")

# transparent FR4 board (solid colour; the 3D text above carries the print)
add_mesh(slab.vertices, slab.faces, [0.05,0.42,0.12,0.46], "board_fr4", double=False)

scene.export(OUT)
print("exported",OUT)