# cycloid_to_dxf_overlays.py
# Generate a cycloidal disk outline and export as DXF with optional overlays.
# Example:
#   python cycloid_to_dxf_overlays.py --N 12 --R 35 --Rr 4 --E 1.5 \
#     --points 2000 --out cycloid.dxf \
#     --pitch --rollers --eccentric \
#     --bc-radius 20 --bc-count 6 --bc-hole-diam 3 \
#     --bore 8

import math
import argparse
import numpy as np
import ezdxf

# ---------------------------
# Geometry generators
# ---------------------------

def cycloid_profile(N: int, R: float, Rr: float, E: float, points: int = 2000):
    """
    Parametric cycloidal disk profile (mm).
      N  - rollers count
      R  - roller pitch circle radius
      Rr - roller radius
      E  - eccentricity
    """
    t = np.linspace(0.0, 2.0 * math.pi, points, endpoint=True)
    one_minus_N_t = (1 - N) * t
    num = np.sin(one_minus_N_t)
    den = (R / (E * N)) - np.cos(one_minus_N_t)
    phi = np.arctan2(num, den)

    x = (R * np.cos(t)
         - Rr * np.cos(t + phi)
         - E * np.cos(N * t))
    y = (-R * np.sin(t)           # sign matches the traditional formula you used
         + Rr * np.sin(t + phi)
         + E * np.sin(N * t))
    return x, y

def add_polyline(msp, pts_xy, layer="PROFILE", close=True):
    msp.add_lwpolyline(pts_xy, format="xy", close=close, dxfattribs={"layer": layer})

def add_circle(msp, center_xy, radius, layer="REF"):
    msp.add_circle((center_xy[0], center_xy[1], 0.0), radius, dxfattribs={"layer": layer})

def add_point(msp, p, layer="REF"):
    msp.add_point((p[0], p[1], 0.0), dxfattribs={"layer": layer})

# ---------------------------
# DXF builder
# ---------------------------

def build_dxf(
    outline_xy,
    outfile: str,
    *,
    add_pitch=False,
    R=None,
    add_rollers=False,
    N=None,
    Rr=None,
    add_eccentric=False,
    E=None,
    bc_radius=None,
    bc_count=None,
    bc_hole_diam=None,
    center_bore=None,
):
    doc = ezdxf.new(dxfversion="R2010")
    doc.header["$INSUNITS"] = 4
    msp = doc.modelspace()

    # Layers
    for name, color in [
        ("PROFILE", 2),       # yellow
        ("PITCH_CIRCLE", 5),  # blue
        ("ROLLERS", 3),       # green
        ("ECCENTRIC", 1),     # red
        ("BOLT_CIRCLE", 4),   # cyan
        ("HOLES", 7),         # white
        ("REF", 8),           # gray
    ]:
        if name not in doc.layers:
            doc.layers.add(name, color=color)

    # Outline
    add_polyline(msp, outline_xy, layer="PROFILE", close=True)

    # Reference: origin point
    add_point(msp, (0.0, 0.0), layer="REF")

    # Optional overlays
    if add_pitch and R is not None:
        add_circle(msp, (0.0, 0.0), R, layer="PITCH_CIRCLE")

    if add_eccentric and E is not None:
        add_circle(msp, (0.0, 0.0), E, layer="ECCENTRIC")

    if add_rollers and (N is not None) and (R is not None) and (Rr is not None):
        # Place N roller circles of radius Rr, centers on the pitch circle.
        for i in range(N):
            th = 2.0 * math.pi * i / N
            # Match the y-sign convention used in the parametric (clockwise y):
            cx = R * math.cos(th)
            cy = -R * math.sin(th)
            add_circle(msp, (cx, cy), Rr, layer="ROLLERS")

    # Bolt circle + holes
    if bc_radius and bc_count and bc_hole_diam:
        add_circle(msp, (0.0, 0.0), bc_radius, layer="BOLT_CIRCLE")
        r_hole = bc_hole_diam / 2.0
        for i in range(bc_count):
            th = 2.0 * math.pi * i / bc_count
            cx = bc_radius * math.cos(th)
            cy = -bc_radius * math.sin(th)
            add_circle(msp, (cx, cy), r_hole, layer="HOLES")

    # Center bore
    if center_bore and center_bore > 0:
        add_circle(msp, (0.0, 0.0), center_bore / 2.0, layer="HOLES")

    doc.saveas(outfile)

# ---------------------------
# Main
# ---------------------------

def main():
    ap = argparse.ArgumentParser(description="Export cycloidal disk outline to DXF with overlays.")
    ap.add_argument("--N", type=int, default=12, help="Number of rollers (N)")
    ap.add_argument("--R", type=float, default=35.0, help="Pitch circle radius R (mm)")
    ap.add_argument("--Rr", type=float, default=4.0, help="Roller radius Rr (mm)")
    ap.add_argument("--E", type=float, default=1.5, help="Eccentricity E (mm)")
    ap.add_argument("--points", type=int, default=2000, help="Number of sampled points for outline")
    ap.add_argument("--out", type=str, default="cycloid.dxf", help="Output DXF filename")

    # Overlays
    ap.add_argument("--pitch", action="store_true", help="Add pitch circle (radius R)")
    ap.add_argument("--rollers", action="store_true", help="Add N roller circles of radius Rr on pitch circle")
    ap.add_argument("--eccentric", action="store_true", help="Add eccentric circle (radius E)")
    ap.add_argument("--bc-radius", type=float, default=None, help="Bolt circle radius (mm)")
    ap.add_argument("--bc-count", type=int, default=None, help="Bolt circle hole count")
    ap.add_argument("--bc-hole-diam", type=float, default=None, help="Bolt circle hole diameter (mm)")
    ap.add_argument("--bore", type=float, default=None, help="Center bore diameter (mm)")

    args = ap.parse_args()

    # Build outline
    x, y = cycloid_profile(args.N, args.R, args.Rr, args.E, points=args.points)
    outline_xy = list(zip(x.tolist(), y.tolist()))

    build_dxf(
        outline_xy,
        args.out,
        add_pitch=args.pitch,
        R=args.R,
        add_rollers=args.rollers,
        N=args.N,
        Rr=args.Rr,
        add_eccentric=args.eccentric,
        E=args.E,
        bc_radius=args.bc_radius,
        bc_count=args.bc_count,
        bc_hole_diam=args.bc-hole_diam if hasattr(args, "bc-hole_diam") else args.bc_hole_diam,  # safety
        center_bore=args.bore,
    )

    print(f"✅ Wrote {args.out}")
    print(
        f"    Params: N={args.N}, R={args.R} mm, Rr={args.Rr} mm, E={args.E} mm, points={args.points}\n"
        f"    Overlays: pitch={args.pitch}, rollers={args.rollers}, eccentric={args.eccentric}, "
        f"BC(radius={args.bc_radius}, count={args.bc_count}, hole_diam={args.bc_hole_diam}), "
        f"bore={args.bore}"
    )

if __name__ == "__main__":
    main()