nodeslicer/core/nodeslicer/nodes/base.py

from __future__ import annotations
from typing import Any, TypeAlias
import bpy
import numpy as np

from ..bpy_ns_types import NfEnumProperty, notify_tree_changed

from ...slice import Pather, slice_at_range, SlicePaths
from ....infra.blender.blender_mesh_capture import SlicingCollection, SlicingObject, capture_mesh_edges, SlicingGroup
from ....infra.blender.blender_bridge import get_all_children
from ....infra.blender.blender_slice_nodes_generic import SliceNode, resolve_input

class SliceIntrinsics:
    bed_x: float
    bed_y: float
    feedrate: float
    travel_feedrate: float
    extrusion_rate: float
    retraction: float
    lift: float
    retraction_feedrate: float

    def __init__(self, bed_x: float = 250.0, bed_y: float = 250.0, feedrate: float = 1200.0, travel_feedrate: float = 2400.0, extrusion_rate: float = 0.017, retraction: float = 0.0, lift: float = 0.0, retraction_feedrate: float = 0.0):
        self.bed_x = bed_x
        self.bed_y = bed_y
        self.feedrate = feedrate
        self.travel_feedrate = travel_feedrate
        self.extrusion_rate = extrusion_rate
        self.retraction = retraction
        self.lift = lift
        self.retraction_feedrate = retraction_feedrate

class InputNode(SliceNode):
    input_type: str = ""
    socket_type: str = ""

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        if "bl_label" not in cls.__dict__: cls.bl_label = cls.input_type

    def init(self, context) -> None: 
        stype = self.socket_type or (f"NodeSocket{self.bl_label}" if not self.bl_label.startswith("NodeSocket") else self.bl_label)
        self.add_output(stype)

    def draw_buttons(self, context, layout):
        assert layout
        if self.outputs:
            sock = self.outputs[0]
            if hasattr(sock, "default_value"):
                layout.prop(sock, "default_value", text="")
            elif hasattr(sock, "value"):
                layout.prop(sock, "value", text="")

    def evaluate_output(self, socket_name: str, **_kwargs): 
        out = self.outputs.get(socket_name) or (self.outputs[0] if self.outputs else None)
        if out is None: return None
        return getattr(out, "value", getattr(out, "default_value", None))

_MATH_OPS: dict[str, Any] = {
    'ADD': lambda a, b, t: a + b,
    'SUBTRACT': lambda a, b, t: a - b,
    'MULTIPLY': lambda a, b, t: a * b,
    'DIVIDE': lambda a, b, t: a / b if b != 0 else 0,
    'ABS': lambda a, b, t: abs(a),
    'NEGATE': lambda a, b, t: -a,
    'POWER': lambda a, b, t: t(a ** b) if not (a == 0 and b < 0) else 0,
    'MIN': lambda a, b, t: min(a, b),
    'MAX': lambda a, b, t: max(a, b),
    'SIGN': lambda a, b, t: (a > 0) - (a < 0),
    'DIVIDE_ROUND': lambda a, b, t: t(round(a / b)) if b != 0 else 0,
    'DIVIDE_FLOOR': lambda a, b, t: a // b if b != 0 else 0,
    'DIVIDE_TRUNCATE': lambda a, b, t: t(a / b) if b != 0 else 0,
    'MODULO_FLOORED': lambda a, b, t: a % b if b != 0 else 0,
    'MODULO_TRUNCATED': lambda a, b, t: a - (b * t(a / b)) if b != 0 else 0,
}

