From 0582bfd8904979051f6244c0834ce21f037f5ae3 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Wed, 4 Sep 2024 16:30:27 -0700
Subject: [PATCH] feat: Add material list

---
 nhf/materials.py | 121 ++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 108 insertions(+), 13 deletions(-)

diff --git a/nhf/materials.py b/nhf/materials.py
index 7ac4c26..bc172c3 100644
--- a/nhf/materials.py
+++ b/nhf/materials.py
@@ -1,37 +1,132 @@
 """
 A catalog of material properties
 """
-from enum import Enum
+from enum import Enum, Flag, auto
+from typing import Union, Optional
 import cadquery as Cq
 
 def _color(name: str, alpha: float) -> Cq.Color:
     r, g, b, _ = Cq.Color(name).toTuple()
     return Cq.Color(r, g, b, alpha)
 
-class Role(Enum):
+
+KEY_ROLE = 'role'
+KEY_MATERIAL = 'material'
+KEY_ITEM = 'item'
+
+class Role(Flag):
     """
     Describes the role of a part
     """
 
-    # Parent and child components in a load bearing joint
-    PARENT = _color('blue4', 0.6)
-    CHILD = _color('darkorange2', 0.6)
-    STRUCTURE = _color('gray', 0.4)
-    DECORATION = _color('lightseagreen', 0.4)
-    ELECTRONIC = _color('mediumorchid', 0.5)
+    # Externally supplied object
+    FIXTURE = auto()
+    # Parent and child sides of joints
+    PARENT = auto()
+    CHILD = auto()
+    CASING = auto()
+    # Springs, cushions
+    DAMPING = auto()
+    # Main structural support
+    STRUCTURE = auto()
+    DECORATION = auto()
+    ELECTRONIC = auto()
+    MOTION = auto()
+
+    # Fasteners, etc.
+    CONNECTION = auto()
+    HANDLE = auto()
+
+    # Parent and child components in a load bearing joint
+
+    def color_avg(self) -> Cq.Color:
+        r, g, b, a = zip(*[ROLE_COLOR_MAP[component].toTuple() for component in self])
+
+        def avg(li):
+            assert li
+            return sum(li) / len(li)
+        r, g, b, a = avg(r), avg(g), avg(b), avg(a)
+        return Cq.Color(r, g, b, a)
+
+    def color_head(self) -> Cq.Color:
+        head = next(iter(self))
+        return ROLE_COLOR_MAP[head]
+
+
+# Maps roles to their colours
+ROLE_COLOR_MAP = {
+    Role.FIXTURE:    _color('black', 0.2),
+    Role.PARENT:     _color('blue4', 0.6),
+    Role.CASING:     _color('dodgerblue3', 0.6),
+    Role.CHILD:      _color('darkorange2', 0.6),
+    Role.DAMPING:    _color('springgreen', 1.0),
+    Role.STRUCTURE:  _color('gray', 0.4),
+    Role.DECORATION: _color('lightseagreen', 0.4),
+    Role.ELECTRONIC: _color('mediumorchid', 0.7),
+    Role.MOTION:     _color('thistle3', 0.7),
+    Role.CONNECTION: _color('steelblue3', 0.8),
+    Role.HANDLE:     _color('tomato4', 0.8),
+}
 
-    def __init__(self, color: Cq.Color):
-        self.color = color
 
 class Material(Enum):
     """
     A catalog of common material properties
+
+    Density listed is in g/cm^3 or mg/mm^3
     """
 
-    WOOD_BIRCH = 0.8, _color('bisque', 0.9)
-    PLASTIC_PLA = 0.5, _color('azure3', 0.6)
-    ACRYLIC_BLACK = 0.5, _color('gray50', 0.6)
+    WOOD_BIRCH = 0.71, _color('bisque', 0.9)
+    PLASTIC_PLA = 1.2, _color('mistyrose', 0.8)
+    RESIN_TRANSPERENT = 1.17, _color('cadetblue2', 0.6)
+    RESIN_TOUGH_1500 = 1.17, _color('seashell3', 0.7)
+    ACRYLIC_BLACK = 1.18, _color('gray5', 0.8)
+    ACRYLIC_TRANSLUSCENT = 1.18, _color('ivory2', 0.8)
+    ACRYLIC_TRANSPARENT = 1.18, _color('ghostwhite', 0.5)
+    STEEL_SPRING = 7.8, _color('gray', 0.8)
 
     def __init__(self, density: float, color: Cq.Color):
         self.density = density
         self.color = color
+
+def add_with_material_role(
+        self: Cq.Assembly,
+        obj: Union[Cq.Shape, Cq.Workplane, None],
+        loc: Optional[Cq.Location] = None,
+        name: Optional[str] = None,
+        material: Optional[Material] = None,
+        role: Optional[Role] = None) -> Cq.Assembly:
+    """
+    Structural add function which allows specifying material and role
+    """
+    metadata = {}
+    color = None
+    if material:
+        metadata[KEY_MATERIAL] = material
+        color = material.color
+    if role:
+        metadata[KEY_ROLE] = role
+        color = role.color_avg()
+    if len(metadata) == 0:
+        metadata = None
+
+    self.add(obj, loc=loc, name=name, color=color, metadata=metadata)
+    return self
+
+Cq.Assembly.addS = add_with_material_role
+
+def color_by_material(self: Cq.Assembly) -> Cq.Assembly:
+    for _, a in self.traverse():
+        if KEY_MATERIAL not in a.metadata:
+            continue
+        a.color = a.metadata[KEY_MATERIAL].color
+    return self
+Cq.Assembly.color_by_material = color_by_material
+def color_by_role(self: Cq.Assembly, avg: bool = True) -> Cq.Assembly:
+    for _, a in self.traverse():
+        if KEY_ROLE not in a.metadata:
+            continue
+        role = a.metadata[KEY_ROLE]
+        a.color = role.color_avg() if avg else role.color_head()
+    return self
+Cq.Assembly.color_by_role = color_by_role