test_widgets.py

# src/napari_deeplabcut/_tests/test_widgets.py
import os
import types
from pathlib import Path

import numpy as np
import pytest
import yaml
from napari.layers import Image, Tracks
from qtpy.QtWidgets import QScrollArea
from vispy import keys

from napari_deeplabcut import _widgets
from napari_deeplabcut.core import io, keypoints
from napari_deeplabcut.core.io import populate_keypoint_layer_properties
from napari_deeplabcut.ui.color_scheme_display import ColorSchemeDisplay
from napari_deeplabcut.ui.dialogs import ShortcutRow
from napari_deeplabcut.ui.labels_and_dropdown import KeypointsDropdownMenu, LabelPair
from napari_deeplabcut.ui.plots.trajectory import TrajectoryMatplotlibCanvas

from .conftest import force_show


def test_guess_continuous():
    import numpy as np
    from napari.layers.utils import color_manager

    # Patch is applied during KeypointControls init (or import-time depending on your setup)
    # Expect float -> continuous
    assert color_manager.guess_continuous(np.array([0.0]))
    # Expect object/categorical -> NOT continuous
    assert not color_manager.guess_continuous(np.array(["a", "b"], dtype=object))


@pytest.mark.usefixtures("qtbot")
def test_keypoint_controls(keypoint_controls):
    controls = keypoint_controls
    controls.label_mode = "loop"
    assert controls._radio_group.checkedButton().text() == "Loop"
    controls.cycle_through_label_modes()
    assert controls._radio_group.checkedButton().text() == "Sequential"


@pytest.mark.usefixtures("qtbot")
def test_save_layers(viewer, keypoint_controls, points):
    viewer.layers.selection.add(points)
    keypoint_controls._save_layers_dialog()


@pytest.mark.usefixtures("qtbot")
def test_show_trails(viewer, keypoint_controls, store):
    keypoint_controls._stores[store.layer] = store
    viewer.layers.selection.active = store.layer
    keypoint_controls._is_saved = True

    keypoint_controls._trail_cb.setChecked(True)

    trails = keypoint_controls._trails_controller.layer
    assert trails is not None
    assert isinstance(trails, Tracks)
    assert trails.visible is True


@pytest.mark.usefixtures("qtbot")
def test_extract_single_frame(keypoint_controls, viewer, images):
    viewer.layers.selection.add(images)
    keypoint_controls._extract_single_frame()


@pytest.mark.usefixtures("qtbot")
def test_store_crop_coordinates(keypoint_controls, viewer, images, config_path):
    viewer.layers.selection.add(images)
    _ = viewer.add_shapes(np.random.random((4, 3)), shape_type="rectangle")
    # _image_meta is expected to be an ImageMetadata instance
    keypoint_controls._image_meta = _widgets.ImageMetadata(name="fake_video")
    # _store_crop_coordinates now uses _project_path instead of reading "project" from _image_meta
    keypoint_controls._project_path = os.path.dirname(config_path)
    # Stores crop coordinates from a rectangle shape into the project's config.yaml
    keypoint_controls._store_crop_coordinates()


@pytest.mark.usefixtures("qtbot")
def test_toggle_face_color(viewer, points):
    viewer.layers.selection.add(points)
    view = viewer.window._qt_viewer
    # Shortcut 'F' toggles coloring between "id" and "label" for multi-animal datasets
    assert points._face.color_properties.name == "id"
    view.canvas.events.key_press(key=keys.Key("F"))
    assert points._face.color_properties.name == "label"
    view.canvas.events.key_press(key=keys.Key("F"))
    assert points._face.color_properties.name == "id"


@pytest.mark.usefixtures("qtbot")
def test_toggle_edge_color(viewer, points):
    viewer.layers.selection.add(points)
    view = viewer.window._qt_viewer
    # Shortcut 'E' toggles border width between 0 and 2
    np.testing.assert_array_equal(points.border_width, 0)
    view.canvas.events.key_press(key=keys.Key("E"))
    np.testing.assert_array_equal(points.border_width, 2)


