AnyImage/tests/tools/nodes/test_arrangement.py

"""Exercise automatic arrangement without changing geometry or publishing assets."""

import bpy
import numpy as np
import pytest


from tools.nodes.groups.image_plane import build_image_plane_group
from tools.nodes.groups.image_depth_plane import build_image_depth_plane_group
from tools.nodes.groups.image_relief_plane import build_image_relief_plane_group
from tests.support.nodes import CUTOUT_BUILDERS
from tools.nodes.arrangement.arrange import BRANCH_GAP, FIELD_HEIGHT, ROUTE_RUN, arrange_nodes
from tools.nodes.arrangement.flow import geometry_paths, origin_socket, peer_links
from tools.nodes.arrangement.dimensions import RULES, estimate_node_dimensions


@pytest.fixture(params=("plane", "depth_plane", "relief_plane", "SOLID", "DEPTH_SOLID", "DEPTH_BALLOON"))
def node_group(request):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    if request.param == "plane":
        return build_image_plane_group()
    if request.param == "depth_plane":
        return build_image_depth_plane_group(build_image_plane_group())
    if request.param == "relief_plane":
        return build_image_relief_plane_group(build_image_plane_group())
    return CUTOUT_BUILDERS[request.param]()


def _source(socket):
    while socket.node.bl_idname == "NodeReroute":
        socket = socket.node.inputs[0].links[0].from_socket
    node = socket.node
    return ("INPUT" if node.bl_idname == "NodeGroupInput" else node.name, socket.identifier)


def _value(value):
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value
    if isinstance(value, bpy.types.ID):
        return value.as_pointer()
    return tuple(value)


def _program(node_group):
    nodes = {}
    for node in node_group.nodes:
        if node.bl_idname in {"NodeReroute", "NodeGroupInput"}:
            continue
        settings = {
            name: _value(getattr(node, name))
            for name in ("operation", "data_type", "input_type", "domain", "mode", "clamp",
                         "interpolation", "interpolation_type", "extension", "pattern_mode")
            if hasattr(node, name)
        }
        inputs = [
            (socket.identifier, _value(socket.default_value) if hasattr(socket, "default_value") else None,
             [_source(link.from_socket) for link in sorted(
                 socket.links, key=lambda link: link.multi_input_sort_id, reverse=True)])
            for socket in node.inputs
        ]
        paired = getattr(node, "paired_output", None)
        nodes[node.name] = (node.bl_idname, settings, inputs, paired.name if paired else None)
    interface = [
        (item.identifier, item.name, item.in_out, item.socket_type,
         _value(item.default_value) if hasattr(item, "default_value") else None)
        for item in node_group.interface.items_tree if item.item_type == "SOCKET"
    ]
    return nodes, interface


