Cosplay/nhf/touhou/houjuu_nue/joints.py

import math
from dataclasses import dataclass, field
from typing import Optional, Tuple
import cadquery as Cq
from nhf import Material, Role
from nhf.build import Model, target, assembly
from nhf.parts.springs import TorsionSpring
from nhf.parts.fasteners import FlatHeadBolt, HexNut, ThreaddedKnob
from nhf.parts.joints import TorsionJoint, HirthJoint
from nhf.parts.box import Hole, MountingBox, box_with_centre_holes
import nhf.utils

TOL = 1e-6

# Parts used

# uxcell 2 Pcs Star Knobs Grips M12 x 30mm Male Thread Steel Zinc Stud Replacement PP
HS_JOINT_KNOB = ThreaddedKnob(
    mass=float('nan'),
    diam_thread=12.0,
    height_thread=30.0,
    diam_knob=50.0,
    # FIXME: Undetermined
    diam_neck=30.0,
    height_neck=10.0,
    height_knob=10.0,
)
# Tom's world 8Pcs M12-1.75 Hex Nut Assortment Set Stainless Steel 304(18-8)
# Metric Hexagon Nut for Bolts, Bright Finish, Full Thread (M12)
HS_JOINT_HEX_NUT = HexNut(
    mass=14.9,
    diam_thread=12.0,
    pitch=1.75,
    thickness=9.7,
    width=18.9,
)
SHOULDER_AXIS_BOLT = FlatHeadBolt(
    # FIXME: measure
    diam_head=10.0,
    height_head=3.0,
    diam_thread=6.0,
    height_thread=20.0,
    mass=float('nan'),
)
# Hoypeyfiy 10 Pieces Torsion Spring Woodworking DIY 90 Degrees Torsional
# Springs Repair Maintenance Spring
SHOULDER_TORSION_SPRING = TorsionSpring(
    mass=2.2,
    # inner diameter = 9
    radius=9/2 + 1.2,
    thickness=1.3,
    height=7.5,
)
# KALIONE 10 Pieces Torsion Spring, Stainless Steel Small Torsion Springs, Tiny
# Torsional Spring, 90° Deflection Compression Spring Kit for Repair Tools
# Woodworking DIY, 50mm
ELBOW_TORSION_SPRING = TorsionSpring(
    mass=1.7,
    radius=9 / 2,
    thickness=1.3,
    height=6.5,
    tail_length=45.0,
    right_handed=False,
)

