"""Involute gear profiles as SVG, sized in millimetres for CAD import.

Generates true involute flanks, not approximations. A gear tooth is the curve
traced by a point on a string unwound from the base circle:

    x = rb (cos t + t sin t)      radius  = rb sqrt(1 + t^2)
    y = rb (sin t - t cos t)      polar a = t - atan(t)

Everything else is bookkeeping: where the flank starts and stops, and how far
round to rotate each tooth.

External and internal gears differ in three ways, all handled below:
  - the tip is outside the pitch circle for external, inside for internal
  - the tooth widens toward the root for external, toward the rim for internal
  - the internal gear needs an outer rim boundary to become a solid annulus

  python gear_svg.py --preset 4.5

Output is 1 SVG user unit = 1 mm, with width/height declared in mm, so Shapr3D
imports at true size. Each file is closed paths only, ready to extrude.
"""

import argparse
import math
from pathlib import Path

OUT = Path(__file__).parent / "gears"


def inv(a):
    return math.tan(a) - a


def involute_pts(rb, r_from, r_to, n=24):
    """Points along one involute flank, from r_from out to r_to."""
    def t_at(r):
        return math.sqrt(max(0.0, (r / rb) ** 2 - 1.0))
    t0, t1 = t_at(max(r_from, rb)), t_at(r_to)
    out = []
    for i in range(n + 1):
        t = t0 + (t1 - t0) * i / n
        out.append((rb * math.sqrt(1 + t * t), t - math.atan(t)))   # (radius, angle)
    return out


def polar(r, a):
    return (r * math.cos(a), r * math.sin(a))


def arc_pts(r, a0, a1, n=8):
    return [polar(r, a0 + (a1 - a0) * i / n) for i in range(n + 1)]


def external_gear(m, z, alpha=20.0, backlash=0.0, addendum=1.0, dedendum=1.25):
    """Closed CCW path for a standard external spur gear."""
    a = math.radians(alpha)
    r = m * z / 2.0
    rb = r * math.cos(a)
    ra = r + addendum * m
    rf = r - dedendum * m
    # Half tooth angle at the pitch circle, less half the backlash.
    half = math.pi / (2 * z) - (backlash / 2.0) / r
    delta = half + inv(a)

    flank = involute_pts(rb, max(rf, rb), ra)
    th_tip = flank[-1][1]
    th_root = flank[0][1]

    pts = []
    for k in range(z):
        phi = 2 * math.pi * k / z
        # Up the left flank (root -> tip).
        if rf < rb:
            pts.append(polar(rf, phi - delta))
        for rad, th in flank:
            pts.append(polar(rad, phi - (delta - th)))
        # Across the tip.
        pts += arc_pts(ra, phi - (delta - th_tip), phi + (delta - th_tip), 4)
        # Down the right flank (tip -> root).
        for rad, th in reversed(flank):
            pts.append(polar(rad, phi + (delta - th)))
        if rf < rb:
            pts.append(polar(rf, phi + delta))
        # Root arc across to the next tooth.
        nxt = 2 * math.pi * (k + 1) / z
        a0 = phi + (delta if rf < rb else (delta - th_root))
        a1 = nxt - (delta if rf < rb else (delta - th_root))
        pts += arc_pts(rf, a0, a1, 6)[1:-1]
    return pts, dict(r=r, rb=rb, ra=ra, rf=rf)


def internal_gear(m, z, alpha=20.0, backlash=0.0, addendum=1.0, dedendum=1.25):
    """Closed path for the toothed bore of a ring gear.

    Tip points inward (ra < r) and the tooth widens outward toward the rim,
    which is why the flank angle adds the involute term instead of subtracting.
    """
    a = math.radians(alpha)
    r = m * z / 2.0
    rb = r * math.cos(a)
    ra = r - addendum * m          # tip, inward
    rf = r + dedendum * m          # root, outward
    half = math.pi / (2 * z) - (backlash / 2.0) / r
    delta = half - inv(a)

    flank = involute_pts(rb, ra, rf)
    th_tip = flank[0][1]
    th_root = flank[-1][1]

    pts = []
    for k in range(z):
        phi = 2 * math.pi * k / z
        for rad, th in reversed(flank):                       # rim -> tip, left
            pts.append(polar(rad, phi - (delta + th)))
        pts += arc_pts(ra, phi - (delta + th_tip), phi + (delta + th_tip), 4)
        for rad, th in flank:                                 # tip -> rim, right
            pts.append(polar(rad, phi + (delta + th)))
        nxt = 2 * math.pi * (k + 1) / z
        pts += arc_pts(rf, phi + (delta + th_root),
                       nxt - (delta + th_root), 6)[1:-1]
    return pts, dict(r=r, rb=rb, ra=ra, rf=rf)


