AareDAQ/tests/unit/gui/test_data_collection_settings.py

"""Tests for the scan settings panels.

Covers where the numbers in the fields come from (the mounted sample's
spreadsheet row, the user's own values, the panel defaults), that a user value
survives a new sample all the way into the scan request while the settings the
user did not touch keep following the sample, and that a number typed but never
committed does not take effect.
"""

import types
from typing import Any, cast

import pytest
from aarecommon.math.coordinate import Coordinate, SmargonCoordinate
from aarecommon.math.diffraction_geometry import DiffractionGeometry
from aarecommon.models.models import (
    BeamlineStateEnum,
    DataCollectionParameters,
    SampleGeometryModel,
    SampleShortInfo,
)
from PySide6.QtCore import Qt
from PySide6.QtWidgets import QLabel

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.panels.scan_settings_panel import SampleParameters
from aare.gui.scan_logic.raster_grid_manager import RasterGridManager
from aare.gui.widgets.number_line_edit import NumberLineEdit


def _commit(field: NumberLineEdit, text: str):
    """Type into a field and commit it, as Enter or leaving the field does."""
    field.setText(text)
    field.on_editing_finished()


def _mount(panel, **params):
    """Mount a sample whose spreadsheet row holds `params`; same two steps
    update_daq_status takes when the parameters of the mounted sample change."""
    panel._db_params = SampleParameters(**params)
    panel._refresh_fields()


# ---------------------------------------------------------------------------
# The spreadsheet row -> SampleParameters
# ---------------------------------------------------------------------------


def test_sample_parameters_translate_the_spreadsheet_row():
    # The real model, so the column names are checked against aareDB and not
    # against a stand-in that agrees with whatever this file happens to say.
    sample = SampleShortInfo(
        db_id=1,
        puck_name="puck1",
        dewar_name="dewar1",
        sample_name="sample1",
        run_number=1,
        pin=1,
        aaredb_params=DataCollectionParameters(
            targetresolution=1.5,
            transmission=0.2,  # 0-to-1 fraction; aareDB converts the percent sheets
            totalangle=180,
            oscillation=0.1,
            exposure=0.02,
        ),
    )
    assert SampleParameters.from_sample(sample) == SampleParameters(
        resolution_a=1.5,
        transmission=0.2,
        total_angle_deg=180.0,
        image_angle_deg=0.1,
        exp_time_s=0.02,
    )


def test_sample_parameters_are_empty_without_a_sample():
    assert SampleParameters.from_sample(None) == SampleParameters()


def test_a_column_aaredb_does_not_have_is_ignored():
    # aareDB owns these names: a renamed column must leave the panel on its
    # defaults, not break the status loop it is read from.
    params = types.SimpleNamespace(targetresolution=1.5)
    sample = cast(Any, types.SimpleNamespace(aaredb_params=params))
    assert SampleParameters.from_sample(sample) == SampleParameters(resolution_a=1.5)


# ---------------------------------------------------------------------------
# Rotation panel
# ---------------------------------------------------------------------------


@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)


@pytest.fixture
def runnable_panel(panel, monkeypatch):
    """A rotation panel whose Run button reaches the request: beamline ready
    and the file path/precondition dialogs answered."""
    panel._beamline_state = BeamlineStateEnum.SampleAlignment
    monkeypatch.setattr(panel, "check_before_run", lambda scan_kind: True)
    return panel


def test_fields_follow_the_mounted_sample(panel):
    _mount(panel, resolution_a=2.5, transmission=0.5, total_angle_deg=180.0, exp_time_s=0.02)

    assert panel.high_res_enter.committed_value == 2.5
    assert panel.dtz_enter.committed_value == pytest.approx(panel._res_to_dtz(2.5), abs=0.01)
    assert panel.transmission_enter.committed_value == 0.5
    assert panel.total_angle.committed_value == 180.0
    assert panel.image_time_enter.committed_value == 0.02
    # An empty spreadsheet cell leaves the panel default in place.
    assert panel.image_angle.committed_value == panel._default_image_angle


def test_editing_resolution_updates_dtz_and_selects_user_values(panel):
    _commit(panel.high_res_enter, "2.00")

    assert panel._show_user_values
    assert panel._user_radio.isChecked()
    assert panel.dtz_enter.committed_value == pytest.approx(panel._res_to_dtz(2.0), abs=0.01)