@dataclass
class RootJoint(Model):
    """
    The Houjuu-Scarlett Mechanism
    """
    knob: ThreaddedKnob = HS_JOINT_KNOB
    hex_nut: HexNut = HS_JOINT_HEX_NUT
    hirth_joint: HirthJoint = field(default_factory=lambda: HirthJoint(
        radius=25.0,
        radius_inner=15.0,
        tooth_height=7.0,
        base_height=5.0,
        n_tooth=24,
    ))
    parent_width: float = 85
    parent_thickness: float = 10
    parent_corner_fillet: float = 5
    parent_corner_cbore_diam: float = 12
    parent_corner_cbore_depth: float = 2
    parent_corner_inset: float = 12
    parent_mount_thickness: float = 25.4 / 16

    child_corner_dx: float = 17.0
    child_corner_dz: float = 24.0

    axis_diam: float = 12.0
    axis_cbore_diam: float = 20
    axis_cbore_depth: float = 3
    corner_hole_diam: float = 6.0

    child_height: float = 60.0
    child_width: float = 50.0
    child_mount_thickness: float = 25.4 / 8

    def corner_pos(self) -> list[tuple[int, int]]:
        """
        Generates a set of points corresponding to the connectorss
        """
        dx = self.parent_width / 2 - self.parent_corner_inset
        return [
            (dx, dx),
            (dx, -dx),
            (-dx, -dx),
            (-dx, dx),
        ]

    @property
    def total_height(self) -> float:
        return self.parent_thickness + self.hirth_joint.total_height

    @target(name="parent")
    def parent(self):
        """
        Parent part of the Houjuu-Scarlett joint, which is composed of a Hirth
        coupling, a cylindrical base, and a mounting base.
        """
        hirth = self.hirth_joint.generate()
        conn = self.corner_pos()
        result = (
            Cq.Workplane('XY')
            .box(
                self.parent_width,
                self.parent_width,
                self.parent_thickness,
                centered=(True, True, False))
            .translate((0, 0, -self.parent_thickness))
            .edges("|Z")
            .fillet(self.parent_corner_fillet)
            .faces(">Z")
            .workplane()
            .pushPoints(conn)
            .cboreHole(
                diameter=self.corner_hole_diam,
                cboreDiameter=self.parent_corner_cbore_diam,
                cboreDepth=self.parent_corner_cbore_depth)
        )
        # Creates a plane parallel to the holes but shifted to the base
        plane = result.faces(">Z").workplane(offset=-self.parent_thickness)

        for i, (px, py) in enumerate(conn):
            plane.moveTo(px, py).tagPoint(f"h{i}")
        result = (
            result
            .faces(">Z")
            .workplane()
            .union(hirth, tol=0.1)
            .clean()
        )
        result = (
            result.faces("<Z")
            .workplane()
            .hole(diameter=self.axis_diam)
            .cut(self.hex_nut.generate().translate((0, 0, -self.parent_thickness)))
        )
        result.faces("<Z").tag("base")
        return result

    @target(name="child")
    def child(self) -> Cq.Workplane:
        hirth = self.hirth_joint.generate(is_mated=True)
        dy = self.child_corner_dx
        dx = self.child_corner_dz
        conn = [
            (-dx, -dy),
            (dx, -dy),
            (dx, dy),
            (-dx, dy),
        ]
        result = (
            Cq.Workplane('XY')
            .box(
                self.child_height,
                self.child_width,
                self.hirth_joint.base_height,
                centered=(True, True, False))
            #.translate((0, 0, -self.base_joint.base_height))
            #.edges("|Z")
            #.fillet(self.hs_joint_corner_fillet)
            .faces(">Z")
            .workplane()
            .pushPoints(conn)
            .hole(self.corner_hole_diam)
        )
        # Creates a plane parallel to the holes but shifted to the base
        plane = result.faces(">Z").workplane(offset=-self.hirth_joint.base_height)

        for i, (px, py) in enumerate(conn):
            plane.moveTo(px, py).tagPlane(f"conn{i}")
        result = (
            result
            .faces(">Z")
            .workplane()
            .union(hirth, tol=0.1)
            .clean()
        )
        result = (
            result.faces("<Z")
            .workplane()
            .hole(self.axis_diam)
        )
        return result

    def assembly(self,
                 offset: int = 0,
                 fastener_pos: float = 0) -> Cq.Assembly:
        """
        Specify knob position to determine the position of the knob from fully
        inserted (0) or fully uninserted (1)
        """
        knob_h = self.hex_nut.thickness
        result = (
            Cq.Assembly()
            .addS(self.parent(), name="parent",
                  material=Material.PLASTIC_PLA,
                  role=Role.PARENT)
            .constrain("parent", "Fixed")
            .addS(self.child(), name="child",
                  material=Material.PLASTIC_PLA,
                  role=Role.CHILD)
            .addS(self.hex_nut.assembly(), name="hex_nut")
            .addS(self.knob.assembly(), name="knob",
                  loc=Cq.Location((0, 0, knob_h * fastener_pos)))
            .constrain("knob/thread", "Fixed")
            .constrain("hex_nut?bot", "parent?base", "Plane", param=0)
            .constrain("hex_nut?dirX", "parent@faces@>X", "Axis", param=0)
        )
        self.hirth_joint.add_constraints(
            result,
            "parent",
            "child",
            offset=offset
        )
        return result.solve()