@pytest.mark.usefixtures("qtbot")
def test_dropdown_menu(qtbot):
    widget = _widgets.DropdownMenu(list("abc"))
    qtbot.add_widget(widget)
    # Ensure update_to selects the correct item
    widget.update_to("c")
    assert widget.currentText() == "c"
    widget.reset()  # Reset should always select the first item
    assert widget.currentText() == "a"


@pytest.mark.usefixtures("qtbot")
def test_keypoints_dropdown_menu_selection_updates_store(store, qtbot):
    widget = KeypointsDropdownMenu(store)
    qtbot.add_widget(widget)
    id_menu = widget.menus.get("id")
    label_menu = widget.menus["label"]

    # If multi-animal, switch ID and ensure the store's id updates
    if id_menu and id_menu.count() > 1:
        id_menu.setCurrentIndex(1)
        assert store.current_id == id_menu.currentText()

    # Switch label and ensure the store's label updates
    if label_menu.count() > 1:
        label_menu.setCurrentIndex(1)
        assert store.current_label == label_menu.currentText()


@pytest.mark.usefixtures("qtbot")
def test_keypoints_dropdown_menu_single_animal_has_no_id_menu(single_animal_store, qtbot):
    widget = KeypointsDropdownMenu(single_animal_store)
    qtbot.add_widget(widget)
    assert "id" not in widget.menus
    assert "label" in widget.menus
    assert widget.menus["label"].count() > 0


@pytest.mark.usefixtures("qtbot")
def test_keypoints_dropdown_menu(store, qtbot):
    widget = KeypointsDropdownMenu(store)
    qtbot.add_widget(widget)
    # Menus for both "id" and "label" should exist; label menu reflects current keypoint
    # This confirms we have multi-animal data
    id_menu = widget.menus["id"]
    label_menu = widget.menus["label"]
    # Baseline: labels for the first ID
    first_id = store.ids[0]
    expected_labels_first = widget.id2label[first_id]
    assert [label_menu.itemText(i) for i in range(label_menu.count())] == expected_labels_first

    # Switch to a different valid ID (if present) and ensure labels update accordingly
    if len(store.ids) > 1:
        # Change selection in the actual menu; this triggers refresh_label_menu via signal
        id_menu.setCurrentIndex(1)
        second_id = id_menu.currentText()
        expected_labels_second = widget.id2label[second_id]
        assert [label_menu.itemText(i) for i in range(label_menu.count())] == expected_labels_second


@pytest.mark.usefixtures("qtbot")
def test_keypoints_dropdown_menu_unknown_id_yields_empty_list(store):
    # If an invalid ID is selected, the label menu should be empty
    widget = KeypointsDropdownMenu(store)
    label_menu = widget.menus["label"]
    widget.refresh_label_menu("__NON_EXISTENT_ID__")
    assert label_menu.count() == 0  # defaultdict(list) → no labels


@pytest.mark.usefixtures("qtbot")
def test_keypoints_dropdown_menu_updates_from_store_current_properties(store, qtbot):
    widget = KeypointsDropdownMenu(store)
    qtbot.add_widget(widget)
    id_menu = widget.menus.get("id")
    label_menu = widget.menus["label"]

    # Pick a valid keypoint (label/id pair) and set it as current
    target = store._keypoints[min(2, len(store._keypoints) - 1)]
    store.current_keypoint = target

    # Simulate event callback
    widget.update_menus(event=None)

    if id_menu:
        assert id_menu.currentText() == target.id
    assert label_menu.currentText() == target.label


