from dataclasses import dataclass
import math
import cadquery as Cq
import nhf.springs as NS
from nhf import Role

def hirth_tooth_angle(n_tooth):
    """
    Angle of one whole tooth
    """
    return 360 / n_tooth

def hirth_joint(radius=60,
                radius_inner=40,
                base_height=20,
                n_tooth=16,
                tooth_height=16,
                tooth_height_inner=2,
                tol=0.01,
                tag_prefix="",
                is_mated=False):
    """
    Creates a cylindrical Hirth Joint

    is_mated: If set to true, rotate the teeth so they line up at 0 degrees.

    FIXME: The curves don't mate perfectly. See if non-planar lofts can solve
    this issue.
    """
    # ensures tangent doesn't blow up
    assert n_tooth >= 5
    assert radius > radius_inner
    assert tooth_height >= tooth_height_inner

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

    c, s, t = math.cos(theta), math.sin(theta), math.tan(theta)
    span = radius * t
    radius_proj = radius / c
    span_inner = radius_inner * s
    # 2 * raise + (inner tooth height) = (tooth height)
    inner_raise = (tooth_height - tooth_height_inner) / 2
    # Outer tooth triangle spans 2*theta radians. This profile is the radial
    # profile projected onto a plane `radius` away from the centre of the
    # cylinder. The y coordinates on the edge must drop to compensate.

    # The drop is equal to, via similar triangles
    drop = inner_raise * (radius_proj - radius) / (radius - radius_inner)
    outer = [
        (span, -tol - drop),
        (span, -drop),
        (0, tooth_height),
        (-span, -drop),
        (-span, -tol - drop),
    ]
    adj = radius_inner * c
    # In the case of the inner triangle, it is projected onto a plane `adj` away
    # from the centre. The apex must extrapolate

    # Via similar triangles
    #
    #  (inner_raise + tooth_height_inner) -
    #    (tooth_height - inner_raise - tooth_height_inner) * ((radius_inner - adj) / (radius - radius_inner))
    apex = (inner_raise + tooth_height_inner)  - \
        inner_raise * (radius_inner - adj) / (radius - radius_inner)
    inner = [
        (span_inner, -tol),
        (span_inner, inner_raise),
        (0, apex),
        (-span_inner, inner_raise),
        (-span_inner, -tol),
    ]
    tooth = (
        Cq.Workplane('YZ')
        .polyline(inner)
        .close()
        .workplane(offset=radius - adj)
        .polyline(outer)
        .close()
        .loft(ruled=False, combine=True)
        .val()
    )
    angle_offset = hirth_tooth_angle(n_tooth) / 2 if is_mated else 0
    teeth = (
        Cq.Workplane('XY')
        .polarArray(
            radius=adj,
            startAngle=angle_offset,
            angle=360,
            count=n_tooth)
        .eachpoint(lambda loc: tooth.located(loc))
        .intersect(Cq.Solid.makeCylinder(
            height=base_height + tooth_height,
            radius=radius,
        ))
        .intersect(Cq.Solid.makeCylinder(
            height=base_height + tooth_height,
            radius=radius,
        ))
        .cut(Cq.Solid.makeCylinder(
            height=base_height + tooth_height,
            radius=radius_inner,
        ))
    )
    base = (
        Cq.Workplane('XY')
        .cylinder(
            height=base_height,
            radius=radius,
            centered=(True, True, False))
        .faces(">Z").tag(f"{tag_prefix}bore")
        .union(teeth.val().move(Cq.Location((0,0,base_height))), tol=tol)
        .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(f"{tag_prefix}mate")
    )
    (
        base
        .polyline([(0, 0, 0), (1, 0, 0)], forConstruction=True)
        .tag(f"{tag_prefix}directrix")
    )
    return base

def hirth_assembly(n_tooth=12):
    """
    Example assembly of two Hirth joints
    """
    #rotate = 180 / 16

    tab = (
        Cq.Workplane('XY')
        .box(100, 10, 2, centered=False)
    )
    obj1 = (
        hirth_joint(n_tooth=n_tooth)
        .faces(tag="bore")
        .cboreHole(
            diameter=10,
            cboreDiameter=20,
            cboreDepth=3)
        .union(tab)
    )
    obj2 = (
        hirth_joint(n_tooth=n_tooth, is_mated=True)
        .union(tab)
    )
    angle = hirth_tooth_angle(n_tooth)
    result = (
        Cq.Assembly()
        .add(obj1, name="obj1", color=Role.PARENT.color)
        .add(obj2, name="obj2", color=Role.CHILD.color)
        .constrain("obj1", "Fixed")
        .constrain("obj1?mate", "obj2?mate", "Plane")
        .constrain("obj1?directrix", "obj2?directrix", "Axis", param=angle)
        .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 comma_assembly():
    joint1 = comma_joint()
    joint2 = comma_joint()
    spring = NS.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?bot", "Plane")
        .constrain("joint2?serrated", "spring?top", "Plane")
        .constrain("joint1?tail", "FixedAxis", (1, 0, 0))
        .constrain("joint2?tail", "FixedAxis", (-1, 0, 0))
        .solve()
    )
    return result

