AareDAQ/tests/unit/gui/test_data_collection_settings.py

"""Tests for the data-collection settings override model.

Covers DbOverrideLineEdit (database value vs user's value, persistence across
samples, go-back-to-DB, recover-my-value) and the panel-level Database/My-values
toggle with the dtz<->resolution coupling that must hold in both modes.
"""

import types

import pytest
from aarecommon.math.coordinate import Coordinate, SmargonCoordinate
from aarecommon.math.diffraction_geometry import DiffractionGeometry
from aarecommon.models.models import SampleGeometryModel
from PySide6.QtCore import Qt

from aare.gui.panels.data_collection_settings import DataCollectionSettings
from aare.gui.panels.raster_data_collection import RasterDataCollectionPanel
from aare.gui.panels.rotation_data_collection import RotationDataCollectionPanel
from aare.gui.scan_logic.raster_grid_manager import RasterGridManager
from aare.gui.widgets.number_line_edit import DbOverrideLineEdit, NumberLineEdit


def _edit(field: DbOverrideLineEdit, text: str):
    """Simulate a user typing into and committing a field."""
    field.editor.setText(text)
    field.editor.on_editing_finished()


# ---------------------------------------------------------------------------
# Widget: DbOverrideLineEdit
# ---------------------------------------------------------------------------


def test_db_override_follows_database_until_edited(qapp):
    w = DbOverrideLineEdit(0, 1000, default=200.0, decimals=2)
    assert w.source() == DbOverrideLineEdit.SOURCE_DB
    assert not w.has_user_value()

    w.set_db_value(150.0)
    assert w.value == 150.0  # untouched field tracks the database


def test_db_override_user_value_persists_and_recovers(qapp):
    w = DbOverrideLineEdit(0, 1000, default=200.0, decimals=2)

    _edit(w, "300.00")
    assert w.has_user_value()
    assert w.source() == DbOverrideLineEdit.SOURCE_MINE
    assert w.value == 300.0

    # A new database value (e.g. next sample) must not clobber the user's value.
    w.set_db_value(120.0)
    assert w.value == 300.0

    # Go back to the database value, then recover the user's value.
    w.set_source(DbOverrideLineEdit.SOURCE_DB, emit=False)
    assert w.value == 120.0
    w.set_source(DbOverrideLineEdit.SOURCE_MINE, emit=False)
    assert w.value == 300.0


def test_db_override_reset_forgets_user_value(qapp):
    w = DbOverrideLineEdit(0, 1000, default=200.0, decimals=2)
    _edit(w, "300.00")
    w.reset()
    assert w.value == 200.0
    assert w.source() == DbOverrideLineEdit.SOURCE_DB
    assert not w.has_user_value()


def test_db_override_emits_value_changed_on_toggle(qapp):
    w = DbOverrideLineEdit(0, 1000, default=200.0, decimals=2)
    _edit(w, "300.00")  # mine = 300, db = 200
    seen = []
    w.valueChanged.connect(lambda v: seen.append(v))
    w.set_source(DbOverrideLineEdit.SOURCE_DB)
    assert seen[-1] == 200.0  # toggling pushes the now-active value downstream
    w.set_source(DbOverrideLineEdit.SOURCE_MINE)
    assert seen[-1] == 300.0


# ---------------------------------------------------------------------------
# Panel: toggle + dtz/resolution coupling
# ---------------------------------------------------------------------------


@pytest.fixture
def diffraction():
    return DiffractionGeometry(
        energy_keV=12.0,
        dtz_mm=150.0,
        pixel_size_mm=0.075,
        beam_center_pxl=(1000.0, 1000.0),
        detector_size_pxl=(2000, 2000),
        detector_description="Eiger 16M",
        detector_serial_number="123",
        poni_rot1_rad=0.0,
        poni_rot2_rad=0.0,
    )


@pytest.fixture
def panel(qapp, diffraction):
    return RotationDataCollectionPanel(diffraction=diffraction)


def test_editing_dtz_updates_resolution_and_switches_to_mine(panel, diffraction):
    _edit(panel.dtz_enter, "250.00")
    assert panel._source == DbOverrideLineEdit.SOURCE_MINE
    assert abs(panel.high_res_enter.value - diffraction.resolution_angstrom(250.0)) < 0.01


def test_editing_resolution_updates_dtz(panel, diffraction):
    _edit(panel.high_res_enter, "2.00")
    expected = max(panel.MIN_DTZ, diffraction.calc_dtz_mm(2.0))
    assert abs(panel.dtz_enter.value - expected) < 0.01


def test_coupling_holds_in_database_mode(panel, diffraction):
    # Set a database resolution; dtz database value must stay consistent.
    panel._apply_db_resolution(2.5)
    assert panel._source == DbOverrideLineEdit.SOURCE_DB
    expected_dtz = max(panel.MIN_DTZ, diffraction.calc_dtz_mm(2.5))
    assert abs(panel.dtz_enter.value - expected_dtz) < 0.01
    assert abs(panel.high_res_enter.value - 2.5) < 0.01


