From 914bc23582f6c23344e3cc8a6513a89e95b283ff Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Fri, 28 Jun 2024 17:21:30 -0400
Subject: [PATCH] feat: Add directrix tag to hirth joint

---
 nhf/__init__.py |  2 +-
 nhf/joints.py   | 72 ++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 54 insertions(+), 20 deletions(-)

diff --git a/nhf/__init__.py b/nhf/__init__.py
index 2175f24..b582e6a 100644
--- a/nhf/__init__.py
+++ b/nhf/__init__.py
@@ -1 +1 @@
-from nhf.materials import Material
+from nhf.materials import Material, Role
diff --git a/nhf/joints.py b/nhf/joints.py
index e7df7e9..7c227b4 100644
--- a/nhf/joints.py
+++ b/nhf/joints.py
@@ -2,6 +2,13 @@ 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,
@@ -9,9 +16,13 @@ def hirth_joint(radius=60,
                 n_tooth=16,
                 tooth_height=16,
                 tooth_height_inner=2,
-                tol=0.01):
+                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.
     """
     # ensures secant doesn't blow up
     assert n_tooth >= 5
@@ -51,9 +62,14 @@ def hirth_joint(radius=60,
         .val()
     )
     tooth_centre_radius = radius_inner * math.cos(theta)
+    angle_offset = hirth_tooth_angle(n_tooth) / 2 if is_mated else 0
     teeth = (
         Cq.Workplane('XY')
-        .polarArray(radius=tooth_centre_radius, startAngle=0, angle=360, count=n_tooth)
+        .polarArray(
+            radius=tooth_centre_radius,
+            startAngle=angle_offset,
+            angle=360,
+            count=n_tooth)
         .eachpoint(lambda loc: tooth.located(loc))
         .intersect(Cq.Solid.makeCylinder(
             height=base_height + tooth_height,
@@ -66,36 +82,54 @@ def hirth_joint(radius=60,
             height=base_height,
             radius=radius,
             centered=(True, True, False))
-        .faces(">Z").tag("bore")
+        .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("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():
+def hirth_assembly(n_tooth=12):
     """
     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
-        )
+    #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=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))
+        .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