def _surface(node_group):
    mesh = bpy.data.meshes.new("Layout Surface")
    mesh.from_pydata(
        [(x - 2, y - 1, 0) for y in range(3) for x in range(5)], [],
        [(y * 5 + x, y * 5 + x + 1, (y + 1) * 5 + x + 1, (y + 1) * 5 + x)
         for y in range(2) for x in range(4)],
    )
    uv = mesh.uv_layers.new(name="UVMap")
    for loop in mesh.loops:
        x, y, _ = mesh.vertices[loop.vertex_index].co
        uv.data[loop.index].uv = ((x + 2) / 4, (y + 1) / 2)
    profile = mesh.attributes.new("o_balloon", "FLOAT", "POINT")
    for vertex, value in zip(mesh.vertices, profile.data):
        value.value = 0.25 * (1 - abs(vertex.co.x) / 2) * (1 - abs(vertex.co.y))
    mesh.materials.append(bpy.data.materials.new("Image"))
    obj = bpy.data.objects.new("Layout Surface", mesh)
    bpy.context.collection.objects.link(obj)
    modifier = obj.modifiers.new("Surface", "NODES")
    modifier.node_group = node_group
    image = bpy.data.images.new("Layout Depth", width=16, height=8, float_buffer=True)
    image.colorspace_settings.name = "Non-Color"
    image.alpha_mode = "CHANNEL_PACKED"
    yy, xx = np.mgrid[:8, :16]
    pixels = np.ones((8, 16, 4), dtype=np.float32)
    pixels[..., 0] = ((xx + 0.5) / 16 - 0.5) * 4
    pixels[..., 1] = -((yy + 0.5) / 8 - 0.5) * 2
    pixels[..., 2] = 2 + xx * 0.02
    image.pixels.foreach_set(pixels.ravel())
    values = {"Subdivide": 1, "Thickness": 0.2, "Depth Image": image,
              "Uniform Scale": 1.0, "Reference Depth": 2.0, "Depth Scale": 0.8,
              "Smooth": 1, "Normal Smooth": 2, "Split Threshold": 0.25}
    mode = None
    for item in node_group.interface.items_tree:
        if item.item_type != "SOCKET" or item.in_out != "INPUT":
            continue
        if item.name in values:
            modifier[item.identifier] = values[item.name]
        elif item.name == "Mode":
            mode = item.identifier

    def evaluate():
        results = []
        for choice in range(2 if mode else 1):
            if mode:
                modifier[mode] = choice
            obj.update_tag(refresh={"DATA"})
            bpy.context.view_layer.update()
            result = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
            evaluated = result.to_mesh()
            try:
                assert len(evaluated.vertices) > 0
                results.append((np.array([vertex.co[:] for vertex in evaluated.vertices]),
                                [tuple(face.vertices) for face in evaluated.polygons]))
            finally:
                result.to_mesh_clear()
        return results

    return evaluate


def test_arrangement_preserves_each_asset_program(node_group):
    program = _program(node_group)
    evaluate = _surface(node_group)
    geometry = evaluate()
    arrange_nodes(node_group)
    assert _program(node_group) == program

    for node in node_group.nodes:
        if node.bl_idname == "NodeGroupInput":
            targets = {link.to_node for socket in node.outputs for link in socket.links}
            assert len(targets) == 1
            assert all(socket.hide for socket in node.outputs if not socket.is_linked)

    for (before_points, before_faces), (after_points, after_faces) in zip(geometry, evaluate()):
        np.testing.assert_allclose(after_points, before_points, atol=1e-6)
        assert after_faces == before_faces
    positions = {
        node.name: tuple(node.location) for node in node_group.nodes
        if node.bl_idname not in {"NodeReroute", "NodeGroupInput"}
    }
    arrange_nodes(node_group)
    assert _program(node_group) == program
    assert {name: tuple(node_group.nodes[name].location) for name in positions} == positions
    for node in node_group.nodes:
        sources = [link.from_socket for link in peer_links(node)]
        if len(sources) < 2:
            continue
        levels = [socket.node.location.y - estimate_node_dimensions(socket.node).sockets[socket] for socket in sources]
        size = estimate_node_dimensions(node)
        output = next(socket for socket in node.outputs if socket in size.sockets)
        assert node.location.y - size.sockets[output] == pytest.approx((min(levels) + max(levels)) / 2, abs=0.001), node.name
        assert all(first >= second - 0.001 for first, second in zip(levels, levels[1:])), node.name


def test_dimensions_include_unlinked_vector_controls_and_collapsed_sockets():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Dimensions", "GeometryNodeTree")
    source = node_group.nodes.new("GeometryNodeInputPosition")
    vector = node_group.nodes.new("ShaderNodeVectorMath")
    vector.operation = "ADD"
    expanded = estimate_node_dimensions(vector)
    node_group.links.new(source.outputs[0], vector.inputs[0])
    linked = estimate_node_dimensions(vector)
    assert expanded.height - linked.height == 3 * RULES["row"]
    vector.hide = True
    collapsed = estimate_node_dimensions(vector)
    assert collapsed.height < linked.height
    assert all(0 < offset < collapsed.height for offset in collapsed.sockets.values())