def test_toggle_pushes_active_value_downstream(panel):
    emitted = []
    panel.dtz_updated.connect(lambda v: emitted.append(round(v, 2)))

    _edit(panel.dtz_enter, "250.00")  # mine dtz = 250
    mine_dtz = panel.dtz_enter.value

    panel.set_source(DbOverrideLineEdit.SOURCE_DB)
    db_dtz = panel.dtz_enter.value
    assert emitted[-1] == round(db_dtz, 2)  # downstream got the db value

    panel.set_source(DbOverrideLineEdit.SOURCE_MINE)
    assert panel.dtz_enter.value == mine_dtz  # user value recovered
    assert emitted[-1] == round(mine_dtz, 2)  # downstream got the user value


def test_user_override_persists_across_samples(panel, diffraction):
    # Sample 1 loads a database resolution.
    panel._sample = types.SimpleNamespace(db_id=1)
    panel._params = types.SimpleNamespace(
        targetresolution=2.5, transmission=0.5, totalrange=180.0, oscillation=0.1, exposure=0.02
    )
    panel.update_data_collection_parameters()

    # User overrides the exposure time.
    _edit(panel.image_time_enter, "0.0500")
    assert panel.image_time_enter.value == 0.05

    # Sample 2 arrives with a different database exposure.
    panel._sample = types.SimpleNamespace(db_id=2)
    panel._params = types.SimpleNamespace(
        targetresolution=1.8, transmission=1.0, totalrange=360.0, oscillation=0.2, exposure=0.01
    )
    panel.update_data_collection_parameters()

    # In My-values mode the user's exposure persists across the sample change.
    assert panel.image_time_enter.value == 0.05
    # Going back to Database shows the new sample's database value.
    panel.set_source(DbOverrideLineEdit.SOURCE_DB)
    assert panel.image_time_enter.value == 0.01


# ---------------------------------------------------------------------------
# Raster grid element size follows the same override model
# ---------------------------------------------------------------------------


@pytest.fixture
def geom():
    return SampleGeometryModel(
        beam_location_pxl=Coordinate(x=500, y=500),
        pixel_in_mm=0.001,
        aerotech=Coordinate(x=0, y=0, z=0),
        aerotech_meas=Coordinate(x=0, y=0, z=0),
        smargon=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0, z=0), phi_deg=0.0, chi_deg=0.0),
        omega_deg=0.0,
        beam_size_mm=Coordinate(x=0.01, y=0.01),
    )


@pytest.fixture
def raster_panel(qapp, diffraction, geom):
    mgr = RasterGridManager(geom)
    return RasterDataCollectionPanel(raster_mgr=mgr, diffraction=diffraction)


def test_grid_element_size_user_value_persists_over_geometry(raster_panel):
    emitted = []
    raster_panel.grid_size_updated.connect(lambda x, y: emitted.append((x, y)))

    _edit(raster_panel.width_enter, "20")
    assert raster_panel._source == DbOverrideLineEdit.SOURCE_MINE
    assert raster_panel.width_enter.value == 20.0
    assert emitted  # editing pushed the new element size downstream

    # The live grid geometry keeps reporting a different element size; the
    # user's value must persist (only the database side changes).
    raster_panel.width_enter.set_db_value(35.0)
    assert raster_panel.width_enter.value == 20.0

    # Revert to the geometry value, then recover the user's value.
    raster_panel.set_source(DbOverrideLineEdit.SOURCE_DB)
    assert raster_panel.width_enter.value == 35.0
    raster_panel.set_source(DbOverrideLineEdit.SOURCE_MINE)
    assert raster_panel.width_enter.value == 20.0


def test_grid_size_field_shares_panel_toggle(raster_panel):
    # Editing the resolution flips the whole panel to "My values", which the
    # grid-size field shares.
    _edit(raster_panel.high_res_enter, "2.50")
    assert raster_panel._source == DbOverrideLineEdit.SOURCE_MINE
    assert raster_panel.width_enter.source() == DbOverrideLineEdit.SOURCE_MINE


# ---------------------------------------------------------------------------
# DataCollectionSettings: per-tab grid button + auto-center after mount
# ---------------------------------------------------------------------------


@pytest.fixture
def settings_panel(qapp, diffraction, geom):
    return DataCollectionSettings(geom, RasterGridManager(geom), diffraction)


def test_grid_button_only_on_raster_tab(settings_panel):
    assert not settings_panel.bounding_box.isHidden()  # Raster is the default tab
    for idx in (1, 2, 3):
        settings_panel._tab_bar.setCurrentIndex(idx)
        assert settings_panel.bounding_box.isHidden()
    settings_panel._tab_bar.setCurrentIndex(0)
    assert not settings_panel.bounding_box.isHidden()