@pytest.mark.usefixtures("qtbot")
def test_keypoints_dropdown_menu_smart_reset(store, qtbot):
    widget = KeypointsDropdownMenu(store)
    qtbot.add_widget(widget)
    force_show(widget, qtbot)
    label_menu = widget.menus["label"]
    label_menu.update_to("kpt_2")
    widget._locked = True
    widget.smart_reset(event=None)
    # Locked state prevents reset; current selection remains unchanged
    assert label_menu.currentText() == "kpt_2"
    widget._locked = False
    # Unlocked: smart_reset picks the first unlabeled keypoint (or defaults to first)
    widget.smart_reset(event=None)
    assert label_menu.currentText() == "kpt_0"


@pytest.mark.usefixtures("qtbot")
def test_color_pair(qtbot):
    pair = LabelPair(color="pink", name="kpt", parent=None)
    qtbot.add_widget(pair)
    # LabelPair couples a color swatch with a clickable label
    # Ensure setters update both UI and tooltip
    assert pair.part_name == "kpt"
    assert pair.color == "pink"
    pair.color = "orange"
    pair.part_name = "kpt2"
    assert pair.color_label.toolTip() == "kpt2"


@pytest.mark.usefixtures("qtbot")
def test_color_scheme_display(qtbot):
    widget = ColorSchemeDisplay(None)
    qtbot.add_widget(widget)
    widget._build()
    # Initially empty: no color scheme entries and no layout widgets
    assert not widget.scheme_dict
    assert not widget._container.layout().count()
    widget.add_entry("keypoint", "red")
    assert widget.scheme_dict["keypoint"] == "red"
    assert widget._container.layout().count() == 1


@pytest.mark.usefixtures("qtbot")
def test_matplotlib_canvas_initialization_and_slider(viewer, points, qtbot):
    # Create the canvas widget
    canvas = TrajectoryMatplotlibCanvas(viewer)
    qtbot.add_widget(canvas)

    # Simulate adding a Points layer (triggers _load_dataframe)
    viewer.layers.selection.add(points)
    canvas._load_dataframe()

    # Ensure dataframe loaded and lines plotted
    assert canvas.df is not None
    assert len(canvas._lines) > 0
    assert canvas.ax.get_xlabel() == "Frame"
    assert canvas.ax.get_ylabel() == "Y position"

    # Test slider updates
    initial_window = canvas._window
    canvas.slider.setValue(initial_window + 100)
    assert canvas._window == initial_window + 100
    assert canvas.slider_value.text() == str(initial_window + 100)

    # Test plot refresh does nothing when plot is hidden
    canvas.update_plot_range(event=type("Event", (), {"value": [5]}))
    assert canvas._n == 0
    # Test plot refresh on frame change (forced as it is hidden)
    canvas.update_plot_range(event=type("Event", (), {"value": [5]}), force=True)
    assert canvas._n == 5
    # Check that x-limits reflect the new window
    start, end = canvas.ax.get_xlim()
    assert start <= 5 <= end


@pytest.fixture(autouse=True)
def _no_autodock(monkeypatch):
    """
    Prevent the QTimer.singleShot in KeypointControls.__init__ from auto-calling
    silently_dock_matplotlib_canvas during tests, which would otherwise race
    with these scenarios and make assertions flaky.
    """
    monkeypatch.setattr(_widgets.QTimer, "singleShot", lambda *args, **kwargs: None)


@pytest.mark.usefixtures("qtbot")
def test_ensure_mpl_canvas_docked_already_docked(keypoint_controls, qtbot, monkeypatch):
    """If already docked, it must be a no-op: do not call add_dock_widget again."""
    controls = keypoint_controls
    controls._mpl_docked = True  # simulate already docked

    called = {"count": 0}

    def fake_add_dock_widget(*args, **kwargs):
        called["count"] += 1

    # Ensure it wouldn't try to dock again
    monkeypatch.setattr(controls.viewer.window, "add_dock_widget", fake_add_dock_widget)

    controls._ensure_traj_canvas_docked()
    assert called["count"] == 0, "add_dock_widget should not be called when already docked"
    assert controls._mpl_docked is True  # stays docked