def test_local_inputs_preserve_identifiers_with_duplicate_display_names():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Controls", "GeometryNodeTree")
    first = node_group.interface.new_socket(name="Thickness", in_out="INPUT", socket_type="NodeSocketFloat")
    second = node_group.interface.new_socket(name="Thickness", in_out="INPUT", socket_type="NodeSocketFloat")
    first.default_value, second.default_value = 0.25, 0.75
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    controls = node_group.nodes.new("NodeGroupInput")
    grid = node_group.nodes.new("GeometryNodeMeshGrid")
    output = node_group.nodes.new("NodeGroupOutput")
    node_group.links.new(controls.outputs[first.identifier], grid.inputs["Size X"])
    node_group.links.new(controls.outputs[second.identifier], grid.inputs["Size Y"])
    node_group.links.new(grid.outputs["Mesh"], output.inputs["Geometry"])
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program


@pytest.mark.parametrize("branch_lengths", [(1, 1), (1, 3), (3, 1), (1, 2, 3)])
def test_geometry_paths_order_join_branches_and_separate_their_shared_input(branch_lengths):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Parallel Geometry", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    source = node_group.nodes.new("GeometryNodeMeshGrid")
    join = node_group.nodes.new("GeometryNodeJoinGeometry")
    output = node_group.nodes.new("NodeGroupOutput")
    branches = []
    for length in branch_lengths:
        geometry = source.outputs["Mesh"]
        branch = []
        for _ in range(length):
            position = node_group.nodes.new("GeometryNodeSetPosition")
            node_group.links.new(geometry, position.inputs["Geometry"])
            geometry = position.outputs["Geometry"]
            branch.append(position)
        node_group.links.new(geometry, join.inputs["Geometry"])
        branches.append(branch)
    node_group.links.new(source.outputs["Mesh"], join.inputs["Geometry"])
    node_group.links.new(join.outputs["Geometry"], output.inputs["Geometry"])

    geometry_parents, paths = geometry_paths(node_group)
    assert set(geometry_parents) == {join, output, *(node for branch in branches for node in branch)}
    assert [join, output] in paths
    assert all(branch in paths for branch in branches)
    assert source not in geometry_parents
    assert sum(len(path) for path in paths) == len(geometry_parents)
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program
    levels = [link.from_node.location.y - estimate_node_dimensions(link.from_node).sockets[link.from_socket]
              for link in peer_links(join)]
    assert levels == sorted(levels, reverse=True)

    if branch_lengths == (1, 2, 3):
        heights = {node: node.location.y for node in geometry_parents}
        field = node_group.nodes.new("GeometryNodeInputPosition").outputs["Position"]
        calculations = [field.node]
        for _ in range(30):
            operation = node_group.nodes.new("ShaderNodeVectorMath")
            operation.operation = "ADD"
            node_group.links.new(field, operation.inputs[0])
            field = operation.outputs["Vector"]
            calculations.append(operation)
        node_group.links.new(field, branches[0][0].inputs["Offset"])
        program = _program(node_group)
        arrange_nodes(node_group)
        assert _program(node_group) == program
        assert {node: node.location.y for node in geometry_parents} == heights
        assert max(node.location.y for node in calculations) - min(
            node.location.y - estimate_node_dimensions(node).height for node in calculations
        ) <= FIELD_HEIGHT


def test_geometry_used_for_field_sampling_stays_with_the_field_branch():
    from tools.nodes.preview.export import _node_group_data
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Sampled Geometry", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    source = node_group.nodes.new("GeometryNodeMeshGrid")
    separate = node_group.nodes.new("GeometryNodeSeparateGeometry")
    sample = node_group.nodes.new("GeometryNodeSampleIndex")
    sample.data_type = "FLOAT_VECTOR"
    position = node_group.nodes.new("GeometryNodeSetPosition")
    output = node_group.nodes.new("NodeGroupOutput")
    node_group.links.new(source.outputs["Mesh"], separate.inputs["Geometry"])
    node_group.links.new(separate.outputs["Selection"], sample.inputs["Geometry"])
    node_group.links.new(sample.outputs["Value"], position.inputs["Position"])
    node_group.links.new(source.outputs["Mesh"], position.inputs["Geometry"])
    node_group.links.new(position.outputs["Geometry"], output.inputs["Geometry"])

    geometry_parents, paths = geometry_paths(node_group)
    assert set(geometry_parents) == {source, position, output}
    assert paths == [[source, position, output]]
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program
    roles = {node["name"]: node["role"] for node in _node_group_data(node_group)["nodes"]}
    assert roles[separate.name] == roles[sample.name] == "field"
    assert roles[position.name] == "geometry"
    node_group.links.new(separate.outputs["Selection"], output.inputs["Geometry"])
    roles = {node["name"]: node["role"] for node in _node_group_data(node_group)["nodes"]}
    assert roles[separate.name] == "geometry"


