import cadquery as Cq
import math

def hirth_joint(radius=60,
                radius_inner=40,
                base_height=20,
                n_tooth=16,
                tooth_height=16,
                tooth_height_inner=2,
                tol=0.01):
    """
    Creates a cylindrical Hirth Joint
    """
    # ensures secant doesn't blow up
    assert n_tooth >= 5

    # angle of half of a single tooth
    theta = math.pi / n_tooth

    # Generate a tooth by lofting between two curves

    inner_raise = (tooth_height - tooth_height_inner) / 2
    # Outer tooth triangle spans a curve of length `2 pi r / n_tooth`. This
    # creates the side profile (looking radially inwards) of each of the
    # triangles.
    outer = [
        (radius * math.tan(theta), 0),
        (radius * math.tan(theta) - tol, 0),
        (0, tooth_height),
        (-radius * math.tan(theta) + tol, 0),
        (-radius * math.tan(theta), 0),
    ]
    inner = [
        (radius_inner * math.sin(theta), 0),
        (radius_inner * math.sin(theta), inner_raise),
        (0, inner_raise + tooth_height_inner),
        (-radius_inner * math.sin(theta), inner_raise),
        (-radius_inner * math.sin(theta), 0),
    ]
    tooth = (
        Cq.Workplane('YZ')
        .polyline(inner)
        .close()
        .workplane(offset=radius - radius_inner)
        .polyline(outer)
        .close()
        .loft(ruled=True, combine=True)
        .val()
    )
    tooth_centre_radius = radius_inner * math.cos(theta)
    teeth = (
        Cq.Workplane('XY')
        .polarArray(radius=tooth_centre_radius, startAngle=0, angle=360, count=n_tooth)
        .eachpoint(lambda loc: tooth.located(loc))
        .intersect(Cq.Solid.makeCylinder(
            height=base_height + tooth_height,
            radius=radius,
        ))
    )
    base = (
        Cq.Workplane('XY')
        .cylinder(
            height=base_height,
            radius=radius,
            centered=(True, True, False))
        .faces(">Z").tag("bore")
        .union(teeth.val().move(Cq.Location((0,0,base_height))))
        .clean()
    )
    #base.workplane(offset=tooth_height/2).circle(radius=radius,forConstruction=True).tag("mate")
    base.polyline([(0, 0, base_height), (0, 0, base_height+tooth_height)], forConstruction=True).tag("mate")
    return base

def hirth_assembly():
    """
    Example assembly of two Hirth joints
    """
    rotate = 180 / 16
    obj1 = hirth_joint().faces(tag="bore").cboreHole(
        diameter=10,
        cboreDiameter=20,
        cboreDepth=3)
    obj2 = (
        hirth_joint()
        .rotate(
            axisStartPoint=(0,0,0),
            axisEndPoint=(0,0,1),
            angleDegrees=rotate
        )
    )
    result = (
        Cq.Assembly()
        .add(obj1, name="obj1", color=Cq.Color(0.8,0.8,0.5,0.3))
        .add(obj2, name="obj2", color=Cq.Color(0.5,0.5,0.5,0.3))
        .constrain("obj1?mate", "obj2?mate", "Plane")
        .solve()
    )
    return result

def comma_joint(radius=30,
                shaft_radius=10,
                height=10,
                flange=10,
                flange_thickness=25,
                n_serration=16,
                serration_angle_offset=0,
                serration_height=5,
                serration_inner_radius=20,
                serration_theta=2 * math.pi / 48,
                serration_tilt=-30,
                right_handed=False):
    """
    Produces a "o_" shaped joint, with serrations to accomodate a torsion spring
    """
    assert flange_thickness <= radius
    flange_poly = [
        (0, radius - flange_thickness),
        (0, radius),
        (flange + radius, radius),
        (flange + radius, radius - flange_thickness)
    ]
    if right_handed:
        flange_poly = [(x, -y) for x,y in flange_poly]
    sketch = (
        Cq.Sketch()
        .circle(radius)
        .polygon(flange_poly, mode='a')
        .circle(shaft_radius, mode='s')
    )
    serration_poly = [
        (0, 0), (radius, 0),
        (radius, radius * math.tan(serration_theta))
    ]
    serration = (
        Cq.Workplane('XY')
        .sketch()
        .polygon(serration_poly)
        .circle(radius, mode='i')
        .circle(serration_inner_radius, mode='s')
        .finalize()
        .extrude(serration_height)
        .translate(Cq.Vector((-serration_inner_radius, 0, height)))
        .rotate(
            axisStartPoint=(0, 0, 0),
            axisEndPoint=(0, 0, height),
            angleDegrees=serration_tilt)
        .val()
    )
    serrations = (
        Cq.Workplane('XY')
        .polarArray(radius=serration_inner_radius,
                    startAngle=0+serration_angle_offset,
                    angle=360+serration_angle_offset,
                    count=n_serration)
        .eachpoint(lambda loc: serration.located(loc))
    )
    result = (
        Cq.Workplane()
        .add(sketch)
        .extrude(height)
        .union(serrations)
        .clean()
    )

    result.polyline([
        (0, 0, height - serration_height),
        (0, 0, height + serration_height)],
        forConstruction=True).tag("serrated")
    result.polyline([
        (0, radius, 0),
        (flange + radius, radius, 0)],
        forConstruction=True).tag("tail")
    result.faces('>X').tag("tail_end")
    return result

def torsion_spring(radius=12,
                   height=20,
                   thickness=2,
                   omega=90,
                   tail_length=25):
    """
    Produces a torsion spring with abridged geometry since sweep is very slow in
    cq-editor.
    """
    base = (
        Cq.Workplane('XY')
        .cylinder(height=height, radius=radius,
                  centered=(True, True, False))
    )
    base.faces(">Z").tag("mate_top")
    base.faces("<Z").tag("mate_bottom")
    result = (
        base
        .cylinder(height=height, radius=radius - thickness, combine='s',
                  centered=(True, True, True))
        .transformed(
            offset=(0, radius-thickness),
            rotate=(0, 0, 0))
        .box(length=tail_length, width=thickness, height=thickness, centered=False)
        .copyWorkplane(Cq.Workplane('XY'))
        .transformed(
            offset=(0, 0, height - thickness),
            rotate=(0, 0, omega))
        .center(-tail_length, radius-thickness)
        .box(length=tail_length, width=thickness, height=thickness, centered=False)
    )

    return result

def comma_assembly():
    joint1 = comma_joint()
    joint2 = comma_joint()
    spring = torsion_spring()
    result = (
        Cq.Assembly()
        .add(joint1, name="joint1", color=Cq.Color(0.8,0.8,0.5,0.3))
        .add(joint2, name="joint2", color=Cq.Color(0.8,0.8,0.5,0.3))
        .add(spring, name="spring", color=Cq.Color(0.5,0.5,0.5,1))
        .constrain("joint1?serrated", "spring?mate_bottom", "Plane")
        .constrain("joint2?serrated", "spring?mate_top", "Plane")
        .constrain("joint1?tail", "FixedAxis", (1, 0, 0))
        .constrain("joint2?tail", "FixedAxis", (-1, 0, 0))
        .solve()
    )
    return result