@dataclass
class ShoulderJoint(Model):

    bolt: FlatHeadBolt = SHOULDER_AXIS_BOLT

    height: float = 70.0
    torsion_joint: TorsionJoint = field(default_factory=lambda: TorsionJoint(
        radius_track=18,
        radius_rider=18,
        groove_depth=4.8,
        groove_radius_outer=16,
        groove_radius_inner=13,
        track_disk_height=5.0,
        rider_disk_height=5.0,
        radius_axle=3.0,
        spring=SHOULDER_TORSION_SPRING,
        rider_slot_begin=0,
        rider_n_slots=1,
        rider_slot_span=0,
    ))

    # On the parent side, drill vertical holes

    parent_conn_hole_diam: float = 6.0
    # Position of the holes relative
    parent_conn_hole_pos: list[Tuple[float, float]] = field(default_factory=lambda: [
        (15, 8),
        (15, -8),
    ])

    parent_lip_length: float = 25.0
    parent_lip_width: float = 30.0
    parent_lip_thickness: float = 5.0
    parent_lip_ext: float = 40.0
    parent_lip_guard_height: float = 8.0

    # Generates a child guard which covers up the internals. The lip length is
    # relative to the +X surface of the guard.
    child_guard_ext: float = 30.0
    child_guard_width: float = 25.0
    # guard length measured from axle
    child_lip_length: float = 40.0
    child_lip_width: float = 20.0
    child_conn_hole_diam: float = 6.0
    # Measured from centre of axle
    child_conn_hole_pos: list[float] = field(default_factory=lambda: [15, 25])
    child_core_thickness: float = 3.0


    # Rotates the torsion joint to avoid collisions or for some other purpose
    axis_rotate_bot: float = 225.0
    axis_rotate_top: float = -225.0

    directrix_id: int = 0
    angle_neutral: float = 10.0
    angle_max_deflection: float = 80.0

    def __post_init__(self):
        assert self.parent_lip_length * 2 < self.height

    @property
    def radius(self):
        return self.torsion_joint.radius

    def parent_arm_loc(self) -> Cq.Location:
        """
        2d location of the arm surface on the parent side, relative to axle
        """
        return Cq.Location.rot2d(self.angle_neutral) * Cq.Location.from2d(self.parent_lip_ext, 0, 0)

    def parent(self, top: bool = False) -> Cq.Assembly:
        joint = self.torsion_joint
        # Thickness of the lip connecting this joint to the wing root
        assert self.parent_lip_width <= joint.radius_track * 2
        assert self.parent_lip_ext > joint.radius_track

        lip_guard = (
            Cq.Solid.makeBox(
                self.parent_lip_ext,
                self.parent_lip_width,
                self.parent_lip_guard_height)
            .located(Cq.Location((0, -self.parent_lip_width/2 , 0)))
            .cut(Cq.Solid.makeCylinder(joint.radius_track, self.parent_lip_guard_height))
        )
        lip = MountingBox(
            length=self.parent_lip_length,
            width=self.parent_lip_width,
            thickness=self.parent_lip_thickness,
            holes=[
                Hole(x=self.height / 2 - x, y=y)
                for x, y in self.parent_conn_hole_pos
            ],
            hole_diam=self.parent_conn_hole_diam,
            generate_side_tags=False,
        )
        # Flip so the lip's holes point to -X
        loc_axis = Cq.Location((0,0,0), (0, 1, 0), -90)
        # so they point to +X
        loc_dir = Cq.Location((0,0,0), (0, 0, 1), 180)
        loc_pos = Cq.Location((self.parent_lip_ext - self.parent_lip_thickness, 0, 0))

        rot = -self.axis_rotate_top if top else self.axis_rotate_bot

        result = (
            Cq.Assembly()
            .add(joint.track(), name="track",
                 loc=Cq.Location((0, 0, 0), (0, 0, 1), rot))
            .add(lip_guard, name="lip_guard")
            .add(lip.generate(), name="lip", loc=loc_pos * loc_dir * loc_axis)
        )
        return result

    @target(name="parent-bot")
    def parent_bot(self) -> Cq.Assembly:
        return self.parent(top=False)
    @target(name="parent-top")
    def parent_top(self) -> Cq.Assembly:
        return self.parent(top=True)

    @property
    def child_height(self) -> float:
        """
        Calculates the y distance between two joint surfaces on the child side
        of the shoulder joint.
        """
        joint = self.torsion_joint
        return self.height - 2 * joint.total_height + 2 * joint.rider_disk_height

    @target(name="child")
    def child(self) -> Cq.Assembly:
        """
        Creates the top/bottom shoulder child joint
        """

        joint = self.torsion_joint
        assert all(r < self.child_lip_length for r in self.child_conn_hole_pos)

        # Half of the height of the bridging cylinder
        dh = self.height / 2 - joint.total_height
        core_start_angle = 30
        core_end_angle1 = 90
        core_end_angle2 = 180

        radius_core_inner = joint.radius_rider - self.child_core_thickness
        core_profile1 = (
            Cq.Sketch()
            .arc((0, 0), joint.radius_rider, core_start_angle, core_end_angle1-core_start_angle)
            .segment((0, 0))
            .close()
            .assemble()
            .circle(radius_core_inner, mode='s')
        )
        core_profile2 = (
            Cq.Sketch()
            .arc((0, 0), joint.radius_rider, -core_start_angle, -(core_end_angle2-core_start_angle))
            .segment((0, 0))
            .close()
            .assemble()
            .circle(radius_core_inner, mode='s')
        )
        core_guard = (
            Cq.Workplane('XY')
            .box(
                length=self.child_guard_ext,
                width=self.child_guard_width,
                height=self.height,
                centered=(False, True, True),
            )
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                radius=self.radius,
                height=self.height,
                combine='cut',
                centered=True,
            )
            .copyWorkplane(Cq.Workplane('XY'))
            .box(
                length=self.child_guard_ext,
                width=self.child_lip_width,
                height=self.height - self.torsion_joint.total_height * 2,
                combine='cut',
                centered=(False, True, True),
            )
        )
        core = (
            Cq.Workplane('XY')
            .placeSketch(core_profile1)
            .toPending()
            .extrude(dh * 2)
            .copyWorkplane(Cq.Workplane('XY'))
            .placeSketch(core_profile2)
            .toPending()
            .extrude(dh * 2)
            .translate(Cq.Vector(0, 0, -dh))
            .union(core_guard)
        )
        assert self.child_lip_width / 2 <= joint.radius_rider
        lip_thickness = joint.rider_disk_height
        lip = box_with_centre_holes(
            length=self.child_lip_length,
            width=self.child_lip_width,
            height=lip_thickness,
            hole_loc=self.child_conn_hole_pos,
            hole_diam=self.child_conn_hole_diam,
        )
        theta = self.torsion_joint.spring.angle_neutral - self.torsion_joint.rider_slot_span
        loc_rotate = Cq.Location((0, 0, 0), (1, 0, 0), 180)
        loc_axis_rotate_bot = Cq.Location((0, 0, 0), (0, 0, 1), self.axis_rotate_bot + self.angle_neutral)
        loc_axis_rotate_top = Cq.Location((0, 0, 0), (0, 0, 1), self.axis_rotate_top + self.angle_neutral)
        result = (
            Cq.Assembly()
            .add(core, name="core", loc=Cq.Location())
            .add(joint.rider(rider_slot_begin=-90, reverse_directrix_label=True), name="rider_top",
                 loc=loc_axis_rotate_top * Cq.Location((0, 0, dh), (0, 0, 1), -90) * Cq.Location((0, 0, 0), (0, 0, 1), theta))
            .add(joint.rider(rider_slot_begin=180), name="rider_bot",
                 loc=loc_axis_rotate_bot * Cq.Location((0, 0, -dh), (0, 0, 1), -90) * loc_rotate)
            .add(lip, name="lip_top",
                 loc=Cq.Location((self.child_guard_ext, 0, dh)))
            .add(lip, name="lip_bot",
                 loc=Cq.Location((self.child_guard_ext, 0, -dh)) * loc_rotate)
        )
        return result

    @assembly()
    def assembly(self, fastener_pos: float = 0.0, deflection: float = 0) -> Cq.Assembly:
        assert deflection <= self.angle_max_deflection
        directrix = self.directrix_id
        mat = Material.RESIN_TRANSPERENT
        mat_spring = Material.STEEL_SPRING
        bolt_z = self.height / 2 + self.bolt.height_thread * (fastener_pos - 1)
        result = (
            Cq.Assembly()
            .addS(self.child(), name="child",
                  role=Role.CHILD, material=mat)
            .constrain("child/core", "Fixed")
            .addS(self.torsion_joint.spring.assembly(deflection=-deflection), name="spring_top",
                  role=Role.DAMPING, material=mat_spring)
            .addS(self.parent_top(),
                  name="parent_top",
                  role=Role.PARENT, material=mat)
            .addS(self.torsion_joint.spring.assembly(deflection=deflection), name="spring_bot",
                  role=Role.DAMPING, material=mat_spring)
            .addS(self.parent_bot(),
                  name="parent_bot",
                  role=Role.PARENT, material=mat)
            # Fasteners
            .addS(self.bolt.assembly(), name="bolt_top",
                  loc=Cq.Location((0, 0, bolt_z)))
            .constrain("bolt_top/thread?root", 'Fixed')
            .addS(self.bolt.assembly(), name="bolt_bot",
                  loc=Cq.Location((0, 0, -bolt_z), (1,0,0), 180))
            .constrain("bolt_bot/thread?root", 'Fixed')
        )
        TorsionJoint.add_constraints(
            result,
            rider="child/rider_top",
            track="parent_top/track",
            spring="spring_top",
            directrix=directrix)
        TorsionJoint.add_constraints(
            result,
            rider="child/rider_bot",
            track="parent_bot/track",
            spring="spring_bot",
            directrix=directrix)
        return result.solve()


