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,
    framed_camera_distance,
)

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 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 snap marker is a DOM overlay, not a trace (a gl3d restyle costs
    # ~90 ms, which made a trace marker lag behind the cursor).
    assert all(trace.get("meta", {}).get("kind") != "hover" for trace in scene.data)


def test_data_bounds_are_unpadded_model_extents() -> None:
    """Camera framing uses raw model bounds, so they must exclude the padding."""
    scene = PlotlyTraceBuilder().build(_load("cantilever"), SceneOptions())
    assert scene.data_bounds["x"] == pytest.approx((0.0, 5.0))
    assert scene.data_bounds["y"] == pytest.approx((0.0, 0.0))
    assert scene.data_bounds["z"] == pytest.approx((0.0, 0.0))
    # Padded frame bounds stay wider than the raw model bounds.
    assert scene.axis_bounds["x"][0] < scene.data_bounds["x"][0]


def test_framed_camera_distance_is_scale_invariant() -> None:
    """plotly's camera is in normalized scene units, so framing must not grow
    with the model's data-unit size (the old bug rendered models as specks)."""
    project = _load("space_frame_3d")
    small = PlotlyTraceBuilder().build(project, SceneOptions())
    scaled = project.model_copy(deep=True)
    for node in scaled.nodes:
        x, y, z = node.coords
        node.coords = (x * 1000.0, y * 1000.0, z * 1000.0)
    large = PlotlyTraceBuilder().build(scaled, SceneOptions())

    assert framed_camera_distance(large) == pytest.approx(framed_camera_distance(small), rel=1e-9)
    # A normalized eye distance stays small (plotly's default eye is 1.25).
    assert 1.0 < framed_camera_distance(small) < 20.0


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"
    # Turntable keeps camera.up pinned to +Z (a level horizon); plotly's
    # ``orbit`` would rotate the up vector and tip the model over.
    assert scene.layout["scene"]["dragmode"] == "turntable"


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

    from otko.core import Project

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


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


# ── scene recipe (opstool-derived framing) ───────────────────────────────
def test_axis_bounds_are_padded_around_the_model() -> None:
    project = _load("cantilever")
    scene = PlotlyTraceBuilder().build(project, SceneOptions())
    points = np.array([node.coords for node in project.nodes], dtype=float)
    for index, axis in enumerate(("x", "y", "z")):
        low, high = scene.axis_bounds[axis]
        assert low < points[:, index].min() or low <= points[:, index].min()
        assert high > points[:, index].max() or high >= points[:, index].max()
    assert scene.diagonal > 0


def test_planar_model_gets_unit_slack_on_the_flat_axis() -> None:
    project = _load("basic_truss")
    scene = PlotlyTraceBuilder().build(project, SceneOptions())
    low, high = scene.axis_bounds["z"]
    assert high - low > 0  # a degenerate axis must not collapse the view


def test_axis_overrides_pin_ranges_with_autorange_off() -> None:
    scene = PlotlyTraceBuilder().build(_load("cantilever"), SceneOptions())
    overrides = scene.axis_overrides()
    assert set(overrides) == {"xaxis", "yaxis", "zaxis"}
    for axis, override in overrides.items():
        assert override["autorange"] is False
        assert len(override["range"]) == 2
        assert override["range"] == list(scene.axis_bounds[axis[0]])


def test_hover_templates_identify_entities() -> None:
    scene = PlotlyTraceBuilder().build(_load("basic_truss"), SceneOptions())
    (nodes,) = _traces(scene, "nodes")
    (frames,) = _traces(scene, "elements")
    assert nodes["hovertemplate"] == "Node #%{customdata}<extra></extra>"
    assert frames["hovertemplate"] == "Element #%{customdata}<extra></extra>"
    # Hover must not fall back to the raw-id "skip" mode.
    assert "hoverinfo" not in nodes


def test_axis_outline_flag_toggles_grid_and_ticks() -> None:
    plain = RenderStyle()
    outlined = RenderStyle(show_axis_outline=True)
    project = _load("cantilever")
    scene_plain = PlotlyTraceBuilder(plain).build(project, SceneOptions())
    scene_outlined = PlotlyTraceBuilder(outlined).build(project, SceneOptions())

    off_axis = scene_plain.layout["scene"]["xaxis"]
    on_axis = scene_outlined.layout["scene"]["xaxis"]
    assert off_axis["showgrid"] is False and off_axis["showticklabels"] is False
    assert on_axis["showgrid"] is True and on_axis["showticklabels"] is True
    # The coloured axis lines stay visible either way (orientation cue).
    assert off_axis["visible"] is True and on_axis["visible"] is True


class _FakeDeformation:
    """Minimal ``DeformationSource`` stand-in (keeps this test pyvista-free)."""

    def __init__(self, rise: float = 0.0) -> None:
        self._rise = rise

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

    def magnitudes(self, node_ids: list[int]) -> np.ndarray:
        return np.linspace(0.0, 1.0, len(node_ids)) if node_ids else np.zeros(0)


def test_loads_scale_with_magnitude_and_colour_by_pattern() -> None:
    scene = PlotlyTraceBuilder().build(_load("space_frame_3d"), SceneOptions())
    cones = [trace for trace in scene.data if trace.get("name") == "nodal-loads"]
    assert len(cones) == 2, "one cone trace per load pattern"
    lengths = sorted(
        float(np.hypot(np.hypot(u, v), w))
        for trace in cones
        for u, v, w in zip(trace["u"], trace["v"], trace["w"], strict=True)
    )
    # 2.5e4 vs 5e4 kN → the second arrow is twice as long.
    assert lengths[0] > 0.0
    assert lengths[-1] == pytest.approx(2.0 * lengths[0], rel=1e-6)
    assert cones[0]["colorscale"][0][1] != cones[1]["colorscale"][0][1]
    assert "Pattern" in cones[0]["customdata"][0] or "|F|" in cones[0]["customdata"][0]


def test_supports_use_dof_glyphs_not_markers() -> None:
    for name in ("space_frame_3d", "basic_truss"):
        scene = PlotlyTraceBuilder().build(_load(name), SceneOptions())
        (supports,) = _traces(scene, "supports")
        assert supports["type"] == "scatter3d"
        assert supports["mode"] == "lines"
        assert len(supports["x"]) > 0


def test_scalar_colouring_adds_arrays_and_colorbar() -> None:
    project = _load("cantilever")
    scalars = np.linspace(0.0, 3.0, len(project.nodes))
    scene = PlotlyTraceBuilder().build(project, SceneOptions(scalars=scalars, scalar_label="|u|"))

    (nodes,) = _traces(scene, "nodes")
    assert isinstance(nodes["marker"]["color"], list)
    assert nodes["marker"]["showscale"] is True
    assert nodes["marker"]["colorbar"]["title"]["text"] == "|u|"
    assert nodes["marker"]["cmin"] == pytest.approx(0.0)
    assert nodes["marker"]["cmax"] == pytest.approx(3.0)

    (frames,) = _traces(scene, "elements")
    assert isinstance(frames["line"]["color"], list)
    assert frames["line"]["cmin"] == pytest.approx(0.0)
    # Scalar mode overrides the family palette.
    assert "showscale" not in frames["line"]


def test_deformation_auto_colours_by_magnitude() -> None:
    project = _load("cantilever")
    scene = PlotlyTraceBuilder().build(
        project, SceneOptions(deformation=_FakeDeformation(rise=0.5))
    )
    (nodes,) = _traces(scene, "nodes")
    assert isinstance(nodes["marker"]["color"], list)
    assert nodes["marker"]["colorbar"]["title"]["text"] == "|u|"


def test_undeformed_reference_overlay_is_opt_in() -> None:
    project = _load("cantilever")
    deformation = _FakeDeformation(rise=0.5)
    without = PlotlyTraceBuilder().build(project, SceneOptions(deformation=deformation))
    assert not _traces(without, "undeformed-reference")

    with_ghost = PlotlyTraceBuilder().build(
        project, SceneOptions(deformation=deformation, show_undeformed=True)
    )
    (ghost,) = _traces(with_ghost, "undeformed-reference")
    assert ghost["mode"] == "lines"
    assert ghost["opacity"] < 1.0