"""
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
    """

    # 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()
    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.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.5),
    Role.CONNECTION: _color('steelblue3', 0.8)
}


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

    Density listed is in g/cm^3 or mg/mm^3
    """

    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["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

def total_mass(self: Cq.Assembly) -> float:
    """
    Calculates the total mass in units of g
    """
    total = 0.0
    for _, a in self.traverse():
        if 'material' not in a.metadata:
            continue
        vol = a.toCompound().Volume()
        total += vol * a.metadata['material'].density
    return total / 1000.0
Cq.Assembly.total_mass = total_mass