From 3ad17f0c3e5522cbd8917b35d571e7c4018c4d31 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Sat, 20 Jul 2024 23:52:11 -0700
Subject: [PATCH] fix: get_abs_location partial

---
 nhf/test.py                   | 17 +++++++++++++++++
 nhf/touhou/houjuu_nue/test.py |  2 +-
 nhf/utils.py                  | 31 +++++++++++++++++++++++++++++--
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/nhf/test.py b/nhf/test.py
index ebef77c..7066a9a 100644
--- a/nhf/test.py
+++ b/nhf/test.py
@@ -163,5 +163,22 @@ class TestUtils(unittest.TestCase):
         self.assertAlmostEqual(bbox.ylen, 15)
         self.assertAlmostEqual(bbox.zlen, 5)
 
+    def test_abs_location(self):
+        box = Cq.Solid.makeBox(1, 1, 1)
+        assembly = (
+            Cq.Assembly()
+            .add(box, name="b1")
+            .add(box, name="b2", loc=Cq.Location((0,0,1)))
+            .add(box, name="b3", loc=Cq.Location((0,0,2)))
+        )
+        (x, y, z), _ = assembly.get_abs_location("b2@faces@>Y").toTuple()
+        self.assertAlmostEqual(x, 0.5)
+        self.assertAlmostEqual(y, 1)
+        self.assertAlmostEqual(z, 1.5)
+        (rx, ry, rz), _ = assembly.get_abs_direction("b2@faces@>Y").toTuple()
+        self.assertAlmostEqual(rx, 0)
+        self.assertAlmostEqual(ry, 1)
+        self.assertAlmostEqual(rz, 0)
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/nhf/touhou/houjuu_nue/test.py b/nhf/touhou/houjuu_nue/test.py
index 9dbb41b..f67c0a2 100644
--- a/nhf/touhou/houjuu_nue/test.py
+++ b/nhf/touhou/houjuu_nue/test.py
@@ -30,7 +30,7 @@ class TestJoints(unittest.TestCase):
         Tests the arm radius
         """
         j = MJ.ShoulderJoint()
-        for deflection in [0, 40, 95, 120]:
+        for deflection in [0, 40, j.angle_max_deflection]:
             with self.subTest(deflection=deflection):
                 a = j.assembly(deflection=deflection)
                 # Axle
diff --git a/nhf/utils.py b/nhf/utils.py
index 3399fce..876f891 100644
--- a/nhf/utils.py
+++ b/nhf/utils.py
@@ -1,9 +1,10 @@
 """
 Utility functions for cadquery objects
 """
-import math
 import functools
+import math
 import cadquery as Cq
+from cadquery.occ_impl.solver import ConstraintSpec
 from nhf import Role
 from typing import Union, Tuple, cast
 from nhf.materials import KEY_ITEM, KEY_MATERIAL
@@ -198,14 +199,40 @@ Cq.Assembly.markPlane = mark_plane
 
 def get_abs_location(self: Cq.Assembly,
                      tag: str) -> Cq.Location:
+    """
+    Gets the location of a tag
+
+    BUG: Currently bugged. See `nhf/test.py` for example
+    """
     name, shape = self._query(tag)
-    loc_self = shape.location()
+    loc_self = Cq.Location(shape.Center())
     loc_parent, _ = self._subloc(name)
     loc = loc_parent * loc_self
     return loc
 
 Cq.Assembly.get_abs_location = get_abs_location
 
+def get_abs_direction(self: Cq.Assembly,
+                      tag: str) -> Cq.Location:
+    """
+    Gets the location of a tag
+    """
+    name, shape = self._query(tag)
+    # Must match `cadquery.occ_impl.solver.ConstraintSpec._getAxis`
+    if isinstance(shape, Cq.Face):
+        vec_dir = shape.normalAt()
+    elif isinstance(shape, Cq.Edge) and shape.geomType() != "CIRCLE":
+        vec_dir = shape.tangentAt()
+    elif isinstance(shape, Cq.Edge) and shape.geomType() == "CIRCLE":
+        vec_dir = shape.normal()
+    else:
+        raise ValueError(f"Cannot construct Axis for {shape}")
+    loc_self = Cq.Location(vec_dir)
+    loc_parent, _ = self._subloc(name)
+    loc = loc_parent * loc_self
+    return loc
+Cq.Assembly.get_abs_direction = get_abs_direction
+
 
 # Tallying functions