@pytest.mark.usefixtures("qtbot")
def test_ensure_mpl_canvas_docked_missing_window(keypoint_controls, qtbot):
    """If viewer has no window attribute, method should safely no-op."""
    controls = keypoint_controls
    qtbot.add_widget(controls)

    # Swap the viewer for a minimal stub object with *no* 'window' attribute
    controls.viewer = types.SimpleNamespace()  # no 'window'

    controls._mpl_docked = False
    controls._ensure_traj_canvas_docked()

    # Nothing should change; crucially, no exceptions should be raised
    assert controls._mpl_docked is False


@pytest.mark.usefixtures("qtbot")
def test_trajectory_loader_ignores_invalid_properties(viewer, keypoint_controls, make_real_header_factory):
    header = make_real_header_factory(individuals=("",))
    md = populate_keypoint_layer_properties(
        header,
        labels=["bodypart1"],
        ids=[""],
        likelihood=np.array([1.0], dtype=float),
        paths=[],
        colormap="viridis",
    )
    md["properties"]["label"] = [np.nan]  # invalid

    layer = viewer.add_points(np.array([[0.0, 10.0, 20.0]]), **md)
    assert layer is not None
    assert keypoint_controls._traj_mpl_canvas.df is None  # loader should have bailed out safely


@pytest.mark.usefixtures("qtbot")
def test_ensure_traj_canvas_docked_missing_qt_window(keypoint_controls, qtbot):
    """If window._qt_window is None, method should safely no-op."""
    controls = keypoint_controls
    qtbot.add_widget(controls)

    class DummyWindow:
        def __init__(self):
            self._qt_window = None  # simulate missing Qt window

        def add_dock_widget(self, *args, **kwargs):
            raise AssertionError("add_dock_widget should not be called when _qt_window is None")

    controls.viewer = types.SimpleNamespace(window=DummyWindow())

    controls._mpl_docked = False
    controls._ensure_traj_canvas_docked()

    # Still undocked, no crash
    assert controls._mpl_docked is False


@pytest.mark.usefixtures("qtbot")
def test_ensure_mpl_canvas_docked_exception_during_docking(keypoint_controls, qtbot):
    """If add_dock_widget raises, method should catch, log, and remain undocked (no crash)."""
    controls = keypoint_controls
    qtbot.add_widget(controls)

    class DummyWindow:
        def __init__(self):
            self._qt_window = object()  # present → attempt docking

        def add_dock_widget(self, *args, **kwargs):
            raise RuntimeError("boom")

    controls.viewer = types.SimpleNamespace(window=DummyWindow())

    controls._mpl_docked = False

    # Should not raise
    controls._ensure_traj_canvas_docked()

    # Docking failed → remains undocked
    assert controls._mpl_docked is False


@pytest.mark.usefixtures("qtbot")
def test_display_shortcuts_dialog(keypoint_controls, qtbot):
    """Ensure that the Shortcuts dialog can be created and shown without errors."""
    controls = keypoint_controls
    qtbot.add_widget(controls)

    # Create the dialog directly
    dlg = _widgets.Shortcuts(controls)
    qtbot.add_widget(dlg)

    # Show it non-modally
    dlg.show()
    qtbot.waitExposed(dlg)

    # Verify it is visible
    assert dlg.isVisible()
    assert dlg.windowTitle() == "Keyboard shortcuts"
    assert dlg.findChildren(QScrollArea)
    assert dlg.findChildren(ShortcutRow)


