"""Ball screw with a DIN 69051 Form B flanged single nut, TBI MOTION's SFU and SFNU series: the screw shaft by length, with its ball track
(detail = full) or plain at the screw's outside diameter (detail = simple), or the flanged ball nut. Shaft axis along Z from the XY plane up. The nut's axis along Z too, its flange face on the XY
plane and the body rising from it, the two flats parallel to Y. Dimensions from TBI MOTION's catalogue, listed in the README."""
import math

import cadquery as cq
from steve import Model

model_no = "SFNU01605-4"  # TBI MOTION model number
part = "screw"  # screw (the shaft) or nut (the flanged ball nut)
length = 500  # mm Shaft length (screw only)
detail = "full"  # full: the shaft with its ball track; simple: a plain shaft at the outside diameter d

# model: (screw diameter d, lead I, ball Da, nut body D, flange A, flange thickness B, nut length L, bolt circle W, across the flats H,
#         bolt hole X), mm. TBI MOTION Ball Screw catalogue (copy hosted by Anaheim Automation, L011848, PDF of 2025-03-12), "SFNU/SFU (DIN 69051
#         FORM B) Series Specifications": rolled ball screws page C76 (the same rows as the precision ground series, page C48).
NUTS = {
    "SFU01204-4": (12, 4, 2.5, 24, 40, 10, 40, 32, 30, 4.5),
    "SFNU01605-4": (16, 5, 3.175, 28, 48, 10, 45, 38, 40, 5.5),
    "SFNU01610-3": (16, 10, 3.175, 28, 48, 10, 57, 38, 40, 5.5),
    "SFNU02005-4": (20, 5, 3.175, 36, 58, 10, 51, 47, 44, 6.6),
    "SFNU02505-4": (25, 5, 3.175, 40, 62, 10, 51, 51, 48, 6.6),
}
# Bolt holes, page C48/C76 drawing for d <= 32: six X holes on the bolt circle W, at 45 degree steps either side of the axis between the flats
# (90 +- 45 and 270 +- 45 degrees, none on the flats' side).
HOLE_ANGLES = (45, 90, 135, 225, 270, 315)
# The ball track, proportions (the catalogue says the groove is a Gothic arch, page C02 Fig 1.1.2, but doesn't dimension it): one circular groove
# of radius TRACK x the ball's radius, its centre on the outside diameter d, so it is that deep; drawn with ARC_STEPS straight steps a side.
TRACK = 1.04
ARC_STEPS = 6
MAX_TURNS = 150  # grooved all the way up to this length / lead, else END_TURNS at each end and a plain middle
END_TURNS = 40


def track(d, lead, ball, run, low=True, high=True):
    """The grooved shaft, one solid sewn from the groove's own helical faces and trimmed flat at both ends (as lead-screw's thread)."""
    top, rg = d / 2, TRACK * ball / 2
    root = top - rg
    up = [(top - rg * math.cos(math.radians(a)), rg * math.sin(math.radians(a))) for a in [90 * i / ARC_STEPS for i in range(ARC_STEPS + 1)]]
    profile = up + [(r, lead - z) for r, z in reversed(up)]   # groove bottom, up its side, the land, down the other side, next bottom

    def helix(r, z):
        return cq.Wire.makeHelix(lead, lead, r).moved(cq.Location(cq.Vector(0, 0, z)))

    wires = [helix(r, z) for r, z in profile]
    piece = [cq.Face.makeRuledSurface(a, b) for a, b in zip(wires, wires[1:])]
    levels = math.ceil(run / lead) + 4
    bottom = -2 * lead
    faces = [face.moved(cq.Location(cq.Vector(0, 0, bottom + k * lead))) for k in range(levels) for face in piece]
    for low_z in (bottom, bottom + levels * lead):
        axis = cq.Wire.assembleEdges([cq.Edge.makeLine(cq.Vector(0, 0, low_z), cq.Vector(0, 0, low_z + lead))])
        faces.append(cq.Face.makeRuledSurface(axis, wires[0].moved(cq.Location(cq.Vector(0, 0, low_z)))))
        outline = [cq.Vector(0, 0, low_z)] + [w.startPoint().add(cq.Vector(0, 0, low_z)) for w in wires] + [cq.Vector(0, 0, low_z + lead)]
        faces.append(cq.Face.makeFromWires(cq.Wire.makePolygon(outline + [outline[0]])))
    body = cq.Workplane(obj=cq.Solid.makeSolid(cq.Shell.makeShell(faces)))
    clear, spare = d + 1, 3 * lead + 1
    body = body.cut(cq.Workplane("XY").workplane(offset=run).circle(clear).extrude(spare + levels * lead))
    body = body.cut(cq.Workplane("XY").workplane(offset=-spare).circle(clear).extrude(spare))
    shape = body.val()
    if len(shape.Solids()) != 1 or not shape.isValid() or not math.pi * root * root * run < shape.Volume() < math.pi * top * top * run:
        raise ValueError(f"The ball track could not be made one solid at d={d}, lead={lead}, length={run}")
    return body