class MathNode(SliceNode):
    bl_label = "Float Math"

    operation: str = NfEnumProperty(
        name="Operation",
        items=[
            ('ADD', "Add", "Add the two values"),
            ('SUBTRACT', "Subtract", "Subtract the second value from the first"),
            ('MULTIPLY', "Multiply", "Multiply the two values"),
            ('DIVIDE', "Divide", "Division"),
            ('ABS', "Absolute", "Absolute value of the first value"),
            ('NEGATE', "Negate", "Negate the first value"),
            ('POWER', "Power", "First value raised to the power of the second"),
            ('MIN', "Minimum", "The smaller of the two values"),
            ('MAX', "Maximum", "The larger of the two values"),
            ('SIGN', "Sign", "The sign of the first value (-1, 0, or 1)"),
            ('DIVIDE_ROUND', "Divide Round", "Division rounded to the nearest integer"),
            ('DIVIDE_FLOOR', "Divide Floor", "Division rounded down"),
            ('DIVIDE_TRUNCATE', "Divide Truncate", "Division rounded towards zero"),
            ('MODULO_FLOORED', "Floored Modulo", "Modulo using floor division"),
            ('MODULO_TRUNCATED', "Truncated Modulo", "Modulo using truncated division"),
        ],
        default='ADD'
    )

    math_type: TypeAlias = float
    socket_label: str = "Float" 
    socket_type: str = "NodeSocketFloat"

    def init(self, context):
        self.add_output(self.socket_type, "Value")

        self.add_input(self.socket_type, "Value")
        self.add_input(self.socket_type, "Value")

    def draw_buttons(self, context, layout):
        assert layout
        layout.prop(self, "operation", text="")

    def evaluate_output(self, socket_name: str, **_kwargs) -> math_type:
        t = self.math_type
        a: t = resolve_input(self.inputs[0]) or 0
        b: t = resolve_input(self.inputs[1]) or 0
        op_fn = _MATH_OPS.get(self.operation)
        if op_fn:
            try:
                return op_fn(a, b, t)
            except Exception:
                return 0
        return 0

class SliceNodeJoiner(SliceNode):
    socket_type: str = ""
    output_socket_type: str = ""
    socket_label: str = ""

    def init(self, context):
        self.add_output(self.output_socket_type or self.socket_type, name=self.socket_label)
        self.add_multi_socket_group(self.socket_type, name=self.socket_label)

    def update(self) -> None:
        super().update()
        from ....infra.blender.list_utils import is_socket_linked_to_list, set_socket_shape
        for inp in self.inputs:
            if getattr(inp, "name", "") == self.socket_label or getattr(inp, "identifier", "").startswith(f"{self.socket_label}_ms_in_"):
                is_list = is_socket_linked_to_list(inp)
                set_socket_shape(inp, "LIST" if is_list else "CIRCLE")

class SliceNodeListJoiner(SliceNode):
    socket_type: str = ""
    output_socket_type: str = ""
    socket_label: str = ""
    output_socket_name: str = ""
    is_list_output: bool = True

    def init(self, context):
        self.add_output(self.output_socket_type or self.socket_type, name=self.output_socket_name or self.socket_label, display_shape='LIST')
        self.add_multi_socket_group(self.socket_type, name=self.socket_label)

    def update(self) -> None:
        super().update()
        from ....infra.blender.list_utils import is_socket_linked_to_list, set_socket_shape
        for inp in self.inputs:
            if getattr(inp, "name", "") == self.socket_label or getattr(inp, "identifier", "").startswith(f"{self.socket_label}_ms_in_"):
                is_list = is_socket_linked_to_list(inp)
                set_socket_shape(inp, "LIST" if is_list else "CIRCLE")

    def evaluate_output(self, socket_name: str, **kwargs) -> list[Any]:
        from ....infra.blender.list_utils import flatten_list, is_list_data
        items: list[Any] = []
        raw_inputs = self.resolve_multi_input(self.socket_label)
        for val in raw_inputs:
            if val is not None:
                if is_list_data(val):
                    items.extend(flatten_list(val))
                else:
                    items.append(val)
        return items

class ExecSliceNode(SliceNode):
    button_label: str = 'Execute'
    
    def _init_exec(self):
        self.add_output("SLICESOCK_ExecutionLink")

    def draw_buttons(self, context, layout):
        assert layout
        layout.operator("slicenode.execute", text=self.button_label).node_name = self.name

    def evaluate_output(self, socket_name: str, **_kwargs) -> Any:
        if socket_name == "Execution Link":
            return self
        return None