@dataclass
class Beam:
    """
    A I-shaped spine with two feet
    """

    foot_length: float = 40.0
    foot_width: float = 20.0
    foot_height: float = 5.0
    spine_thickness: float = 4.0
    spine_length: float = 10.0
    total_height: float = 50.0

    hole_diam: float = 6.0
    # distance between the centres of the two holes
    hole_dist: float = 24.0

    def __post_init__(self):
        assert self.spine_height > 0
        assert self.hole_diam + self.hole_dist < self.foot_length
        assert self.hole_dist - self.hole_diam >= self.spine_length

    @property
    def spine_height(self):
        return self.total_height - self.foot_height * 2

    def foot(self) -> Cq.Workplane:
        """
        just one foot
        """
        dx = self.hole_dist / 2
        result = (
            Cq.Workplane('XZ')
            .box(self.foot_length, self.foot_width, self.foot_height,
                 centered=(True, True, False))
            .faces(">Y")
            .workplane()
            .pushPoints([(dx, 0), (-dx, 0)])
            .hole(self.hole_diam)
        )
        plane = result.faces(">Y").workplane()
        plane.moveTo(dx, 0).tagPlane("conn1")
        plane.moveTo(-dx, 0).tagPlane("conn0")
        return result

    def generate(self, flip: bool = False) -> Cq.Assembly:
        beam = (
            Cq.Workplane('XZ')
            .box(self.spine_length, self.spine_thickness, self.spine_height)
        )
        h = self.spine_height / 2 + self.foot_height

        tag_p, tag_n = "top", "bot"
        if flip:
            tag_p, tag_n = tag_n, tag_p
        result = (
            Cq.Assembly()
            .add(beam, name="beam")
            .add(self.foot(), name=tag_p,
                 loc=Cq.Location((0, h, 0)))
            .add(self.foot(), name=tag_n,
                 loc=Cq.Location((0, -h, 0), (1, 0, 0), 180))
        )
        return result