# NOTE SuperAnimal keypoints functionality and testing may need an overhaul in the future:
# these tests currently exercise only a narrow "everything fine" path and rely on specific metadata
# layout and SuperAnimal conversion-table conventions, which makes them susceptible to API changes
@pytest.mark.usefixtures("qtbot")
def test_widget_load_superkeypoints_diagram(keypoint_controls, viewer, qtbot, points, monkeypatch):
    # Arrange: conversion table uses *realistic* keys (not SK1/SK2),
    # and does not depend on any asset conventions.
    layer = points
    super_animal = "superanimal_quadruped"
    layer.metadata["tables"] = {super_animal: {"kp1": "nose", "kp2": "upper_jaw"}}

    # Arrange: stub I/O so the test doesn't depend on installed assets
    dummy_img = np.zeros((8, 8), dtype=np.uint8)
    dummy_superkpts = {
        "nose": [1.0, 2.0],
        "upper_jaw": [3.0, 4.0],
    }
    monkeypatch.setattr(io, "load_superkeypoints_diagram", lambda name: dummy_img)
    monkeypatch.setattr(io, "load_superkeypoints", lambda name: dummy_superkpts)

    n_layers_before = len(viewer.layers)

    # Act
    keypoint_controls.load_superkeypoints_diagram()

    # Assert: one new image layer is added
    assert len(viewer.layers) == n_layers_before + 1
    assert isinstance(viewer.layers[-1], Image)
    assert viewer.layers[-1].data.shape == dummy_img.shape

    # Assert: labels match the table keys (reference keypoints)
    assert list(layer.properties["label"]) == ["kp1", "kp2"]

    # Assert: points data updated to [0, x, y] for each mapping
    assert layer.data.shape == (2, 3)
    assert np.allclose(layer.data[:, 0], 0.0)
    assert np.allclose(layer.data[:, 1:], np.array([[1.0, 2.0], [3.0, 4.0]]))

    # Assert: UI updated
    assert keypoint_controls._keypoint_mapping_button.text() == "Map keypoints"
    assert keypoint_controls._keypoint_mapping_button.text() == "Map keypoints"


@pytest.mark.usefixtures("qtbot")
def test_widget_map_keypoints_writes_to_config(keypoint_controls, qtbot, points, config_path, monkeypatch):
    controls = keypoint_controls
    qtbot.add_widget(controls)

    # Arrange: ensure the points layer has some data (shape: [t, x, y])
    points.data = np.array(
        [
            [0.0, 10.0, 20.0],
            [0.0, 30.0, 40.0],
        ],
        dtype=float,
    )

    # Arrange: provide the metadata that _map_keypoints expects
    # _map_keypoints builds config_path as Path(project)/"config.yaml"
    project_dir = Path(config_path).parent
    points.metadata["project"] = str(project_dir)
    points.metadata["tables"] = {"superanimal_quadruped": {}}

    import pandas as pd

    from napari_deeplabcut.config.models import DLCHeaderModel

    cols = pd.MultiIndex.from_product(
        [["S"], [""], ["bp1", "bp2"], ["x", "y"]],
        names=["scorer", "individuals", "bodyparts", "coords"],
    )
    points.metadata["header"] = DLCHeaderModel(
        columns=cols,
    )

    # Ensure config file exists (some setups create it already; this is safe)
    Path(config_path).write_text("{}", encoding="utf-8")

    # Arrange: stub superkeypoints + nearest-neighbor results to be deterministic
    # Your JSON is dict(key -> [x,y]) so we mimic that.
    dummy_superkpts = {"nose": [0.0, 0.0], "upper_jaw": [1.0, 1.0]}
    monkeypatch.setattr(io, "load_superkeypoints", lambda name: dummy_superkpts)

    # neighbors indices correspond to ordering of list(dummy_superkpts)
    # Here: ["nose", "upper_jaw"] -> indices [0, 1]
    monkeypatch.setattr(keypoints, "_find_nearest_neighbors", lambda xy, xy_ref: np.array([0, 1]))

    # If your io.load_config / io.write_config do more than YAML I/O,
    # you can keep them. Otherwise stubbing them makes the test isolated.
    def _load_config(path):
        with open(path, encoding="utf-8") as fh:
            return yaml.safe_load(fh) or {}

    def _write_config(path, cfg):
        with open(path, "w", encoding="utf-8") as fh:
            yaml.safe_dump(cfg, fh, sort_keys=False)

    monkeypatch.setattr(io, "load_config", _load_config)
    monkeypatch.setattr(io, "write_config", _write_config)

    # Act
    controls._map_keypoints("superanimal_quadruped")

    # Assert
    with open(config_path, encoding="utf-8") as fh:
        cfg = yaml.safe_load(fh) or {}

    assert "SuperAnimalConversionTables" in cfg

    # Optional stronger assertion: verify the mapping is written as expected
    assert cfg["SuperAnimalConversionTables"]["superanimal_quadruped"] == {
        "bp1": "nose",
        "bp2": "upper_jaw",
    }


