classes.py

"""
File consists of several classes for the different elements of a device.
"""
from __future__ import annotations
from enum import Enum, auto
import uuid
from typing import List, Dict
import json
import copy


# TODO
# - Docstrings
# - Minimalbeispiel für KPIs -> halb umgesetzt -> mit get_components veranschaulichen
# - Erlaube Clone bei Assembly (jedes child muss durch durch Klon ersetzt werden)
# - Änderungen an Beispiel umsetzen

# - Erlaube Listen bei add_component und add_assembly  (-> Nä Semester)
# - Gute String Darstellung -> Ist so schon ok bisher? -> Nä Semester
# - Export als GraphViz -> Nä Semeseter



class ComponentCategory(Enum):
    BATTERY = auto()
    MOTOR = auto()
    FRAME = auto()
    WHEEL = auto()
    AXLE = auto()
    GEAR = auto()


class AggregationLayer(Enum):
    SYSTEM = auto()
    ASSEMBLY = auto()
    SUBASSEMBLY = auto()
    COMPONENT = auto()


class LegoComponent:
    def __init__(
            self,
            label: str,
            properties: dict,
            layer: AggregationLayer = AggregationLayer.COMPONENT,
    ) -> None:
        self.uuid: uuid.UUID = uuid.uuid4()
        self.parent: None | LegoAssembly = None
        self.label: str = label
        self.properties: dict = properties
        self.layer: AggregationLayer = layer

    def clone(self, new_description: str = None) -> LegoComponent:
        if new_description is None:
            new_description = self.label
        clone = LegoComponent(new_description, copy.deepcopy(self.properties), self.layer)
        return clone

    def get_root_assembly(self):
        if self.parent is None:
            return None
        current_assembly = self.parent
        while current_assembly.parent is not None:
            current_assembly = current_assembly.parent
        return current_assembly

    def to_dict(self) -> Dict:
        ATTRIBUTES = [
            "uuid",
            "name",
            "category",
            "lego_id",
            "cost",
            "mass",
            "delivery_time",
            "layer",
            "properties",
        ]
        dict_ = {}
        # store attributes
        for attr in ATTRIBUTES:
            dict_[attr] = getattr(self, attr)

        dict_ = {"component": dict_}
        return dict_

    # TODO good string representation
    def __str__(self):
        return self.__repr__()
        return (
            f"Item(id={self.uuid}, item_number={self.lego_id}, "
            f"mass={self.mass}, delivery_time={self.delivery_time}, "
            f"parent_id={self.parent})"
        )

    # TODO good repr representation
    def __repr__(self):
        return f"LegoComponent {self.label} [{self.uuid}]"


class LegoAssembly:
    def __init__(self, label: str, properties: dict, layer: AggregationLayer) -> None:
        self.uuid: uuid.UUID = uuid.uuid4()
        self.parent: None | LegoAssembly = None
        self.label: str = label
        self.properties: dict = properties
        self.layer: AggregationLayer = layer
        self.components: List[LegoComponent] = []
        self.assemblies: List[LegoAssembly] = []

    def add_component(self, component: LegoComponent) -> None:
        if not isinstance(component, LegoComponent):
            raise TypeError(
                f"Argument should be of type {LegoComponent.__name__}, "
                f"got {type(component).__name__} instead."
            )
        if self.get_root_assembly().contains_uuid(component.uuid):
            raise AssertionError(
                f"This assembly or a subassembly already contains the component with ID "
                f"{component.uuid}."
            )
        component.parent = self
        self.components.append(component)

    def add_assembly(self, assembly: LegoAssembly) -> None:
        if not isinstance(assembly, LegoAssembly):
            raise TypeError(
                f"Argument should be of type {LegoAssembly.__name__}, "
                f"got {type(assembly).__name__} instead."
            )
        if self.get_root_assembly().contains_uuid(assembly.uuid):
            raise AssertionError(
                f"This assembly or a subassembly already contains the assembly with ID "
                f"{assembly.uuid}."
            )
        assembly.parent = self
        self.assemblies.append(assembly)

    def children(self) -> Dict[str, List[LegoComponent] | List[LegoAssembly]]:
        return {"components": self.components, "assemblies": self.assemblies}

    def get_component_list(self, max_depth: int = -1) -> List[LegoComponent]:
        component_list = []
        component_list.extend(self.components)
        if max_depth > 0:
            for assembly in self.assemblies:
                component_list.extend(assembly.get_component_list(max_depth - 1))
        return component_list

    def get_root_assembly(self) -> LegoAssembly:
        current_assembly = self
        while current_assembly.parent is not None:
            current_assembly = current_assembly.parent
        return current_assembly

    def contains_uuid(self, uuid_: uuid.UUID):
        # check component ids
        component_ids = list(map(lambda c: c.uuid, self.components))
        if uuid_ in component_ids:
            return True
        # check assembly ids
        assembly_ids = list(map(lambda a: a.uuid, self.assemblies))
        if uuid_ in assembly_ids:
            return True
        # recursively check assemblies
        for assembly in self.assemblies:
            if assembly.contains_uuid(uuid_):
                return True
        return False

    def to_dict(self) -> Dict:
        ATTRIBUTES = ["uuid", "name", "layer", "properties"]
        dict_ = {}
        # store attributes
        for attr in ATTRIBUTES:
            dict_[attr] = getattr(self, attr)
        # store components
        dict_["components"] = [component.to_dict() for component in self.components]
        dict_["assemblies"] = [assembly.to_dict() for assembly in self.assemblies]

        return {"assembly": dict_}

    # TODO find good string representation
    def __repr__(self):
        return f"LegoAssembly {self.label} [{self.uuid}]"


def print_assembly_tree(root, level=0, is_last=False):
    # print component
    assembly_padding = ""
    if level > 0:
        assembly_padding += "│   " * (level - 1)
        if is_last:
            assembly_padding += "└── "
        else:
            assembly_padding += "├── "
    print(f"{assembly_padding}{root}")
    # recursively print child components
    for i, assembly in enumerate(root.assemblies):
        is_last_ = i == len(root.assemblies) - 1 and len(root.components) == 0
        print_assembly_tree(assembly, level + 1, is_last_)
    # print items
    for i, item in enumerate(root.components):
        component_padding = "│   " * level if not is_last else "    "
        component_padding += "├── " if i < len(root.components) - 1 else "└── "
        print(f"{component_padding}{item}")


class KPIEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, uuid.UUID):
            return "kpi-" + str(o)
        if isinstance(o, (ComponentCategory, AggregationLayer)):
            return "kpi-" + o.name
        return super().default(o)