@dataclass
class DiskJoint(Model):
    """
    Sandwiched disk joint for the wrist and elbow

    We embed a spring inside the joint, with one leg in the disk and one leg in
    the housing. This provides torsion resistance.
    """
    spring: TorsionSpring = ELBOW_TORSION_SPRING

    radius_housing: float = 22.0
    radius_disk: float = 20.0
    radius_axle: float = 3.0

    housing_thickness: float = 4.0
    disk_thickness: float = 7.0

    # Amount by which the wall carves in
    wall_inset: float = 2.0
    # Height of the spring hole; if you make it too short the spring can't enter
    spring_tail_hole_height: float = 2.0

    # Spring angle at 0 degrees of movement
    spring_angle_at_0: float = 90.0
    spring_slot_offset: float = 5.0

    # Angular span of movement
    movement_angle: float = 120.0
    # Angular span of tongue on disk
    tongue_span: float = 30.0
    tongue_length: float = 10.0

    generate_inner_wall: bool = False


    def __post_init__(self):
        super().__init__(name="disk-joint")
        assert self.radius_housing > self.radius_disk > self.radius_axle
        assert self.spring.height < self.housing_thickness + self.disk_thickness

        assert self.housing_upper_carve_offset > 0
        assert self.spring_tail_hole_height > self.spring.thickness

    @property
    def neutral_movement_angle(self) -> Optional[float]:
        a = self.spring.angle_neutral - self.spring_angle_at_0
        if 0 <= a and a <= self.movement_angle:
            return a
        return None

    @property
    def total_thickness(self) -> float:
        return self.housing_thickness * 2 + self.disk_thickness
    @property
    def disk_bot_thickness(self) -> float:
        """
        Pads the bottom of the disk up to spring height
        """
        return max(0, self.disk_thickness + self.spring.thickness - self.spring.height)

    @property
    def opening_span(self) -> float:
        return self.movement_angle + self.tongue_span

    @property
    def housing_upper_carve_offset(self) -> float:
        """
        Distance between the spring track and the outside of the upper housing
        """
        return self.spring_tail_hole_height + (self.disk_thickness - self.disk_bot_thickness) - self.spring.height

    @property
    def housing_upper_dz(self) -> float:
        """
        Distance between the default upper housing location and the median line
        """
        return self.total_thickness / 2 - self.housing_thickness

    def _disk_cut(self) -> Cq.Workplane:
        return (
            Cq.Solid.makeBox(
                length=self.spring.tail_length,
                width=self.spring.thickness,
                height=self.spring.height-self.disk_bot_thickness,
            )
            .located(Cq.Location((0, self.spring.radius_inner, self.disk_bot_thickness)))
            .rotate((0, 0, 0), (0, 0, 1), self.spring_slot_offset)
        )

    @target(name="disk")
    def disk(self) -> Cq.Workplane:
        radius_tongue = self.radius_disk + self.tongue_length
        tongue = (
            Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=radius_tongue,
                angleDegrees=self.tongue_span,
            ).cut(Cq.Solid.makeCylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
            ))
        )
        result = (
            Cq.Workplane('XY')
            .cylinder(
                height=self.disk_thickness,
                radius=self.radius_disk,
                centered=(True, True, False)
            )
            .union(tongue, tol=TOL)
            .copyWorkplane(Cq.Workplane('XY'))
            .cylinder(
                height=self.disk_thickness,
                radius=self.spring.radius,
                centered=(True, True, False),
                combine='cut',
            )
            .cut(self._disk_cut())
        )
        plane = result.copyWorkplane(Cq.Workplane('XY'))
        theta = math.radians(self.spring_slot_offset)
        plane.tagPlane("dir", direction=(math.cos(theta), math.sin(theta), 0))
        plane.workplane(offset=self.disk_thickness).tagPlane("mate_top")
        plane.workplane(offset=self.disk_bot_thickness).tagPlane("mate_spring")
        result.copyWorkplane(Cq.Workplane('YX')).tagPlane("mate_bot")

        return result

    def wall(self) -> Cq.Compound:
        height = self.disk_thickness + self.wall_inset
        wall = Cq.Solid.makeCylinder(
            radius=self.radius_housing,
            height=height,
            angleDegrees=360 - self.opening_span,
        ).cut(Cq.Solid.makeCylinder(
            radius=self.radius_disk,
            height=height,
        )).rotate((0, 0, 0), (0, 0, 1), self.opening_span)
        return wall

    @target(name="housing-lower")
    def housing_lower(self) -> Cq.Workplane:
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
            .cut(Cq.Solid.makeCylinder(
                radius=self.radius_axle,
                height=self.housing_thickness,
            ))
        )
        result.faces(">Z").tag("mate")
        result.faces(">Z").workplane().tagPlane("dirX", direction="+X")
        result = result.cut(
            self
            .wall()
            .located(Cq.Location((0, 0, self.housing_thickness - self.wall_inset)))
            #.rotate((0, 0, 0), (1, 0, 0), 180)
            #.located(Cq.Location((0, 0, self.disk_thickness + self.housing_thickness)))
        )
        return result

    @target(name="housing-upper")
    def housing_upper(self) -> Cq.Workplane:
        carve_angle = -(self.spring_angle_at_0 - self.spring_slot_offset)
        carve = (
            Cq.Solid.makeCylinder(
                radius=self.spring.radius,
                height=self.spring_tail_hole_height,
            ).fuse(Cq.Solid.makeBox(
                length=self.spring.tail_length,
                width=self.spring.thickness,
                height=self.spring_tail_hole_height,
            ).located(Cq.Location((0, -self.spring.radius, 0))))
        ).rotate((0, 0, 0), (0, 0, 1), carve_angle)
        result = (
            Cq.Workplane('XY')
            .cylinder(
                radius=self.radius_housing,
                height=self.housing_thickness,
                centered=(True, True, False),
            )
        )
        theta = math.radians(carve_angle)
        result.faces("<Z").tag("mate")
        p_xy = result.copyWorkplane(Cq.Workplane('XY'))
        p_xy.tagPlane("dirX", direction="+X")
        p_xy.tagPlane("dir", direction=(math.cos(theta), math.sin(theta), 0))
        result = result.faces(">Z").hole(self.radius_axle * 2)

        # tube which holds the spring interior
        if self.generate_inner_wall:
            tube = (
                Cq.Solid.makeCylinder(
                    radius=self.radius_spring_internal,
                    height=self.disk_thickness + self.housing_thickness,
                ).cut(Cq.Solid.makeCylinder(
                    radius=self.radius_axle,
                    height=self.disk_thickness + self.housing_thickness,
                ))
            )
            result = result.union(tube)
        wall = (
            self.wall()
            .located(Cq.Location((0, 0, -self.disk_thickness-self.wall_inset)))
        )
        result = (
            result
            .union(wall, tol=TOL)
            #.cut(carve)
            .cut(carve.located(Cq.Location((0, 0, -self.housing_upper_carve_offset))))
        )
        return result.clean()

    def add_constraints(self,
                        assembly: Cq.Assembly,
                        housing_lower: str,
                        housing_upper: str,
                        disk: str,
                        angle: float = 0.0,
                        ) -> Cq.Assembly:
        assert 0 <= angle <= self.movement_angle
        deflection = angle - self.neutral_movement_angle
        spring_name = disk.replace("/", "__Z") + "_spring"
        (
            assembly
            .addS(
                self.spring.assembly(deflection=-deflection),
                name=spring_name,
                role=Role.DAMPING,
                material=Material.STEEL_SPRING)
            .constrain(f"{disk}?mate_bot", f"{housing_lower}?mate", "Plane")
            .constrain(f"{disk}?mate_top", f"{housing_upper}?mate", "Plane")
            .constrain(f"{housing_lower}?dirX", f"{housing_upper}?dirX", "Axis", param=0)
            .constrain(f"{housing_upper}?dir", f"{spring_name}?dir_top", "Axis", param=0)
            .constrain(f"{spring_name}?dir_bot", f"{disk}?dir", "Axis", param=0)
            .constrain(f"{disk}?mate_spring", f"{spring_name}?bot", "Plane")
            #.constrain(f"{housing_lower}?dirX", f"{housing_upper}?dir", "Axis", param=0)
            #.constrain(f"{housing_lower}?dirX", f"{disk}?dir", "Axis", param=angle)
            #.constrain(f"{housing_lower}?dirY", f"{disk}?dir", "Axis", param=angle - 90)
        )
        return (
            assembly
        )


    def assembly(self, angle: Optional[float] = 0) -> Cq.Assembly:
        if angle is None:
            angle = self.neutral_movement_angle
            if angle is None:
                angle = 0
        else:
            assert 0 <= angle <= self.movement_angle
        result = (
            Cq.Assembly()
            .addS(self.disk(), name="disk", role=Role.CHILD)
            .addS(self.housing_lower(), name="housing_lower", role=Role.PARENT)
            .addS(self.housing_upper(), name="housing_upper", role=Role.CASING)
            .constrain("housing_lower", "Fixed")
        )
        result = self.add_constraints(
            result,
            housing_lower="housing_lower",
            housing_upper="housing_upper",
            disk="disk",
            angle=angle,
        )
        return result.solve()


