"""Gravity counterbalance SKETCH (Cameron 2026-07-11): yoke_0 tail extension (behind, over the base)
with an eye hook + eye hook on connector_2's mid back face (servo_2 link) + spring drawn at zero pose.
Output: two_servo_cbtail.urdf (visual sketch — spring is a static cylinder, real one stretches)."""
import numpy as np, trimesh
import xml.etree.ElementTree as ET

TAIL_LEN, TAIL_W = 70.0, 14.0
def rpy_T(xyz, rpy):
    T = trimesh.transformations.euler_matrix(*rpy, axes="sxyz"); T[:3, 3] = xyz; return T
def eye(at, axis):
    t = trimesh.creation.torus(major_radius=6.0, minor_radius=2.5)
    if axis == "x": t.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2, [0,1,0]))
    if axis == "y": t.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2, [1,0,0]))
    t.apply_translation(at); return t
def U(p): return trimesh.boolean.union(p, engine="manifold")

# tail on yoke_0: horizontal beam extending back past the bridge (-x), eye at the tip
tail = trimesh.creation.box(extents=[TAIL_LEN, TAIL_W, TAIL_W])
tail.apply_translation([-53.0 - TAIL_LEN/2 + 6.0, -9.0, 0.0])
tail = U([tail, eye([-53.0 - TAIL_LEN + 2.0, -9.0, 0.0], "y")])
tail.export("gnode_cb_tail.stl")
tip_yoke0 = np.array([-53.0 - TAIL_LEN + 2.0, -9.0, 0.0]) / 1000.0
# saddle + eye on connector_2 mid, back (-x) face (servo_2 link frame)
c2 = trimesh.load("connector_beauty_2.stl", force="mesh")
lo, hi = c2.bounds; mid_z = (lo[2] + hi[2]) / 2
saddle = trimesh.creation.box(extents=[6.0, 20.0, 26.0])
saddle.apply_translation([lo[0] - 1.0, -8.8, mid_z])
hook2 = U([saddle, eye([lo[0] - 8.0, -8.8, mid_z], "y")])
hook2.export("gnode_cb_hook2.stl")
pt_s2 = np.array([lo[0] - 8.0, -8.8, mid_z]) / 1000.0
# chain transform servo_2 -> yoke_0 at zero pose (walk URDF joints)
arm = ET.parse("two_servo_beauty.urdf").getroot()
joints = {j.find("child").get("link"): j for j in arm.findall("joint")}
def T_to_world(link):
    T = np.eye(4); cur = link
    while cur in joints:
        j = joints[cur]; o = j.find("origin")
        T = rpy_T([float(v) for v in (o.get("xyz") or "0 0 0").split()],
                  [float(v) for v in (o.get("rpy") or "0 0 0").split()]) @ T
        cur = j.find("parent").get("link")
    return T
T_hook2_in_y0 = np.linalg.inv(T_to_world("yoke_0")) @ T_to_world("servo_2")
p2 = (T_hook2_in_y0 @ np.append(pt_s2, 1))[:3]
d = p2 - tip_yoke0; L = np.linalg.norm(d)
spring = trimesh.creation.cylinder(radius=0.004, height=L, sections=24)
z = np.array([0, 0, 1.0]); ax = np.cross(z, d / L)
if np.linalg.norm(ax) > 1e-6:
    spring.apply_transform(trimesh.transformations.rotation_matrix(
        np.arcsin(np.clip(np.linalg.norm(ax), -1, 1)) if d[2]/L >= 0 else np.pi - np.arcsin(np.linalg.norm(ax)), ax))
spring.apply_translation((tip_yoke0 + p2) / 2)
spring.apply_scale(1000.0)
spring.export("gnode_cb_spring.stl")
print(f"spring at zero pose: length {L*1000:.0f}mm, yoke_0 tip {np.round(tip_yoke0*1000,0).tolist()} -> hook2 {np.round(p2*1000,0).tolist()}")
# URDF
t = ET.parse("two_servo_beauty.urdf"); r = t.getroot(); r.set("name", "two_servo_cbtail")
def vis(link_name, fn, rgba, scale="0.001 0.001 0.001"):
    link = next(l for l in r.findall("link") if l.get("name") == link_name)
    v = ET.SubElement(link, "visual"); ET.SubElement(v, "origin", {"xyz": "0 0 0", "rpy": "0 0 0"})
    ET.SubElement(ET.SubElement(v, "geometry"), "mesh", {"filename": fn, "scale": scale})
    ET.SubElement(ET.SubElement(v, "material", {"name": "m" + fn}), "color", {"rgba": rgba})
vis("yoke_0", "gnode_cb_tail.stl", "0.85 0.35 0.25 1")
vis("servo_2", "gnode_cb_hook2.stl", "0.85 0.35 0.25 1")
vis("yoke_0", "gnode_cb_spring.stl", "0.75 0.75 0.30 1")
t.write("two_servo_cbtail.urdf", encoding="unicode", xml_declaration=True)
print("wrote two_servo_cbtail.urdf")