def circle_pts(rad, n=180, cw=False):
    step = -1 if cw else 1
    return [polar(rad, step * 2 * math.pi * i / n) for i in range(n)]


def path_d(pts):
    d = f"M {pts[0][0]:.4f} {pts[0][1]:.4f} "
    d += " ".join(f"L {x:.4f} {y:.4f}" for x, y in pts[1:])
    return d + " Z"


def write_svg(name, subpaths, extent, title=""):
    """One SVG, mm units, 1 user unit = 1 mm so CAD imports at true size."""
    OUT.mkdir(parents=True, exist_ok=True)
    pad = 2.0
    size = 2 * (extent + pad)
    body = "\n".join(
        f'  <path d="{path_d(p)}" fill="none" stroke="#000" stroke-width="0.2"/>'
        for p in subpaths)
    svg = f'''<svg xmlns="http://www.w3.org/2000/svg"
     width="{size:.3f}mm" height="{size:.3f}mm"
     viewBox="{-extent-pad:.3f} {-extent-pad:.3f} {size:.3f} {size:.3f}">
  <title>{title}</title>
{body}
</svg>
'''
    p = OUT / f"{name}.svg"
    p.write_text(svg)
    return p


def write_dxf(name, subpaths):
    """R12-style DXF: closed POLYLINEs in millimetres.

    DXF is the safer route into Shapr3D than SVG. It carries real units, so
    there is no 96-dpi scaling guesswork, and closed polylines import directly
    as sketch profiles ready to extrude. R12 POLYLINE/VERTEX is more verbose
    than LWPOLYLINE but every CAD package reads it.
    """
    e = ["0\nSECTION\n2\nHEADER\n"
         "9\n$ACADVER\n1\nAC1009\n"
         "9\n$INSUNITS\n70\n4\n"          # 4 = millimetres
         "0\nENDSEC\n",
         "0\nSECTION\n2\nENTITIES\n"]
    for pts in subpaths:
        e.append("0\nPOLYLINE\n8\nGEAR\n66\n1\n70\n1\n"
                 "10\n0.0\n20\n0.0\n30\n0.0\n")
        for x, y in pts:
            e.append(f"0\nVERTEX\n8\nGEAR\n10\n{x:.5f}\n20\n{y:.5f}\n30\n0.0\n")
        e.append("0\nSEQEND\n8\nGEAR\n")
    e.append("0\nENDSEC\n0\nEOF\n")
    p = OUT / f"{name}.dxf"
    p.write_text("".join(e))
    return p


def make_set(m, z_sun, z_planet, z_ring, n_planets, bore_sun, bore_planet,
             rim, backlash, tag):
    made = []

    pts, g = external_gear(m, z_sun, backlash=backlash)
    subs = [pts]
    if bore_sun:
        subs.append(circle_pts(bore_sun / 2, cw=True))
    p = write_svg(f"{tag}_sun_{z_sun}T", subs, g["ra"],
                  f"Sun {z_sun}T module {m} - pitch dia {2*g['r']:.2f}mm")
    write_dxf(f"{tag}_sun_{z_sun}T", subs)
    made.append((p, f"sun {z_sun}T", 2 * g["r"], 2 * g["ra"]))

    pts, g = external_gear(m, z_planet, backlash=backlash)
    subs = [pts]
    if bore_planet:
        subs.append(circle_pts(bore_planet / 2, cw=True))
    p = write_svg(f"{tag}_planet_{z_planet}T", subs, g["ra"],
                  f"Planet {z_planet}T module {m} - pitch dia {2*g['r']:.2f}mm "
                  f"- need {n_planets}")
    write_dxf(f"{tag}_planet_{z_planet}T", subs)
    made.append((p, f"planet {z_planet}T x{n_planets}", 2 * g["r"], 2 * g["ra"]))

    pts, g = internal_gear(m, z_ring, backlash=backlash)
    outer = g["rf"] + rim
    p = write_svg(f"{tag}_ring_{z_ring}T", [circle_pts(outer), pts], outer,
                  f"Ring {z_ring}T module {m} internal - pitch dia {2*g['r']:.2f}mm")
    write_dxf(f"{tag}_ring_{z_ring}T", [circle_pts(outer), pts])
    made.append((p, f"ring {z_ring}T internal", 2 * g["r"], 2 * outer))
    return made


