otko/tests/unit/test_plotly_trace_builder.py

"""Unit tests for the pure Plotly trace builder (no Qt, no plotly import)."""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

from otko.services import load_project
from otko.views.canvas3d.style import RenderStyle
from otko.views.canvas_plotly.trace_builder import (
    PlotlyTraceBuilder,
    SceneOptions,
)

EXAMPLES = Path(__file__).resolve().parents[2] / "examples"


def _load(name: str):  # type: ignore[no-untyped-def]
    return load_project(EXAMPLES / f"{name}.osmodel")


def _traces(scene, name: str) -> list[dict]:  # type: ignore[no-untyped-def]
    return [trace for trace in scene.data if trace.get("name") == name]


def _kinds(scene) -> list[str]:  # type: ignore[no-untyped-def]
    return [trace.get("meta", {}).get("kind", trace["type"]) for trace in scene.data]


def test_builds_grid_nodes_and_frames() -> None:
    scene = PlotlyTraceBuilder().build(_load("basic_truss"), SceneOptions())
    names = {trace.get("name") for trace in scene.data}
    assert "grid-active" in names
    assert "nodes" in names
    assert "elements" in names
    assert scene.diagonal > 0
    # The hover-snap marker is always present so JS can restyle it.
    assert scene.data[scene.hover_trace]["meta"]["kind"] == "hover"


def test_nodes_carry_ids_as_customdata() -> None:
    project = _load("basic_truss")
    scene = PlotlyTraceBuilder().build(project, SceneOptions())
    (nodes,) = _traces(scene, "nodes")
    assert nodes["meta"]["kind"] == "node"
    assert nodes["customdata"] == [node.id for node in project.nodes]
    assert len(nodes["marker"]["color"]) == len(project.nodes)


def _color_for_index(trace: dict, index: float) -> str:
    """Map a ``line.color`` scalar index to its colour via the colorscale."""
    scale = sorted(trace["line"]["colorscale"], key=lambda item: item[0])
    return scale[int(index)][1]


def test_frames_are_one_trace_with_per_element_colours() -> None:
    """A single line trace carries the palette through its ``line.color`` array."""
    project = _load("cantilever")
    style = RenderStyle()
    scene = PlotlyTraceBuilder(style).build(
        project, SceneOptions(selection_elements=frozenset({2}))
    )
    frames = _traces(scene, "elements")
    assert len(frames) == 1, "selection must not split the frame trace"
    line = frames[0]["line"]
    assert isinstance(line["color"], list)
    assert len(line["color"]) == len(frames[0]["x"])
    assert line["colorscale"]

    colors = line["color"]
    customdata = frames[0]["customdata"]

    def index_of(element_id: int) -> set[float]:
        return {colors[k] for k, value in enumerate(customdata) if value == element_id}

    selected = index_of(2)
    unselected = index_of(1)
    assert len(selected) == 1 and len(unselected) == 1
    assert selected != unselected
    assert _color_for_index(frames[0], next(iter(selected))) == style.selected_color
    assert _color_for_index(frames[0], next(iter(unselected))) == style.element_beam_color


def test_element_palette_distinguishes_beam_from_truss() -> None:
    style = RenderStyle()
    builder = PlotlyTraceBuilder(style)
    beam = _traces(builder.build(_load("cantilever"), SceneOptions()), "elements")[0]
    truss = _traces(builder.build(_load("basic_truss"), SceneOptions()), "elements")[0]
    assert _color_for_index(beam, beam["line"]["color"][0]) == style.element_beam_color
    assert _color_for_index(truss, truss["line"]["color"][0]) == style.element_truss_color


def test_selected_node_is_highlighted_by_colour() -> None:
    project = _load("basic_truss")
    target = project.nodes[1].id
    scene = PlotlyTraceBuilder().build(project, SceneOptions(selection_nodes=frozenset({target})))
    (nodes,) = _traces(scene, "nodes")
    ids = nodes["customdata"]
    assert nodes["marker"]["color"][ids.index(target)] != nodes["marker"]["color"][0]