def test_shared_field_references_branch_near_their_consumers():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Shared Field", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    source = node_group.nodes.new("GeometryNodeInputPosition")
    geometry = node_group.nodes.new("GeometryNodeMeshGrid").outputs["Mesh"]
    consumers = []
    for index in range(8):
        node = node_group.nodes.new("GeometryNodeSetPosition")
        node_group.links.new(geometry, node.inputs["Geometry"])
        if index in (0, 4, 7):
            node_group.links.new(source.outputs["Position"], node.inputs["Offset"])
            consumers.append(node)
        geometry = node.outputs["Geometry"]
    output = node_group.nodes.new("NodeGroupOutput")
    node_group.links.new(geometry, output.inputs["Geometry"])
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program
    for node in consumers[1:]:
        branch = node.inputs["Offset"].links[0].from_node
        assert branch.bl_idname == "NodeReroute"
        assert node.location.x - branch.location.x >= ROUTE_RUN
    assert any(node.bl_idname == "NodeReroute" and len(node.outputs[0].links) > 1
               for node in node_group.nodes)
    assert all(link.to_node.location.x > link.from_node.location.x
               for link in node_group.links
               if link.from_node.bl_idname == link.to_node.bl_idname == "NodeReroute")


@pytest.mark.parametrize("side", ["upper", "lower"])
def test_field_modules_reserve_lane_and_exit_clearance(side):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Branch Rectangles", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    source = node_group.nodes.new("GeometryNodeMeshGrid")
    join = node_group.nodes.new("GeometryNodeJoinGeometry")
    output = node_group.nodes.new("NodeGroupOutput")
    branches = []
    for _ in range(2):
        geometry, branch = source.outputs["Mesh"], []
        for _ in range(5):
            node = node_group.nodes.new("GeometryNodeSetPosition")
            node_group.links.new(geometry, node.inputs["Geometry"])
            geometry = node.outputs["Geometry"]
            branch.append(node)
        branches.append(branch)
        node_group.links.new(geometry, join.inputs["Geometry"])
    node_group.links.new(join.outputs["Geometry"], output.inputs["Geometry"])
    arrange_nodes(node_group)
    branches.sort(key=lambda branch: branch[0].location.y, reverse=True)
    lane = branches[0 if side == "upper" else -1]
    heights = [node.location.y for branch in branches for node in branch]
    fields = []
    for target in (lane[1], lane[3]):
        value = node_group.nodes.new("GeometryNodeInputPosition").outputs[0]
        branch = [value.node]
        for _ in range(8):
            node = node_group.nodes.new("ShaderNodeVectorMath")
            node.operation = "ADD"
            node_group.links.new(value, node.inputs[0])
            value = node.outputs["Vector"]
            branch.append(node)
        node_group.links.new(value, target.inputs["Offset"])
        fields.append(branch)
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program
    assert [node.location.y for branch in branches for node in branch] == heights
    boxes = [(min(node.location.x for node in branch),
              max(node.location.x + estimate_node_dimensions(node).width for node in branch),
              max(node.location.y for node in branch),
              min(node.location.y - estimate_node_dimensions(node).height for node in branch)) for branch in fields]
    core, _ = geometry_paths(node_group)
    core_left = min(node.location.x for node in core)
    for _, right, top, bottom in boxes:
        if right + BRANCH_GAP <= core_left + 0.001:
            continue
        if side == "upper":
            assert bottom >= lane[0].location.y + BRANCH_GAP - 0.001
        else:
            assert top <= lane[0].location.y - estimate_node_dimensions(lane[0]).height - BRANCH_GAP + 0.001
    first, second = boxes
    assert (first[1] + BRANCH_GAP <= second[0] + 0.001 or second[1] + BRANCH_GAP <= first[0] + 0.001
            or first[3] >= second[2] + BRANCH_GAP - 0.001 or second[3] >= first[2] + BRANCH_GAP - 0.001)
    if side == "lower":
        assert any(right + BRANCH_GAP <= core_left + 0.001 for _, right, _, _ in boxes)
    for branch, target, other in zip(fields, (lane[1], lane[3]), reversed(boxes)):
        terminal = branch[-1]
        assert terminal.location.x + terminal.width < target.location.x
        socket = terminal.outputs["Vector"]
        x1 = terminal.location.x + terminal.width
        y1 = terminal.location.y - estimate_node_dimensions(terminal).sockets[socket]
        x2 = target.location.x
        y2 = target.location.y - estimate_node_dimensions(target).sockets[target.inputs["Offset"]]
        bend = max(40, (x2 - x1) * 0.42)
        t = np.linspace(0, 1, 200)
        x = (1-t)**3*x1 + 3*(1-t)**2*t*(x1+bend) + 3*(1-t)*t**2*(x2-bend) + t**3*x2
        y = (1-t)**3*y1 + 3*(1-t)**2*t*y1 + 3*(1-t)*t**2*y2 + t**3*y2
        left, right, top, bottom = other
        assert not np.any((left < x) & (x < right) & (bottom < y) & (y < top))