def test_editing_dtz_updates_resolution(panel, diffraction):
    _commit(panel.dtz_enter, "250.00")

    assert panel.high_res_enter.committed_value == pytest.approx(
        diffraction.resolution_angstrom(250.0), abs=0.01
    )


def test_user_value_survives_the_next_sample(runnable_panel):
    panel = runnable_panel
    _mount(panel, resolution_a=2.5, transmission=0.5, total_angle_deg=180.0, exp_time_s=0.02)

    # The user overrides the resolution the spreadsheet asks for.
    _commit(panel.high_res_enter, "3.00")
    user_dtz = panel.dtz_enter.committed_value

    # Next sample, same spreadsheet values: the user's resolution stays put...
    _mount(panel, resolution_a=2.5, transmission=0.5, total_angle_deg=180.0, exp_time_s=0.02)
    assert panel.high_res_enter.committed_value == 3.00
    assert panel.dtz_enter.committed_value == user_dtz

    # ...and it is what the scan actually collects at.
    requests = []
    panel.rotation_scan.connect(requests.append)
    panel.run_measurement()
    assert requests[-1].dtz == user_dtz


def test_untouched_settings_keep_following_the_sample(panel):
    _mount(panel, total_angle_deg=180.0, exp_time_s=0.02)
    _commit(panel.image_time_enter, "0.0500")  # only the image time is the user's

    _mount(panel, total_angle_deg=360.0, exp_time_s=0.02)
    assert panel.image_time_enter.committed_value == 0.05
    assert panel.total_angle.committed_value == 360.0


def test_toggle_swaps_between_the_two_sets(panel):
    _mount(panel, resolution_a=2.5, exp_time_s=0.02)
    _commit(panel.image_time_enter, "0.0500")

    panel._database_radio.setChecked(True)  # back to Database values
    assert panel.image_time_enter.committed_value == 0.02
    assert panel.high_res_enter.committed_value == 2.5

    panel._user_radio.setChecked(True)  # the user's set is remembered
    assert panel.image_time_enter.committed_value == 0.05


def test_a_new_sample_is_shown_while_on_database_values(panel):
    _mount(panel, exp_time_s=0.02)
    assert panel.image_time_enter.committed_value == 0.02
    _mount(panel, exp_time_s=0.03)
    assert panel.image_time_enter.committed_value == 0.03


def test_text_that_was_never_entered_does_not_take_effect(runnable_panel):
    panel = runnable_panel
    _mount(panel, resolution_a=2.5, transmission=0.5, total_angle_deg=180.0, exp_time_s=0.02)
    database_dtz = panel.dtz_enter.committed_value

    requests = []
    panel.rotation_scan.connect(requests.append)

    panel.high_res_enter.setText("3.00")  # typed, never entered
    panel.run_measurement()
    assert not panel._show_user_values
    assert requests[-1].dtz == database_dtz

    # Committing it is what makes it the user's.
    panel.high_res_enter.on_editing_finished()
    panel.run_measurement()
    assert panel._show_user_values
    assert requests[-1].dtz == pytest.approx(panel._res_to_dtz(3.00), abs=0.01)


def test_screening_transmission_is_the_panels_own(runnable_panel):
    panel = runnable_panel
    _mount(panel, transmission=0.5)
    _commit(panel.screening_transmission_enter, "0.1000")

    requests = []
    panel.rotation_scan.connect(requests.append)
    panel.run_screening()

    assert requests[-1].transmission == 0.1  # not the 0.5 the rotation uses


def test_downstream_gets_the_active_values(panel):
    dtz_seen = []
    transmission_seen = []
    panel.dtz_updated.connect(dtz_seen.append)
    panel.transmission_updated.connect(transmission_seen.append)

    _commit(panel.dtz_enter, "250.00")
    assert dtz_seen[-1] == 250.0

    # The transmission was never touched, so it keeps following the sample.
    _mount(panel, transmission=0.4)
    assert transmission_seen[-1] == 0.4

    panel._database_radio.setChecked(True)
    assert dtz_seen[-1] == pytest.approx(panel.dtz_enter.committed_value)


