"""Planetary gearbox sizing for a joint.

Sun driven by the motor, ring fixed to the housing, carrier is the output.
That layout gives

    ratio = 1 + N_ring / N_sun          and     N_ring = N_sun + 2 N_planet

so for 6:1 you need N_ring = 5 N_sun exactly, and the planets fill the gap.
The ring's pitch radius is m * N_ring / 2, which is why the ring diameter, not
the ratio, is what actually constrains the design.

  python gearbox.py --ratio 6 --ring-radius 60
"""

import argparse
import math


def combos(ratio, ring_radius_mm, n_planets=(3, 4, 5), modules=(0.8, 1.0, 1.25, 1.5, 2.0, 2.5, 3.0)):
    """Every tooth combination that hits the ratio exactly and fits the ring."""
    out = []
    k = ratio - 1.0                       # N_ring / N_sun
    for m in modules:
        n_ring = 2 * ring_radius_mm / m
        if abs(n_ring - round(n_ring)) > 1e-9:
            continue                      # ring teeth must be a whole number
        n_ring = int(round(n_ring))
        n_sun = n_ring / k
        if abs(n_sun - round(n_sun)) > 1e-9:
            continue
        n_sun = int(round(n_sun))
        if (n_ring - n_sun) % 2:
            continue                      # planets must be a whole number
        n_planet = (n_ring - n_sun) // 2
        if n_sun < 8 or n_planet < 8:
            continue

        for n in n_planets:
            # Equal angular spacing is only possible when this divides evenly.
            if (n_sun + n_ring) % n:
                continue
            # Adjacent planets must not touch: compare the gap between planet
            # centres against the planet tip diameter.
            a = m * (n_sun + n_planet) / 2.0          # centre distance
            gap = 2 * a * math.sin(math.pi / n)
            tip = m * (n_planet + 2)                  # planet outside diameter
            clearance = gap - tip
            if clearance <= 0.5:
                continue
            out.append({
                "module": m, "n_sun": n_sun, "n_planet": n_planet,
                "n_ring": n_ring, "planets": n,
                "ratio": 1 + n_ring / n_sun,
                "sun_pcd": m * n_sun, "planet_pcd": m * n_planet,
                "ring_pcd": m * n_ring, "centre": a,
                "clearance": clearance,
                "undercut": n_sun < 17,
            })
    return out


def tooth_stress(m_mm, n_sun, planets, torque_nm, face_mm=10.0):
    """Lewis bending stress at the sun mesh, the most loaded gear."""
    r = (m_mm * n_sun / 2) / 1000.0
    Ft = torque_nm / r / planets                 # load shared across planets
    Y = 0.29 + 0.0035 * min(n_sun, 40)           # rough Lewis form factor
    return Ft / ((face_mm / 1000) * (m_mm / 1000) * Y) / 1e6, Ft


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ratio", type=float, default=6.0)
    ap.add_argument("--ring-radius", type=float, default=60.0,
                    help="pitch radius of the ring gear, mm")
    ap.add_argument("--motor-torque", type=float, default=1.2, help="rated, Nm")
    ap.add_argument("--motor-peak", type=float, default=2.5, help="peak, Nm")
    ap.add_argument("--motor-rpm", type=float, default=7200, help="no-load at 24V")
    ap.add_argument("--face", type=float, default=10.0, help="face width, mm")
    args = ap.parse_args()

    R = args.ring_radius
    k = args.ratio - 1

    print(f"\nPLANETARY, {args.ratio:g}:1, sun in / carrier out / ring fixed")
    print("=" * 68)
    print(f"  ratio = 1 + N_ring/N_sun  ->  N_ring must be {k:g}x N_sun")
    print(f"  ring pitch radius {R:.0f} mm  ->  ring pitch dia {2*R:.0f} mm\n")

    print(f"  If {R:.0f} mm is the RING pitch radius:")
    S = R / k
    P = (R - S) / 2
    print(f"    sun pitch radius    {S:6.2f} mm  (dia {2*S:.1f} mm)")
    print(f"    planet pitch radius {P:6.2f} mm  (dia {2*P:.1f} mm)")
    print(f"    centre distance     {S + P:6.2f} mm\n")

    print(f"  If {R:.0f} mm were the SUN pitch radius instead:")
    print(f"    ring pitch radius would be {R*k:.0f} mm "
          f"-> a {2*R*k:.0f} mm gearbox. Not buildable on a leg.\n")

    rows = combos(args.ratio, R)
    if not rows:
        print("  No exact tooth combination fits. Try another ring radius.")
        return

    print(f"  {'mod':>5} {'sun':>4} {'plan':>5} {'ring':>5} {'x':>3} "
          f"{'centre':>7} {'clear':>6} {'stress':>8}  note")
    print("  " + "-" * 64)
    best = None
    for c in rows:
        st, Ft = tooth_stress(c["module"], c["n_sun"], c["planets"],
                              args.motor_peak, args.face)
        note = []
        if c["undercut"]:
            note.append(f"{c['n_sun']}T sun undercuts at 20deg PA")
        if st > 150:
            note.append("highly stressed")
        if not note:
            note.append("clean")
            if best is None or (c["module"] > best["module"]
                                and c["clearance"] > 2.0):
                best = c
        print(f"  {c['module']:5.2f} {c['n_sun']:4d} {c['n_planet']:5d} "
              f"{c['n_ring']:5d} {c['planets']:3d} {c['centre']:7.2f} "
              f"{c['clearance']:6.2f} {st:7.0f}MPa  {'; '.join(note)}")

    best = best or rows[0]
    print("\n  RECOMMENDED")
    print("  " + "-" * 64)
    print(f"    module {best['module']:g}, {best['planets']} planets")
    print(f"    sun     {best['n_sun']:3d} teeth   pitch dia {best['sun_pcd']:6.1f} mm  (on the motor shaft)")
    print(f"    planet  {best['n_planet']:3d} teeth   pitch dia {best['planet_pcd']:6.1f} mm  x{best['planets']}")
    print(f"    ring    {best['n_ring']:3d} teeth   pitch dia {best['ring_pcd']:6.1f} mm  (fixed to housing)")
    print(f"    planet centres on a {best['centre']:.2f} mm radius circle, "
          f"{360/best['planets']:.0f} deg apart")
    print(f"    exact ratio {best['ratio']:.4f} : 1")

    st, Ft = tooth_stress(best["module"], best["n_sun"], best["planets"],
                          args.motor_peak, args.face)
    print(f"\n    output torque   {args.motor_torque*best['ratio']*0.9:6.2f} Nm continuous, "
          f"{args.motor_peak*best['ratio']*0.9:.2f} Nm peak (90% efficient)")
    print(f"    output speed    {args.motor_rpm/best['ratio']:6.0f} rpm no-load "
          f"({args.motor_rpm/best['ratio']*2*math.pi/60:.0f} rad/s)")
    print(f"    tooth load      {Ft:6.0f} N per mesh at peak, "
          f"{st:.0f} MPa bending on a {args.face:.0f} mm face")


if __name__ == "__main__":
    main()