def test_boolean_operands_follow_socket_order_and_center_the_junction():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Boolean Operands", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    comparisons = [node_group.nodes.new("FunctionNodeCompare") for _ in range(2)]
    operation = node_group.nodes.new("FunctionNodeBooleanMath")
    operation.operation = "OR"
    grid = node_group.nodes.new("GeometryNodeMeshGrid")
    output = node_group.nodes.new("NodeGroupOutput")
    for source, socket in zip(comparisons, operation.inputs):
        node_group.links.new(source.outputs[0], socket)
    node_group.links.new(operation.outputs[0], grid.inputs["Size X"])
    node_group.links.new(grid.outputs["Mesh"], output.inputs["Geometry"])
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program
    levels = [node.location.y - estimate_node_dimensions(node).sockets[node.outputs[0]] for node in comparisons]
    assert levels[0] > levels[1]
    assert comparisons[0].location.x + comparisons[0].width == comparisons[1].location.x + comparisons[1].width
    assert operation.location.y - estimate_node_dimensions(operation).sockets[operation.outputs[0]] == pytest.approx(sum(levels) / 2)


def test_preview_is_read_only_and_escapes_node_names(tmp_path):
    from tools.nodes.preview.export import write_node_preview

    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = build_image_plane_group()
    arrange_nodes(node_group)
    node_group.nodes[0].name = '</script><script>alert("node")</script>'
    program = _program(node_group)
    positions = {node.name: tuple(node.location) for node in node_group.nodes}
    destination = write_node_preview([node_group], tmp_path / "nodes.html")
    markup = destination.read_text(encoding="utf-8")
    assert markup.startswith("<!doctype html>")
    assert '__NODE_DATA__' not in markup
    assert '<script>alert("node")</script>' not in markup
    assert '\\u003c/script>' in markup
    assert _program(node_group) == program
    assert {node.name: tuple(node.location) for node in node_group.nodes} == positions


def test_geometry_bypass_has_two_level_corners_next_to_processing():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Geometry Bypass", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    source = node_group.nodes.new("GeometryNodeMeshGrid")
    transform = node_group.nodes.new("GeometryNodeTransform")
    switch = node_group.nodes.new("GeometryNodeSwitch")
    switch.input_type = "GEOMETRY"
    output = node_group.nodes.new("NodeGroupOutput")
    node_group.links.new(source.outputs["Mesh"], transform.inputs["Geometry"])
    node_group.links.new(source.outputs["Mesh"], switch.inputs["False"])
    node_group.links.new(transform.outputs["Geometry"], switch.inputs["True"])
    node_group.links.new(switch.outputs["Output"], output.inputs["Geometry"])
    program = _program(node_group)
    arrange_nodes(node_group)
    exit_port = switch.inputs["False"].links[0].from_node
    entry_port = exit_port.inputs[0].links[0].from_node
    assert entry_port.bl_idname == exit_port.bl_idname == "NodeReroute"
    assert entry_port.location.y == exit_port.location.y
    assert source.location.x + source.width < entry_port.location.x < exit_port.location.x < switch.location.x
    assert exit_port.location.x + 16 == pytest.approx(transform.location.x + transform.width)
    assert _program(node_group) == program