def assembly(m, z_sun, z_planet, z_ring, n_planets, rim, backlash, tag,
             bore_sun=0.0, bore_planet=0.0):
    """All parts in mesh, as a single sheet, for checking before you extrude.

    Includes the bores: this is the drawing you use to lay out the carrier, so
    it needs the planet pin positions on it.
    """
    a = m * (z_sun + z_planet) / 2.0
    subs = []
    pts, gs = external_gear(m, z_sun, backlash=backlash)
    subs.append(pts)
    if bore_sun:
        subs.append(circle_pts(bore_sun / 2, cw=True))
    pr, gp = external_gear(m, z_planet, backlash=backlash)
    for i in range(n_planets):
        ang = 2 * math.pi * i / n_planets
        # Rotate each planet so its teeth actually mesh with the sun at that
        # angular position, which is what the assembly condition guarantees.
        spin = ang * (1 + z_sun / z_planet)
        cx, cy = a * math.cos(ang), a * math.sin(ang)
        subs.append([(cx + x * math.cos(spin) - y * math.sin(spin),
                      cy + x * math.sin(spin) + y * math.cos(spin))
                     for x, y in pr])
        if bore_planet:
            subs.append([(cx + x, cy + y)
                         for x, y in circle_pts(bore_planet / 2, cw=True)])
    ring, gr = internal_gear(m, z_ring, backlash=backlash)
    # Half a ring tooth of offset, and the planets phased by the rolling law
    # theta*(1 + z_sun/z_planet). Both are needed or the teeth interpenetrate.
    ca, sa = math.cos(math.pi / z_ring), math.sin(math.pi / z_ring)
    ring = [(x * ca - y * sa, x * sa + y * ca) for x, y in ring]
    outer = gr["rf"] + rim
    subs += [ring, circle_pts(outer)]
    return write_svg(f"{tag}_ASSEMBLY", subs, outer,
                     f"{1 + z_ring/z_sun:.4f}:1 assembly check")


PRESETS = {
    "4.5":  dict(m=1.25, z_sun=28, z_planet=35, z_ring=98, n_planets=3),
    "4.43": dict(m=1.25, z_sun=28, z_planet=34, z_ring=96, n_planets=4),
    "3.0":  dict(m=1.0,  z_sun=38, z_planet=19, z_ring=76, n_planets=3),
    "6.0":  dict(m=1.25, z_sun=28, z_planet=56, z_ring=140, n_planets=3),
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--preset", default="4.5", choices=sorted(PRESETS))
    ap.add_argument("--bore-sun", type=float, default=8.0,
                    help="motor shaft bore, mm (0 for none)")
    ap.add_argument("--bore-planet", type=float, default=10.0,
                    help="planet bearing bore, mm (0 for none)")
    ap.add_argument("--rim", type=float, default=6.0, help="ring rim thickness, mm")
    ap.add_argument("--backlash", type=float, default=0.06,
                    help="total backlash at the pitch circle, mm")
    args = ap.parse_args()

    cfg = PRESETS[args.preset]
    tag = f"r{args.preset.replace('.', 'p')}"
    ratio = 1 + cfg["z_ring"] / cfg["z_sun"]

    made = make_set(bore_sun=args.bore_sun, bore_planet=args.bore_planet,
                    rim=args.rim, backlash=args.backlash, tag=tag, **cfg)
    asm = assembly(rim=args.rim, backlash=args.backlash, tag=tag,
                   bore_sun=args.bore_sun, bore_planet=args.bore_planet, **cfg)

    print(f"\n{ratio:.4f}:1  module {cfg['m']}  "
          f"{cfg['z_sun']}/{cfg['z_planet']}/{cfg['z_ring']}  "
          f"x{cfg['n_planets']} planets")
    print(f"centre distance {cfg['m']*(cfg['z_sun']+cfg['z_planet'])/2:.3f} mm   "
          f"backlash {args.backlash} mm\n")
    for p, what, pitch, od in made:
        print(f"  {p.name:26s} {what:22s} pitch {pitch:6.2f}mm  outline {od:6.2f}mm")
    print(f"  {asm.name:26s} {'all parts in mesh':22s}")
    print(f"\nwritten to {OUT}")


if __name__ == "__main__":
    main()
