#!/usr/bin/env python3
"""planetary.py — print-in-place herringbone planetary gearbox, parametric.

One STL, printed assembled: sun + 3 planets + ring gear, already meshed.
Herringbone teeth lock the planets axially, so nothing falls out and no
carrier is needed for the demo. Cannot be made subtractively: the gears are
inside each other from the first layer.

Everything that matters is a parameter at the top. CLR is the one you set from
the design-rule coupon (the gap that printed free without fusing).

Ratios (ring fixed, sun in, planet carrier out):  1 + Zr/Zs
"""
import math, json
import cadquery as cq

# ---------------- parameters ----------------
M      = 1.0      # module, mm
ZS, ZP = 12, 12   # sun, planet teeth  (ring = ZS + 2*ZP)
NPL    = 3        # planets; needs (ZS+ZR) % NPL == 0
H      = 12.0     # total gear height, mm
CLR    = 0.35     # printed clearance per side, mm  <- from your coupon
PA     = 20.0     # pressure angle, deg
TWIST  = 22.0     # herringbone twist of the PLANET per half-height, deg
WALL   = 3.5      # ring gear wall thickness beyond root, mm
SUN_HEX = 5.0     # hex socket across flats in the sun (for a key or servo horn adapter)
CHAMFER_BOTTOM = 0.3   # elephant-foot relief on the bottom layer

ZR = ZS + 2 * ZP
assert (ZS + ZR) % NPL == 0, "planets will not assemble evenly"
RATIO = 1 + ZR / ZS

# ---------------- involute profile ----------------
def gear_profile(z, m, backlash=0.0, internal=False, n=7):
    """2D polygon of an external involute gear. backlash > 0 thins the teeth.
    For the ring we cut an 'internal' profile: addendum/dedendum swapped, teeth thickened."""
    pa = math.radians(PA)
    rp = m * z / 2
    rb = rp * math.cos(pa)
    if internal:
        ra = rp + 1.25 * m + CLR    # cutter tip -> ring root, deepened for tip clearance
        rf = rp - 1.0 * m           # cutter root -> ring tip
    else:
        ra = rp + 1.0 * m
        rf = rp - 1.25 * m - CLR    # root deepened so mating tips clear
    # half tooth thickness angle at pitch circle, with backlash
    half_t = (math.pi / (2 * z)) - (backlash / 2) / rp
    inv = lambda a: math.tan(a) - a
    phi_p = inv(pa)
    pts = []
    for i in range(z):
        c = 2 * math.pi * i / z
        flank = []
        rr = max(rf, rb)
        for k in range(n + 1):
            r = rr + (ra - rr) * k / n
            a = math.acos(min(1.0, rb / r))
            th = half_t + phi_p - inv(a)
            flank.append((r, th))
        # root point (if root is below base circle, drop radially)
        up = [(rf, flank[0][1])] + flank if rf < rb else flank      # root -> tip, |angle| decreasing
        # traverse with increasing angle: trailing flank up (at -th), leading flank down (at +th)
        tooth = [(r, -th) for (r, th) in up] + [(r, th) for (r, th) in reversed(up)]
        for (r, th) in tooth:
            pts.append((r * math.cos(c + th), r * math.sin(c + th)))
    return pts

def herringbone(profile_pts, h, twist_deg):
    """Two half-height twist extrudes with opposite hand -> herringbone."""
    half = h / 2
    lo = cq.Workplane("XY").polyline(profile_pts).close().twistExtrude(half, twist_deg)
    hi = (cq.Workplane("XY").workplane(offset=half)
          .polyline(_rot(profile_pts, twist_deg)).close().twistExtrude(half, -twist_deg))
    return lo.union(hi)

def _rot(pts, deg):
    a = math.radians(deg); c, s = math.cos(a), math.sin(a)
    return [(x * c - y * s, x * s + y * c) for x, y in pts]

# twist must scale as 1/Z so helix angles match at the pitch circles:
tw_p = TWIST
tw_s = -TWIST * ZP / ZS
tw_r = TWIST * ZP / ZR

# ---------------- sun ----------------
sun = herringbone(gear_profile(ZS, M, backlash=CLR), H, tw_s)
sun = sun.cut(cq.Workplane("XY").polygon(6, SUN_HEX / math.cos(math.pi / 6)).extrude(H))

# ---------------- planets ----------------
rc = M * (ZS + ZP) / 2            # carrier radius
planets = None
for i in range(NPL):
    ang = 2 * math.pi * i / NPL
    # phase the planet so its teeth mesh with the sun at this position
    phase = math.degrees(ang) * (ZS / ZP) * 1.0 + (180.0 / ZP)   # half tooth offset
    prof = _rot(gear_profile(ZP, M, backlash=CLR), phase)
    p = herringbone(prof, H, tw_p).translate((rc * math.cos(ang), rc * math.sin(ang), 0))
    planets = p if planets is None else planets.union(p)

# ---------------- ring ----------------
r_root_ring = M * ZR / 2 + 1.25 * M
r_out = r_root_ring + WALL
cutter = herringbone(_rot(gear_profile(ZR, M, backlash=-CLR, internal=True), 180.0 / ZR), H, tw_r)
ring = cq.Workplane("XY").circle(r_out).extrude(H).cut(cutter)
# scalloped grip: part of the ring's outline, one extrude, no boolean loop
def ring_outline(r, n_scallop=12, depth=0.9, pts_per=10):
    out = []
    for i in range(n_scallop * pts_per):
        t = 2 * math.pi * i / (n_scallop * pts_per)
        rr = r - depth * max(0.0, math.cos(n_scallop * t)) ** 2
        out.append((rr * math.cos(t), rr * math.sin(t)))
    return out
ring = cq.Workplane("XY").polyline(ring_outline(r_out)).close().extrude(H).cut(cutter)

asm = ring.union(sun).union(planets)

cq.exporters.export(asm, "planetary.stl", tolerance=0.01, angularTolerance=0.1)
cq.exporters.export(asm, "planetary.step")

vol = asm.val().Volume() / 1000.0
bb = asm.val().BoundingBox()
info = dict(module=M, sun=ZS, planet=ZP, ring=ZR, planets=NPL, ratio=RATIO, clearance=CLR,
            height=H, outer_d=2 * r_out, volume_cm3=round(vol, 1),
            bbox_mm=[round(bb.xlen, 1), round(bb.ylen, 1), round(bb.zlen, 1)],
            twist_deg=dict(sun=tw_s, planet=tw_p, ring=tw_r))
json.dump(info, open("planetary_info.json", "w"), indent=1)
print(json.dumps(info, indent=1))