@dataclass(kw_only=True)
class ElbowJoint(Model):
    """
    Creates the elbow and wrist joints.

    This consists of a disk joint, where each side of the joint has mounting
    holes for connection to the exoskeleton. Each side 2 mounting feet on the
    top and bottom, and each foot has 2 holes.

    On the parent side, additional bolts are needed to clamp the two sides of
    the housing together.
    """

    disk_joint: DiskJoint = field(default_factory=lambda: DiskJoint(
        movement_angle=60,
    ))

    # Distance between the child/parent arm to the centre
    child_arm_radius: float = 40.0
    parent_arm_radius: float = 40.0

    lip_thickness: float = 5.0
    lip_length: float = 60.0
    hole_pos: list[float] = field(default_factory=lambda: [15, 25])
    parent_arm_width: float = 10.0
    # Angle of the beginning of the parent arm
    parent_arm_angle: float = 180.0

    # Size of the mounting holes
    hole_diam: float = 4.0

    material: Material = Material.RESIN_TRANSPERENT

    angle_neutral: float = 30.0

    def __post_init__(self):
        assert self.child_arm_radius > self.disk_joint.radius_housing
        assert self.parent_arm_radius > self.disk_joint.radius_housing
        self.disk_joint.tongue_length = self.child_arm_radius - self.disk_joint.radius_disk - self.lip_thickness / 2

    @property
    def total_thickness(self):
        return self.disk_joint.total_thickness

    def parent_arm_loc(self) -> Cq.Location:
        """
        2d Location of the centre of the arm surface on the parent side, assuming
        axle is at position 0, and parent direction is -X
        """
        return Cq.Location.from2d(-self.parent_arm_radius, 0, 0)
    def child_arm_loc(self, flip: bool = False) -> Cq.Location:
        """
        2d Location of the centre of the arm surface on the child side, assuming
        axle is at position 0, and parent direction is -X

        Set `flip=True` to indicate that the joint is supposed to be installed upside down
        """
        result = Cq.Location.rot2d(self.angle_neutral) * Cq.Location.from2d(self.child_arm_radius, 0, 180)
        return result.flip_y() if flip else result

    def lip(self) -> Cq.Workplane:
        holes = [
            h
            for i, x in enumerate(self.hole_pos)
            for h in [
                    Hole(x=x, tag=f"conn_top{i}"),
                    Hole(x=-x, tag=f"conn_bot{i}")
            ]
        ]
        mbox = MountingBox(
            length=self.lip_length,
            width=self.disk_joint.total_thickness,
            thickness=self.lip_thickness,
            holes=holes,
            hole_diam=self.hole_diam,
            centred=(True, True),
            generate_side_tags=False,
        )
        return mbox.generate()

    @target(name="child")
    def child_joint(self) -> Cq.Assembly:
        angle = -self.disk_joint.tongue_span / 2
        dz = self.disk_joint.disk_thickness / 2
        # We need to ensure the disk is on the "other" side so
        flip_x = Cq.Location((0, 0, 0), (1, 0, 0), 180)
        flip_z = Cq.Location((0, 0, 0), (0, 0, 1), 180)
        lip_dz = self.lip_thickness
        loc_lip = (
            Cq.Location((0, 0, 0), (0, 1, 0), 180) *
            Cq.Location((-lip_dz, 0, 0), (1, 0, 0), 90) *
            Cq.Location((0, 0, 0), (0, 1, 0), 90)
        )
        loc_disk = flip_x * flip_z * Cq.Location((-self.child_arm_radius, 0, -dz), (0, 0, 1), angle)
        loc_cut_rel = Cq.Location((0, self.disk_joint.spring.radius_inner, -self.disk_joint.disk_bot_thickness))
        disk_cut = self.disk_joint._disk_cut().located(
            loc_lip.inverse * loc_cut_rel * loc_disk)
        result = (
            Cq.Assembly()
            .add(self.lip().cut(disk_cut), name="lip", loc=loc_lip)
            .add(self.disk_joint.disk(), name="disk", loc=loc_disk)
        )
        return result

    @target(name="parent-lower")
    def parent_joint_lower(self) -> Cq.Workplane:
        return self.disk_joint.housing_lower()

    @target(name="parent-upper")
    def parent_joint_upper(self):
        axial_offset = Cq.Location((self.parent_arm_radius, 0, 0))
        housing_dz = self.disk_joint.housing_upper_dz
        conn_h = self.disk_joint.total_thickness
        conn_w = self.parent_arm_width
        connector = (
            Cq.Solid.makeBox(
                length=self.parent_arm_radius,
                width=conn_w,
                height=conn_h,
            ).located(Cq.Location((0, -conn_w/2, 0)))
            #Cq.Solid.makeCylinder(
            #    height=conn_h,
            #    radius=self.parent_arm_radius - self.lip_thickness / 2,
            #    angleDegrees=self.parent_arm_span)
            .cut(Cq.Solid.makeCylinder(
                height=conn_h,
                radius=self.disk_joint.radius_housing,
            ))
            .located(Cq.Location((0, 0, -conn_h / 2)))
            .rotate((0,0,0), (0,0,1), 180)
            #.rotate((0,0,0), (0,0,1), 180-self.parent_arm_span / 2)
        )
        housing = self.disk_joint.housing_upper()
        housing_loc = Cq.Location(
            (0, 0, housing_dz),
            (0, 0, 1),
            -self.disk_joint.tongue_span / 2 + self.angle_neutral
        )
        lip_dz = self.lip_thickness
        result = (
            Cq.Assembly()
            .add(self.lip(), name="lip", loc=
                 Cq.Location((0, 0, 0), (0, 1, 0), 180) *
                 Cq.Location((-lip_dz, 0, 0), (1, 0, 0), 90) *
                 Cq.Location((0, 0, 0), (0, 1, 0), 90))
            .add(housing, name="housing",
                 loc=axial_offset * housing_loc)
            .add(connector, name="connector",
                 loc=axial_offset)
            #.constrain("housing", "Fixed")
            #.constrain("connector", "Fixed")
            #.solve()
        )
        return result

    @assembly()
    def assembly(self, angle: float = 0) -> Cq.Assembly:
        result = (
            Cq.Assembly()
            .addS(self.child_joint(), name="child",
                  role=Role.CHILD, material=self.material)
            .addS(self.parent_joint_lower(), name="parent_lower",
                  role=Role.CASING, material=self.material)
            .addS(self.parent_joint_upper(), name="parent_upper",
                  role=Role.PARENT, material=self.material)
            #.constrain("child/disk?mate_bot", "Fixed")
        )
        result = self.disk_joint.add_constraints(
            result,
            housing_lower="parent_lower",
            housing_upper="parent_upper/housing",
            disk="child/disk",
            angle=angle,
        )
        return result.solve()

if __name__ == '__main__':
    p = ShoulderJoint()
    p.build_all()
    p = DiskJoint()
    p.build_all()