@pytest.mark.parametrize("zone_type", ["Repeat", "Simulation", "ForeachGeometryElement"])
def test_preview_zone_contains_body_and_excludes_shared_external_sources(zone_type):
    from tools.nodes.preview.export import _node_group_data

    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Zone Preview", "GeometryNodeTree")
    source = node_group.nodes.new("GeometryNodeMeshGrid")
    entry = node_group.nodes.new(f"GeometryNode{zone_type}Input")
    end = node_group.nodes.new(f"GeometryNode{zone_type}Output")
    entry.pair_with_output(end)
    body = node_group.nodes.new("GeometryNodeSetPosition")
    node_group.links.new(source.outputs["Mesh"], entry.inputs["Geometry"])
    geometry = next(socket for socket in entry.outputs if socket.type == "GEOMETRY")
    node_group.links.new(geometry, body.inputs["Geometry"])
    geometry_input = next(socket for socket in end.inputs if socket.type == "GEOMETRY")
    node_group.links.new(body.outputs["Geometry"], geometry_input)
    entry.location, body.location, end.location = (0, 0), (300, 0), (700, 0)
    data = _node_group_data(node_group)
    zone = data["zones"][0]
    names = {data["nodes"][index]["name"] for index in zone["nodes"]}
    assert names == {entry.name, body.name, end.name}
    for index in zone["nodes"]:
        node = data["nodes"][index]
        assert zone["x"] <= node["x"] and node["x"] + node["w"] <= zone["x"] + zone["w"]
        assert zone["y"] <= node["y"] and node["y"] + node["h"] <= zone["y"] + zone["h"]


def test_repeat_geometry_collector_stays_in_the_local_body():
    from tools.nodes.arrangement.zones import node_zones

    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = CUTOUT_BUILDERS["SOLID"]()
    arrange_nodes(node_group)
    entry = node_group.nodes["Repeat Input"]
    end = entry.paired_output
    collector = node_group.nodes["Join Geometry.001"]
    assert collector in next(members for start, _, members in node_zones(node_group) if start == entry)
    assert entry.location.x < collector.location.x < end.location.x
    assert abs(collector.location.y - entry.location.y) <= FIELD_HEIGHT + BRANCH_GAP
    assert entry.location.y == end.location.y


def test_shared_geometry_uses_one_forward_trunk_with_late_branches():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Shared Geometry References", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    shared = node_group.nodes.new("GeometryNodeMeshGrid")
    geometry = node_group.nodes.new("GeometryNodeMeshGrid").outputs["Mesh"]
    samplers = []
    for index in range(10):
        node = node_group.nodes.new("GeometryNodeSetPosition")
        node_group.links.new(geometry, node.inputs["Geometry"])
        geometry = node.outputs["Geometry"]
        if index in (0, 5, 9):
            sample = node_group.nodes.new("GeometryNodeSampleIndex")
            sample.data_type = "FLOAT_VECTOR"
            node_group.links.new(shared.outputs["Mesh"], sample.inputs["Geometry"])
            node_group.links.new(sample.outputs["Value"], node.inputs["Offset"])
            samplers.append(sample)
    output = node_group.nodes.new("NodeGroupOutput")
    node_group.links.new(geometry, output.inputs["Geometry"])
    program = _program(node_group)
    arrange_nodes(node_group)
    for sample in samplers[1:]:
        branch = sample.inputs["Geometry"].links[0].from_node
        assert branch.bl_idname == "NodeReroute"
        assert branch.location.x > shared.location.x + (sample.location.x - shared.location.x) / 2
    assert any(node.bl_idname == "NodeReroute" and len(node.outputs[0].links) > 1
               for node in node_group.nodes if node.inputs and node.inputs[0].is_linked
               and origin_socket(node.inputs[0].links[0].from_socket).node == shared)
    assert _program(node_group) == program