def test_working_plane_filters_grid_segments() -> None:
    project = _load("basic_truss")
    builder = PlotlyTraceBuilder()
    iso = builder.build(project, SceneOptions())
    plane = builder.build(project, SceneOptions(working_plane=("XY", 0.0)))
    (iso_grid,) = _traces(iso, "grid-active")
    (plane_grid,) = _traces(plane, "grid-active")
    assert len(plane_grid["x"]) <= len(iso_grid["x"])


def test_snap_targets_carry_world_coordinates() -> None:
    scene = PlotlyTraceBuilder().build(_load("basic_truss"), SceneOptions())
    (snap,) = _traces(scene, "snap")
    assert snap["meta"]["kind"] == "snap"
    assert snap["customdata"], "expected at least one snappable intersection"
    assert all(len(point) == 3 for point in snap["customdata"])
    assert all(isinstance(coord, float) for point in snap["customdata"] for coord in point)


def test_extrusions_add_a_mesh3d_trace() -> None:
    project = _load("cantilever")
    scene = PlotlyTraceBuilder().build(project, SceneOptions(show_extrusions=True))
    (mesh,) = _traces(scene, "extrusions")
    assert mesh["type"] == "mesh3d"
    assert mesh["i"] and mesh["j"] and mesh["k"]
    assert max(mesh["i"]) < len(mesh["x"])


def test_local_axes_add_coloured_cones() -> None:
    project = _load("cantilever")
    scene = PlotlyTraceBuilder().build(project, SceneOptions(show_local_axes=True))
    names = {trace.get("name") for trace in scene.data}
    assert {"local-x", "local-y", "local-z"} <= names
    for trace in scene.data:
        if trace.get("name", "").startswith("local-"):
            assert trace["type"] == "cone"
            assert len(trace["u"]) == len(trace["x"]) > 0


def test_labels_traces_are_text_only() -> None:
    project = _load("cantilever")
    scene = PlotlyTraceBuilder().build(
        project, SceneOptions(show_node_labels=True, show_element_labels=True)
    )
    (node_labels,) = _traces(scene, "node-labels")
    (element_labels,) = _traces(scene, "element-labels")
    assert node_labels["mode"] == "text"
    assert len(node_labels["text"]) == len(project.nodes)
    assert element_labels["mode"] == "text"


def test_deformation_shifts_node_coordinates() -> None:
    project = _load("cantilever")

    class _Shift:
        def shifted(self, points: np.ndarray, node_ids: list[int]) -> np.ndarray:
            out = points.copy()
            out[:, 2] += 1.0
            return out

    base = PlotlyTraceBuilder().build(project, SceneOptions())
    moved = PlotlyTraceBuilder().build(project, SceneOptions(deformation=_Shift()))
    (base_nodes,) = _traces(base, "nodes")
    (moved_nodes,) = _traces(moved, "nodes")
    assert moved_nodes["z"] == pytest.approx([z + 1.0 for z in base_nodes["z"]])


def test_payload_is_json_serialisable() -> None:
    scene = PlotlyTraceBuilder().build(
        _load("cantilever"),
        SceneOptions(
            selection_nodes=frozenset({1}),
            selection_elements=frozenset({1}),
            working_plane=("XY", 0.0),
            show_extrusions=True,
            show_local_axes=True,
            show_node_labels=True,
            show_element_labels=True,
        ),
    )
    payload = json.dumps(scene.to_payload())
    assert '"data"' in payload and '"layout"' in payload
    assert scene.layout["scene"]["aspectmode"] == "data"


def test_empty_project_yields_only_the_hover_marker() -> None:
    scene = PlotlyTraceBuilder().build(None, SceneOptions())
    assert scene.data == []
    assert scene.hover_trace == -1

    from otko.core import Project

    empty = PlotlyTraceBuilder().build(Project(ndm=3, ndf=6), SceneOptions())
    assert _kinds(empty) == ["hover"]


def test_frame_trace_meta_marks_elements_pickable() -> None:
    project = _load("cantilever")
    scene = PlotlyTraceBuilder().build(project, SceneOptions())
    (frames,) = _traces(scene, "elements")
    assert frames["meta"]["kind"] == "element"
    # None separators break the line into per-element segments.
    assert None in frames["x"]
    assert len(frames["customdata"]) == len(frames["x"])