# ---------------------------------------------------------------------------
# Raster panel
# ---------------------------------------------------------------------------


@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_is_reported_and_echoed_back(raster_panel):
    emitted = []
    raster_panel.grid_size_updated.connect(lambda x, y: emitted.append((x, y)))

    _commit(raster_panel.width_enter, "20")
    assert emitted[-1][0] == pytest.approx(0.020)

    # The grid manager answers with the geometry it settled on; the field
    # follows it and does not emit again.
    before = len(emitted)
    raster_panel.grid_scan_size_change(4, 5, 0.035, 0.020)
    assert raster_panel.width_enter.committed_value == 35.0
    assert len(emitted) == before


def test_raster_exposure_follows_the_toggle(raster_panel):
    _mount(raster_panel, exp_time_s=0.04)
    assert raster_panel.image_time_enter.committed_value == 0.04

    _commit(raster_panel.image_time_enter, "0.0800")
    _mount(raster_panel, exp_time_s=0.04)
    assert raster_panel.image_time_enter.committed_value == 0.08

    raster_panel._database_radio.setChecked(True)
    assert raster_panel.image_time_enter.committed_value == 0.04


def test_grid_element_size_stays_out_of_the_toggle(raster_panel):
    _commit(raster_panel.width_enter, "20")
    # A panel-only field must not take the panel to user values.
    assert not raster_panel._show_user_values


# ---------------------------------------------------------------------------
# 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
    assert tracked_box.committed_value == 100.0
    qtbot.keyClick(tracked_box, Qt.Key.Key_Return)
    assert seen == [300.0]
    assert tracked_box.committed_value == 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_set_committed_value_keeps_a_value_that_did_not_change(tracked_box, qtbot):
    _type(qtbot, tracked_box, "300")  # typing, no commit
    tracked_box.set_committed_value(100.0)  # what the field already holds
    assert tracked_box.text() == "300"  # the typing survives
    tracked_box.set_committed_value(120.0)  # a real change wins
    assert tracked_box.committed_value == 120.0
    assert tracked_box.property("movestate") == ""


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


def _grid_widget(grid, row: int, col: int):
    item = grid.itemAtPosition(row, col)
    assert item is not None
    widget = item.widget()
    assert widget is not None
    return widget


def test_transmission_rows_are_per_mode(panel, qapp, diffraction):
    # Rotation tab: the rotation transmission moved out of the shared top
    # rows into the Rotation section (row 11, right under the header), below
    # the Screening block and its Run button (row 9).
    grid = panel._layout
    label = _grid_widget(grid, 11, 0)
    assert isinstance(label, QLabel)
    assert label.text() == "Rotation transmission"
    assert _grid_widget(grid, 11, 1) is panel.transmission_enter
    assert _grid_widget(grid, 9, 0) is panel.screening_button

    # Gridscan tab keeps its own transmission in the top rows, without the
    # misleading "Rotation" prefix.
    geom = SampleGeometryModel(
        beam_location_pxl=Coordinate(x=500, y=500),
        pixel_in_mm=0.001,
        aerotech=Coordinate(),
        aerotech_meas=Coordinate(),
        smargon=SmargonCoordinate(sh_mm=Coordinate(), phi_deg=0.0, chi_deg=0.0),
        omega_deg=0.0,
        beam_size_mm=Coordinate(x=0.01, y=0.01),
    )
    raster = RasterDataCollectionPanel(RasterGridManager(geom), diffraction)
    raster_label = _grid_widget(raster._layout, 2, 0)
    assert isinstance(raster_label, QLabel)
    assert raster_label.text() == "Transmission"


# ---------------------------------------------------------------------------
# File path panel: the Run buttons must not overwrite an existing dataset
# ---------------------------------------------------------------------------


@pytest.fixture
def file_panel(qapp):
    from aare.gui.panels.file_path_panel import FilePathPanel

    p = FilePathPanel()
    p.directory_edit.setText("d")
    p.file_prefix_edit.setText("x")
    p.run_number_edit.setValue(1)
    return p


def _taken(monkeypatch, *suffixes):
    from aare.gui.panels import file_path_panel

    monkeypatch.setattr(
        file_path_panel.os.path, "exists", lambda path: any(path.endswith(s) for s in suffixes)
    )