if model_no not in NUTS:
    raise ValueError(f"no TBI MOTION model {model_no!r}: {', '.join(NUTS)}")
d, lead, ball, D, A, B, L, W, H, X = NUTS[model_no]

model = Model()
if part == "screw" and detail == "simple":
    # The screw shaft as its outside diameter d (the catalogue entries' plain "simple" shaft).
    with model.sketch("shaft_sketch", on="XY", name="Shaft") as sk:
        sk.circle("shaft_circle", (0, 0), d / 2)
    model.extrude("shaft", "shaft_sketch", distance=length, name="Screw shaft")
    model.material("steel-generic", color="#b9bdc2")
elif part == "screw":
    if length <= MAX_TURNS * lead:
        model.add("shaft", track(d, lead, ball, length), name="Screw shaft")
    else:
        end = END_TURNS * lead
        model.add("shaft", track(d, lead, ball, end), name="Screw shaft")
        model.add("shaft_end", track(d, lead, ball, end).translate((0, 0, length - end)), name="Screw shaft")
        model.plane("middle_plane", on="XY", offset=end, name="Plain middle")
        with model.sketch("middle_sketch", on="middle_plane", name="Plain middle") as sk:
            sk.circle("middle_circle", (0, 0), d / 2)
        model.extrude("middle", "middle_sketch", distance=length - 2 * end, name="Plain middle")
    model.material("steel-generic", color="#b9bdc2")
elif part == "nut":
    with model.sketch("flange_sketch", on="XY", name="Flange") as sk:
        sk.circle("flange_circle", (0, 0), A / 2)
    model.extrude("flange", "flange_sketch", distance=B, name="Flange")
    # the two flats, H apart
    model.box("flat_left", A, A, B, at=(-H / 2 - A / 2, 0, B / 2), op="cut", name="Flat")
    model.box("flat_right", A, A, B, at=(H / 2 + A / 2, 0, B / 2), op="cut", name="Flat")
    model.plane("body_plane", on="XY", offset=B, name="Body")
    with model.sketch("body_sketch", on="body_plane", name="Body") as sk:
        sk.circle("body_circle", (0, 0), D / 2)
    model.extrude("body", "body_sketch", distance=L - B, op="join", name="Nut body")
    with model.sketch("holes_sketch", on="XY", name="Bolt holes") as sk:
        for i, a in enumerate(HOLE_ANGLES):
            sk.circle(f"hole_{i}", (W / 2 * math.cos(math.radians(a)), W / 2 * math.sin(math.radians(a))), X / 2)
    model.extrude("holes", "holes_sketch", distance=B, op="cut", name="Bolt holes")
    # The bore: the screw's outside diameter d, a proportion (the catalogue doesn't dimension the nut's inside or its ball track).
    with model.sketch("bore_sketch", on="XY", name="Bore") as sk:
        sk.circle("bore_circle", (0, 0), d / 2)
    model.extrude("bore", "bore_sketch", distance=L, op="cut", name="Bore")
    model.material("steel-generic", color="#8d9298")
else:
    raise ValueError("part is screw or nut")
result = model.result()