def test_read_config_injects_tables_metadata(tmp_path):
    cfg = {
        "Task": "demo",
        "scorer": "Tester",
        "date": "2026-03-27",
        "multianimalproject": False,
        "identity": "",
        "project_path": str(tmp_path),
        "bodyparts": ["bp1", "bp2"],
        "skeleton": [],
        "pcutoff": 0.6,
        "dotsize": 8,
        "colormap": "viridis",
        "SuperAnimalConversionTables": {
            "superanimal_quadruped": {
                "bp1": "nose",
                "bp2": "upper_jaw",
            }
        },
    }

    config_path = tmp_path / "config.yaml"
    config_path.write_text(yaml.safe_dump(cfg), encoding="utf-8")

    layers = io.read_config(str(config_path))
    _, layer_props, layer_type = layers[0]

    assert layer_type == "points"
    assert "tables" in layer_props["metadata"]
    assert layer_props["metadata"]["tables"] == {
        "superanimal_quadruped": {
            "bp1": "nose",
            "bp2": "upper_jaw",
        }
    }


@pytest.mark.usefixtures("qtbot")
def test_points_layer_with_tables_shows_superkeypoints_button(keypoint_controls, qtbot, points):
    controls = keypoint_controls
    qtbot.add_widget(controls)

    assert not controls._keypoint_mapping_button.isVisible()

    points.metadata["tables"] = {"superanimal_quadruped": {"bp1": "nose", "bp2": "upper_jaw"}}

    # Simulate the same setup path that real inserted/adopted layers use
    controls._setup_points_layer(points, allow_merge=False)

    assert not controls._keypoint_mapping_button.isHidden()
    assert controls._keypoint_mapping_button.text() == "Load superkeypoints diagram"


@pytest.mark.usefixtures("qtbot")
def test_points_layer_with_tables_button_not_lost_on_merge_path(keypoint_controls, qtbot, points, monkeypatch):
    controls = keypoint_controls
    qtbot.add_widget(controls)

    points.metadata["tables"] = {"superanimal_quadruped": {"bp1": "nose"}}

    # Force the merge branch to happen
    monkeypatch.setattr(controls, "_maybe_merge_config_points_layer", lambda layer: True)

    controls._setup_points_layer(points, allow_merge=True)

    assert controls._keypoint_mapping_button.isHidden()


@pytest.mark.usefixtures("qtbot")
def test_video_panel_has_extraction_options(keypoint_controls):
    controls = keypoint_controls
    panel = controls._video_group
    assert panel.extract_button.text() == "Extract current frame"
    assert panel.crop_button.text() == "Save crop to config"
    assert panel.export_labels_cb.text() == "Also export labels"
    assert panel.apply_crop_cb.text() == "Crop to rectangle"


@pytest.mark.usefixtures("qtbot")
def test_extract_single_frame_warns_without_image_layer(keypoint_controls, qtbot, monkeypatch):
    controls = keypoint_controls
    qtbot.addWidget(controls)

    seen = {}

    monkeypatch.setattr(
        "napari_deeplabcut.ui.cropping.show_warning",
        lambda msg: seen.setdefault("warning", msg),
    )

    controls._extract_single_frame()

    assert "No image/video layer is active." in seen["warning"]