def test_auto_center_fires_on_mount_only_when_armed(settings_panel):
    clicks = []
    settings_panel.find_tip.clicked.connect(lambda: clicks.append(1))
    settings_panel.auto_center_after_mount.setChecked(True)

    # First status with a sample already mounted = GUI (re)start: hardware
    # must not move, only the id is recorded.
    settings_panel._track_sample(1)
    assert not clicks

    # Real mount after an unmount fires the centering.
    settings_panel._track_sample(None)
    settings_panel._track_sample(2)
    assert len(clicks) == 1

    # Direct sample exchange (no unmount tick in between) also fires.
    settings_panel._track_sample(3)
    assert len(clicks) == 2

    # Remounting the SAME sample counts as a fresh mount.
    settings_panel._track_sample(None)
    settings_panel._track_sample(3)
    assert len(clicks) == 3

    # Disarmed: mounts no longer trigger.
    settings_panel.auto_center_after_mount.setChecked(False)
    settings_panel._track_sample(None)
    settings_panel._track_sample(4)
    assert len(clicks) == 3


# ---------------------------------------------------------------------------
# Pending color while typing (movestate property, committed on Enter/focus-out)
# ---------------------------------------------------------------------------


def _type(qtbot, box, text):
    box.clear()
    qtbot.keyClicks(box, text)


@pytest.fixture
def tracked_box(qtbot):
    box = NumberLineEdit(-720, 720, default=100.0, decimals=2, track_pending=True)
    qtbot.addWidget(box)
    return box


def test_typing_marks_pending_and_enter_commits(tracked_box, qtbot):
    seen = []
    tracked_box.newValue.connect(seen.append)
    _type(qtbot, tracked_box, "300")
    assert tracked_box.property("movestate") == "pending"
    assert not seen  # nothing is applied until Enter/focus-out
    qtbot.keyClick(tracked_box, Qt.Key.Key_Return)
    assert seen == [300.0]
    assert tracked_box.property("movestate") == ""


def test_typing_back_the_applied_value_cancels_pending(tracked_box, qtbot):
    _type(qtbot, tracked_box, "300")
    qtbot.keyClick(tracked_box, Qt.Key.Key_Return)
    _type(qtbot, tracked_box, "300")  # same as what is applied now
    assert tracked_box.property("movestate") == ""


def test_incomplete_entry_is_pending(tracked_box, qtbot):
    _type(qtbot, tracked_box, "-")  # not a number (yet)
    assert tracked_box.property("movestate") == "pending"


def test_untracked_box_never_touches_movestate(qtbot):
    # MotorMoveGroup owns movestate on the boxes it registers; the default
    # NumberLineEdit must therefore stay away from the property entirely.
    box = NumberLineEdit(-720, 720, default=100.0, decimals=2)
    qtbot.addWidget(box)
    _type(qtbot, box, "300")
    assert box.property("movestate") is None


def test_programmatic_update_clears_pending(tracked_box, qtbot):
    _type(qtbot, tracked_box, "300")
    tracked_box.force_update_value(120.0)
    assert tracked_box.property("movestate") == ""
    _type(qtbot, tracked_box, "300")
    tracked_box.update_value(500.0)
    assert tracked_box.property("movestate") == ""


def test_db_override_typing_pending_until_commit(qapp, qtbot):
    w = DbOverrideLineEdit(0, 1000, default=200.0, decimals=2)
    qtbot.addWidget(w)
    _type(qtbot, w.editor, "300")
    assert w.editor.property("movestate") == "pending"
    qtbot.keyClick(w.editor, Qt.Key.Key_Return)
    assert w.editor.property("movestate") == ""
    assert w.value == 300.0
    assert w.source() == DbOverrideLineEdit.SOURCE_MINE


def test_db_override_source_toggle_clears_pending(qapp, qtbot):
    w = DbOverrideLineEdit(0, 1000, default=200.0, decimals=2)
    qtbot.addWidget(w)
    _type(qtbot, w.editor, "300")  # typing, no commit
    assert w.editor.property("movestate") == "pending"
    # a panel refresh (radio toggle, DB push) discards the uncommitted text
    w.set_source(DbOverrideLineEdit.SOURCE_DB, emit=False)
    assert w.editor.property("movestate") == ""
    assert w.value == 200.0


def test_energy_spin_motor_move_semantics(settings_panel, daq_status_factory):
    # same placeholder limits as the Beamline setup row (shared constants)
    assert settings_panel.energy_spin.minimum() == 4.0
    assert settings_panel.energy_spin.maximum() == 20.0
    assert "Beamline Scientist" in settings_panel.energy_spin.toolTip()

    # Same MotorMoveGroup wiring as the Beamline setup row (full state walk
    # tested there); here: stage -> pending, button sends -> moving, readback
    # arrival -> neutral.
    box = settings_panel.energy_spin  # movestate sits on the spin itself
    assert not settings_panel.change_energy_button.isEnabled()

    settings_panel.update_daq_status(daq_status_factory())  # readback 12.0 keV
    assert settings_panel.energy_spin.value() == 12.0

    settings_panel.energy_spin.setValue(12.4)
    assert box.property("movestate") == "pending"
    assert settings_panel.change_energy_button.isEnabled()

    sent = []
    settings_panel.change_energy.connect(sent.append)
    settings_panel.change_energy_button.click()
    assert sent and sent[-1] == pytest.approx(12400.0)
    assert box.property("movestate") == "moving"

    settings_panel._energy_state.update_actual(12.3995)
    assert box.property("movestate") == ""