def test_run_is_blocked_with_a_popup_when_the_file_exists(file_panel, monkeypatch):
    from PySide6.QtWidgets import QMessageBox

    _taken(monkeypatch, "data/d/x_001_master.h5")
    boxes = []
    monkeypatch.setattr(QMessageBox, "warning", lambda *a, **k: boxes.append(a))

    # Editing never skips a taken run silently: the name stays literal.
    file_panel.set_scan_kind("rotation")
    assert file_panel.run_number_edit.value() == 1

    assert file_panel.file_path_error_box(scan_kind="rotation") is False
    assert len(boxes) == 1
    assert "data/d/x_001_master.h5" in boxes[0][2]
    # Run number went up by one, so the next click can go ahead.
    assert file_panel.run_number_edit.value() == 2
    assert file_panel.file_path_error_box(scan_kind="rotation") is True
    assert len(boxes) == 1


def test_run_proceeds_when_the_file_is_free(file_panel, monkeypatch):
    from PySide6.QtWidgets import QMessageBox

    _taken(monkeypatch)
    monkeypatch.setattr(QMessageBox, "warning", lambda *a, **k: pytest.fail("no popup expected"))

    assert file_panel.file_path_error_box(scan_kind="rotation") is True
    assert file_panel.run_number_edit.value() == 1


def test_a_derived_dataset_counts_as_taken(file_panel, monkeypatch):
    # X-ray centering writes <run>_raster2d_master.h5, not <run>_master.h5.
    from aare.gui.panels import file_path_panel

    _taken(monkeypatch, "raster/d/x_001_raster2d_master.h5")
    monkeypatch.setattr(file_path_panel.QMessageBox, "warning", lambda *a, **k: None)

    assert file_panel.file_path_error_box(scan_kind="raster") is False
    assert file_panel.run_number_edit.value() == 2


def test_a_directory_or_a_sibling_run_is_not_taken(file_panel, monkeypatch):
    # Only the exact master files count; anything else would block runs that
    # nothing will write.
    _taken(monkeypatch, "raster/d/x_001", "raster/d/x_0011_master.h5", "raster/d/x_001_data.h5")
    file_panel.set_scan_kind("raster")
    assert file_panel.file_path_error_box(scan_kind="raster") is True


def test_every_tab_reaches_the_file_guard(settings_panel, monkeypatch):
    # The tabs live in a QStackedWidget, which reparents them: a parent()
    # lookup of file_path_panel found nothing and every Run skipped the guard.
    from aare.gui.panels import scan_settings_panel

    asked = []
    monkeypatch.setattr(
        settings_panel.file_path_panel,
        "file_path_error_box",
        lambda scan_kind: asked.append(scan_kind) or False,
    )
    monkeypatch.setattr(scan_settings_panel, "precondition_check", lambda *a, **k: True)
    scans = []
    for panel in (settings_panel.screening, settings_panel.raster):
        panel._beamline_state = BeamlineStateEnum.SampleAlignment
    settings_panel.screening.rotation_scan.connect(scans.append)
    settings_panel.simple.rotation_scan.connect(scans.append)
    settings_panel.raster.evaluate_grid.connect(lambda: scans.append("grid"))
    settings_panel.raster.evaluate_grid_auto.connect(lambda: scans.append("auto"))

    settings_panel.screening.run_screening()
    settings_panel.screening.run_measurement()
    settings_panel.simple.run_measurement()
    settings_panel.raster._on_evaluate_clicked()
    settings_panel.raster._on_evaluate_auto_clicked()

    assert asked == ["screening", "rotation", "rotation", "raster", "raster"]
    assert scans == []


def test_simple_tab_run_is_blocked_by_the_file_guard(qapp):
    from PySide6.QtWidgets import QWidget

    from aare.gui.panels.smart_rotation_panel import SimpleRotationSettingsPanel

    holder = QWidget()
    cast(Any, holder).file_path_panel = types.SimpleNamespace(
        file_path_error_box=lambda scan_kind: False
    )
    panel = SimpleRotationSettingsPanel(parent=holder)
    requests = []
    panel.rotation_scan.connect(requests.append)

    panel.run_measurement()

    assert requests == []