@dataclass(frozen=True)
class TorsionJoint:
    """
    This jonit consists of a rider puck on a track puck. IT is best suited if
    the radius has to be small and vertical space is abundant.
    """

    # Radius limit for rotating components
    radius = 40
    disk_height = 10

    radius_spring = 15
    radius_axle = 10

    # Offset of the spring hole w.r.t. surface
    spring_hole_depth = 4

    # Also used for the height of the hole for the spring
    spring_thickness = 2
    spring_height = 15

    spring_tail_length = 40

    groove_radius_outer = 35
    groove_radius_inner = 20
    groove_depth = 5
    rider_gap = 2
    n_slots = 8

    right_handed: bool = False

    def __post_init__(self):
        assert self.disk_height > self.spring_hole_depth
        assert self.radius > self.groove_radius_outer
        assert self.groove_radius_outer > self.groove_radius_inner
        assert self.groove_radius_inner > self.radius_spring
        assert self.spring_height > self.groove_depth, "Groove is too deep"

    @property
    def total_height(self):
        return 2 * self.disk_height + self.spring_height

    @property
    def _radius_spring_internal(self):
        return self.radius_spring - self.spring_thickness

    def _slot_polygon(self, flip: bool=False):
        r1 = self.radius_spring - self.spring_thickness
        r2 = self.radius_spring
        flip = flip != self.right_handed
        if flip:
            r1 = -r1
            r2 = -r2
        return [
            (0, r2),
            (self.spring_tail_length, r2),
            (self.spring_tail_length, r1),
            (0, r1),
        ]
    def _directrix(self, height, theta=0):
        c, s = math.cos(theta), math.sin(theta)
        r2 = self.radius_spring
        l = self.spring_tail_length
        if self.right_handed:
            r2 = -r2
        # This is (0, r2) and (l, r2) transformed by rotation matrix
        # [[c, s], [-s, c]]
        return [
            (s * r2,         -s * l + c * r2, height),
            (c * l + s * r2, -s * l + c * r2, height),
        ]


    def spring(self):
        return NS.torsion_spring(
            radius=self.radius_spring,
            height=self.spring_height,
            thickness=self.spring_thickness,
            tail_length=self.spring_tail_length,
        )

    def track(self):
        groove_profile = (
            Cq.Sketch()
            .circle(self.radius)
            .circle(self.groove_radius_outer, mode='s')
            .circle(self.groove_radius_inner, mode='a')
            .circle(self.radius_spring, mode='s')
        )
        spring_hole_profile = (
            Cq.Sketch()
            .circle(self.radius)
            .polygon(self._slot_polygon(flip=False), mode='s')
            .circle(self.radius_spring, mode='s')
        )
        result = (
            Cq.Workplane('XY')
            .cylinder(radius=self.radius, height=self.disk_height)
            .faces('>Z')
            .tag("spring")
            .placeSketch(spring_hole_profile)
            .extrude(self.spring_thickness)
            # If the spring hole profile is not simply connected, this workplane
            # will have to be created from the `spring-mate` face.
            .faces('>Z')
            .placeSketch(groove_profile)
            .extrude(self.groove_depth)
            .faces('>Z')
            .hole(self.radius_axle)
        )
        # Insert directrix`
        result.polyline(self._directrix(self.disk_height),
                        forConstruction=True).tag("directrix")
        return result

    def rider(self):
        def slot(loc):
            wire = Cq.Wire.makePolygon(self._slot_polygon(flip=False))
            face = Cq.Face.makeFromWires(wire)
            return face.located(loc)
        wall_profile = (
            Cq.Sketch()
            .circle(self.radius, mode='a')
            .circle(self.radius_spring, mode='s')
            .parray(
                r=0,
                a1=0,
                da=360,
                n=self.n_slots)
            .each(slot, mode='s')
            #.circle(self._radius_wall, mode='a')
        )
        contact_profile = (
            Cq.Sketch()
            .circle(self.groove_radius_outer, mode='a')
            .circle(self.groove_radius_inner, mode='s')
            #.circle(self._radius_wall, mode='a')
            .parray(
                r=0,
                a1=0,
                da=360,
                n=self.n_slots)
            .each(slot, mode='s')
        )
        middle_height = self.spring_height - self.groove_depth - self.rider_gap
        result = (
            Cq.Workplane('XY')
            .cylinder(radius=self.radius, height=self.disk_height)
            .faces('>Z')
            .tag("spring")
            .placeSketch(wall_profile)
            .extrude(middle_height)
            # The top face might not be in one piece.
            #.faces('>Z')
            .workplane(offset=middle_height)
            .placeSketch(contact_profile)
            .extrude(self.groove_depth + self.rider_gap)
            .faces(tag="spring")
            .circle(self._radius_spring_internal)
            .extrude(self.spring_height)
            .faces('>Z')
            .hole(self.radius_axle)
        )
        for i in range(self.n_slots):
            theta = 2 * math.pi * i / self.n_slots
            result.polyline(self._directrix(self.disk_height, theta),
                            forConstruction=True).tag(f"directrix{i}")
        return result

    def rider_track_assembly(self):
        rider = self.rider()
        track = self.track()
        spring = self.spring()
        result = (
            Cq.Assembly()
            .add(spring, name="spring", color=Cq.Color(0.5,0.5,0.5,1))
            .add(track, name="track", color=Cq.Color(0.5,0.5,0.8,0.3))
            .constrain("track?spring", "spring?top", "Plane")
            .add(rider, name="rider", color=Cq.Color(0.8,0.8,0.5,0.3))
            .constrain("rider?spring", "spring?bot", "Plane")
            .constrain("track?directrix", "spring?directrix_bot", "Axis")
            .constrain("rider?directrix0", "spring?directrix_top", "Axis")
            .solve()
        )
        return result