def test_reused_geometry_processing_keeps_its_upstream_result_path():
    from tools.nodes.arrangement.flow import node_roles

    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Reused Processing", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    grid = node_group.nodes.new("GeometryNodeMeshGrid")
    transform = node_group.nodes.new("GeometryNodeTransform")
    position = node_group.nodes.new("GeometryNodeSetPosition")
    output = node_group.nodes.new("NodeGroupOutput")
    node_group.links.new(grid.outputs["Mesh"], transform.inputs["Geometry"])
    node_group.links.new(transform.outputs["Geometry"], position.inputs["Geometry"])
    node_group.links.new(position.outputs["Geometry"], output.inputs["Geometry"])
    for name in ("Selection", "Offset"):
        separate = node_group.nodes.new("GeometryNodeSeparateGeometry")
        sample = node_group.nodes.new("GeometryNodeSampleIndex")
        node_group.links.new(transform.outputs["Geometry"], separate.inputs["Geometry"])
        node_group.links.new(separate.outputs["Selection"], sample.inputs["Geometry"])
        node_group.links.new(sample.outputs["Value"], position.inputs[name])
    arrange_nodes(node_group)
    roles = node_roles(node_group)
    assert all(roles[node] == "geometry" for node in (grid, transform, position, output))
    assert grid.location.y == transform.location.y == position.location.y == output.location.y
    assert all(roles[node] == "field" for node in node_group.nodes
               if node.bl_idname in {"GeometryNodeSampleIndex", "GeometryNodeSeparateGeometry"})


def test_geometry_preparation_shared_by_independent_paths_becomes_a_public_input():
    from tools.nodes.arrangement.flow import node_roles

    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Public Prepared Geometry", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    grid = node_group.nodes.new("GeometryNodeMeshGrid")
    prepare = node_group.nodes.new("GeometryNodeTransform")
    join = node_group.nodes.new("GeometryNodeJoinGeometry")
    output = node_group.nodes.new("NodeGroupOutput")
    node_group.links.new(grid.outputs["Mesh"], prepare.inputs["Geometry"])
    branches = []
    for _ in range(3):
        node = node_group.nodes.new("GeometryNodeSetPosition")
        node_group.links.new(prepare.outputs["Geometry"], node.inputs["Geometry"])
        node_group.links.new(node.outputs["Geometry"], join.inputs["Geometry"])
        branches.append(node)
    node_group.links.new(join.outputs["Geometry"], output.inputs["Geometry"])
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program
    roles = node_roles(node_group)
    assert roles[grid] == roles[prepare] == "field"
    assert all(roles[node] == "geometry" for node in branches)
    assert prepare.location.x + prepare.width < min(node.location.x for node in branches)


