Cosplay/nhf/materials.py

"""
A catalog of material properties
"""
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(Flag):
    """
    Describes the role of a part
    """

    PARENT = auto()
    CHILD = auto()
    CASING = auto()
    DAMPING = auto()
    STRUCTURE = auto()
    DECORATION = auto()
    ELECTRONIC = auto()
    CONNECTION = 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.PARENT:     _color('blue4', 0.6),
    Role.CASING:     _color('dodgerblue3', 0.6),
    Role.CHILD:      _color('darkorange2', 0.6),
    Role.DAMPING:    _color('springgreen', 0.8),
    Role.STRUCTURE:  _color('gray', 0.4),
    Role.DECORATION: _color('lightseagreen', 0.4),
    Role.ELECTRONIC: _color('mediumorchid', 0.5),
    Role.CONNECTION: _color('steelblue3', 0.8)
}


class Material(Enum):
    """
    A catalog of common material properties
    """

    WOOD_BIRCH = 0.8, _color('bisque', 0.9)
    PLASTIC_PLA = 0.5, _color('mistyrose', 0.8)
    RESIN_TRANSPERENT = 1.1, _color('cadetblue2', 0.6)
    RESIN_TOUGH_1500 = 1.1, _color('seashell3', 0.7)
    ACRYLIC_BLACK = 0.5, _color('gray5', 0.8)
    ACRYLIC_TRANSLUSCENT = 0.5, _color('ivory2', 0.8)
    ACRYLIC_TRANSPARENT = 0.5, _color('ghostwhite', 0.5)
    STEEL_SPRING = 1.0, _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["material"] = material
        color = material.color
    if role:
        metadata["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 'material' not in a.metadata:
            continue
        a.color = a.metadata["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 'role' not in a.metadata:
            continue
        role = a.metadata["role"]
        a.color = role.color_avg() if avg else role.color_head()
    return self
Cq.Assembly.color_by_role = color_by_role