def test_repeat_field_connectors_reserve_horizontal_run_at_both_ends():
    from tools.nodes.arrangement.connections import link_endpoints
    from tools.nodes.arrangement.zones import node_zones

    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = CUTOUT_BUILDERS["SOLID"]()
    arrange_nodes(node_group)
    core, _ = geometry_paths(node_group)
    sizes = {node: estimate_node_dimensions(node) for node in node_group.nodes}
    positions = {node: tuple(node.location) for node in node_group.nodes}
    checked = 0
    for _, _, members in node_zones(node_group):
        for link in node_group.links:
            if link.from_node not in members or link.to_node not in members:
                continue
            if (link.from_node in core) == (link.to_node in core):
                continue
            start, end = link_endpoints(link, positions, sizes)
            assert end[0] - start[0] >= abs(end[1] - start[1]) * 0.8 - 0.01, (link.from_node.name, link.to_node.name)
            checked += 1
    assert checked > 0
    entry = node_group.nodes["Repeat Input.001"]
    end = entry.paired_output
    assert end.location.x - entry.location.x < 6500
    source = entry.outputs["Geometry"]
    assert len(source.links) == 1
    branch = source.links[0].to_node
    while branch.bl_idname == "NodeReroute" and len(branch.outputs[0].links) == 1:
        branch = branch.outputs[0].links[0].to_node
    assert branch.bl_idname == "NodeReroute"
    assert len(branch.outputs[0].links) > 1
    deliveries = [link for link in branch.outputs[0].links
                  if link.to_node not in core and link.to_node.bl_idname != "NodeReroute"]
    assert deliveries
    for link in deliveries:
        start, target = link_endpoints(link, positions, sizes)
        assert target[0] - start[0] == pytest.approx(max(ROUTE_RUN, abs(target[1] - start[1]) * 0.8), abs=0.01)
    for port in core:
        if not port.get("arrangement_input") or port.outputs[0].type != "GEOMETRY":
            continue
        source = origin_socket(port.outputs[0]).node
        if source not in core:
            continue
        left = source.location.x + sizes[source].width
        right, level = port.location.x + 8, port.location.y - 8
        for node in node_group.nodes:
            if node in core or node.bl_idname == "NodeReroute":
                continue
            if left < node.location.x + sizes[node].width and node.location.x < right:
                assert not node.location.y - sizes[node].height < level < node.location.y, node.name
    for link in node_group.links:
        if link.from_socket.type != "GEOMETRY":
            continue
        start, target = link_endpoints(link, positions, sizes)
        if abs(start[1] - target[1]) > 0.01:
            continue
        for node in node_group.nodes:
            if node in core or node.bl_idname == "NodeReroute" or node in {link.from_node, link.to_node}:
                continue
            if start[0] < node.location.x + sizes[node].width and node.location.x < target[0]:
                assert not node.location.y - sizes[node].height < start[1] < node.location.y, node.name


def test_expression_skip_links_have_passages_between_intermediate_nodes():
    from tools.nodes.arrangement.connections import link_endpoints

    bpy.ops.wm.read_factory_settings(use_empty=True)
    node_group = bpy.data.node_groups.new("Expression Passages", "GeometryNodeTree")
    node_group.interface.new_socket(name="Geometry", in_out="OUTPUT", socket_type="NodeSocketGeometry")
    grid = node_group.nodes.new("GeometryNodeMeshGrid")
    position = node_group.nodes.new("GeometryNodeSetPosition")
    output = node_group.nodes.new("NodeGroupOutput")
    seed = node_group.nodes.new("GeometryNodeInputIndex")
    previous = node_group.nodes.new("GeometryNodeInputMeshFaceArea").outputs[0]
    fields = [seed, previous.node]
    for _ in range(7):
        operation = node_group.nodes.new("ShaderNodeMath")
        operation.operation = "ADD"
        node_group.links.new(previous, operation.inputs[0])
        node_group.links.new(seed.outputs[0], operation.inputs[1])
        previous = operation.outputs[0]
        fields.append(operation)
    node_group.links.new(grid.outputs["Mesh"], position.inputs["Geometry"])
    node_group.links.new(previous, position.inputs["Offset"])
    node_group.links.new(position.outputs["Geometry"], output.inputs["Geometry"])
    program = _program(node_group)
    arrange_nodes(node_group)
    assert _program(node_group) == program
    sizes = {node: estimate_node_dimensions(node) for node in node_group.nodes}
    positions = {node: tuple(node.location) for node in node_group.nodes}
    assert any(node.get("arrangement_branch") for node in node_group.nodes)
    for link in node_group.links:
        if origin_socket(link.from_socket).node not in fields:
            continue
        (x1, y1), (x2, y2) = link_endpoints(link, positions, sizes)
        bend = max(40, abs(x2 - x1) * 0.42)
        for t in np.linspace(0, 1, 100):
            x = (1-t)**3*x1 + 3*(1-t)**2*t*(x1+bend) + 3*(1-t)*t*t*(x2-bend) + t**3*x2
            y = (1-t)**3*y1 + 3*(1-t)**2*t*y1 + 3*(1-t)*t*t*y2 + t**3*y2
            for node in fields:
                if node in {link.from_node, link.to_node}:
                    continue
                left, top = positions[node]
                assert not (left < x < left + sizes[node].width and top - sizes[node].height < y < top), node.name