bec_widgets/tests/unit_tests/test_curve_settings.py

import json
from unittest.mock import MagicMock, patch

import pytest
from bec_lib.scan_history import ScanHistory
from qtpy.QtWidgets import QComboBox, QVBoxLayout

from bec_widgets.tests.client_mocks import (
    dap_plugin_message,
    inject_scan_history,
    mocked_client,
    mocked_client_with_dap,
)
from bec_widgets.widgets.plots.waveform.settings.curve_settings.curve_setting import CurveSetting
from bec_widgets.widgets.plots.waveform.settings.curve_settings.curve_tree import CurveTree
from bec_widgets.widgets.plots.waveform.waveform import Waveform
from tests.unit_tests.conftest import create_widget

##################################################
# CurveSetting
##################################################


@pytest.fixture
def curve_setting_fixture(qtbot, mocked_client):
    """
    Creates a CurveSetting widget targeting a mock or real Waveform widget.
    """
    wf = create_widget(qtbot, Waveform, client=mocked_client)
    wf.x_mode = "auto"
    curve_setting = create_widget(qtbot, CurveSetting, parent=None, target_widget=wf)
    return curve_setting, wf


def test_curve_setting_init(curve_setting_fixture):
    """
    Ensure CurveSetting constructs properly, with a CurveTree inside
    and an x-axis group box for modes.
    """
    curve_setting, wf = curve_setting_fixture

    # The layout should be QVBoxLayout
    assert isinstance(curve_setting.layout, QVBoxLayout)

    # There's an x_axis_box group and a y_axis_box group
    assert hasattr(curve_setting, "x_axis_box")
    assert hasattr(curve_setting, "y_axis_box")

    # The x_axis_box should contain a QComboBox for mode
    mode_combo = curve_setting.mode_combo
    assert isinstance(mode_combo, QComboBox)
    # Should contain these items: ["auto", "index", "timestamp", "device"]
    expected_modes = ["auto", "index", "timestamp", "device"]
    for m in expected_modes:
        assert m in [
            curve_setting.mode_combo.itemText(i) for i in range(curve_setting.mode_combo.count())
        ]

    # Check that there's a curve_manager inside y_axis_box
    assert hasattr(curve_setting, "curve_manager")
    assert curve_setting.y_axis_box.layout.count() > 0
    assert not curve_setting.device_x.isEnabled()
    assert not curve_setting.signal_x.isEnabled()
    assert "red" not in curve_setting.device_x.styleSheet()
    assert "red" not in curve_setting.signal_x.styleSheet()


def test_curve_setting_accept_changes(curve_setting_fixture, qtbot):
    """
    Test that calling accept_changes() applies x-axis mode changes
    and triggers the CurveTree to send its curve JSON to the target waveform.
    """
    curve_setting, wf = curve_setting_fixture

    # Suppose user chooses "index" from the combo
    curve_setting.mode_combo.setCurrentText("index")
    # The device_x is disabled if not device mode

    # Spy on 'send_curve_json' from the curve_manager
    send_spy = MagicMock()
    curve_setting.curve_manager.send_curve_json = send_spy

    # Call accept_changes()
    curve_setting.accept_changes()

    # Check that we updated the waveform
    assert wf.x_mode == "index"
    # Check that the manager send_curve_json was called
    send_spy.assert_called_once()


def test_curve_setting_switch_device_mode(curve_setting_fixture, qtbot):
    """
    If user chooses device mode from the combo, the device_x line edit should be enabled
    and set to the current wavefrom.x_axis_mode["name"].
    """
    curve_setting, wf = curve_setting_fixture

    # Initially we assume "auto"
    assert curve_setting.mode_combo.currentText() == "auto"
    # Switch to device
    curve_setting.mode_combo.setCurrentText("device")
    assert curve_setting.device_x.isEnabled()

    # This line edit should reflect the waveform.x_axis_mode["name"], or be blank if none
    assert curve_setting.device_x.currentText() == ""


def test_curve_setting_refresh(curve_setting_fixture, qtbot):
    """
    Test that calling refresh() refreshes the embedded CurveTree
    and re-reads the x axis mode from the waveform.
    """
    curve_setting, wf = curve_setting_fixture

    # Suppose the waveform changed x_mode from "auto" to "timestamp" behind the scenes
    wf.x_mode = "timestamp"
    # Spy on the curve_manager
    refresh_spy = MagicMock()
    curve_setting.curve_manager.refresh_from_waveform = refresh_spy

    # Call refresh
    curve_setting.refresh()

    refresh_spy.assert_called_once()
    # The combo should now read "timestamp"
    assert curve_setting.mode_combo.currentText() == "timestamp"


def test_change_device_from_target_widget(curve_setting_fixture, qtbot):
    curve_setting, wf = curve_setting_fixture

    wf.x_mode = "samx"

    # Call refresh
    curve_setting.refresh()

    assert curve_setting.mode_combo.currentText() == "device"
    assert curve_setting.device_x.isEnabled()
    assert curve_setting.device_x.currentText() == wf.x_axis_mode["name"]
    assert curve_setting.signal_x.currentText() == f"{wf.x_axis_mode['entry']} (readback)"


##################################################
# CurveTree
##################################################


@pytest.fixture
def curve_tree_fixture(qtbot, mocked_client_with_dap):
    """
    Creates a CurveTree widget referencing a mocked or real Waveform.
    """
    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
    wf.color_palette = "plasma"
    curve_tree = create_widget(qtbot, CurveTree, parent=None, waveform=wf)
    return curve_tree, wf


def test_curve_tree_init(curve_tree_fixture):
    """
    Test that the CurveTree initializes properly with references to the waveform,
    sets up the toolbar, and an empty QTreeWidget.
    """
    curve_tree, wf = curve_tree_fixture
    assert curve_tree.waveform == wf
    assert curve_tree.color_palette == "plasma"
    assert curve_tree.tree.columnCount() == 8

    assert curve_tree.toolbar.components.exists("add")
    assert curve_tree.toolbar.components.exists("expand")
    assert curve_tree.toolbar.components.exists("collapse")
    assert curve_tree.toolbar.components.exists("renormalize_colors")


def test_add_new_curve(curve_tree_fixture):
    """
    Test that add_new_curve() adds a top-level item with a device curve config,
    assigns it a color from the buffer, and doesn't modify existing rows.
    """
    curve_tree, wf = curve_tree_fixture
    curve_tree.color_buffer = ["#111111", "#222222", "#333333", "#444444", "#555555"]

    assert curve_tree.tree.topLevelItemCount() == 0

    with patch.object(curve_tree, "_ensure_color_buffer_size") as ensure_spy:
        new_item = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
        ensure_spy.assert_called_once()

    assert curve_tree.tree.topLevelItemCount() == 1
    last_item = curve_tree.all_items[-1]
    assert last_item is new_item
    assert new_item.config.source == "device"
    assert new_item.config.signal.device == "bpm4i"
    assert new_item.config.signal.signal == "bpm4i"
    assert new_item.config.color in curve_tree.color_buffer


def test_renormalize_colors(curve_tree_fixture):
    """
    Test that renormalize_colors overwrites colors for all items in creation order.
    """
    curve_tree, wf = curve_tree_fixture
    # Add multiple curves
    c1 = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    c2 = curve_tree.add_new_curve(device="bpm3a", signal="bpm3a")
    curve_tree.color_buffer = []

    set_color_spy_c1 = patch.object(c1.color_button, "set_color")
    set_color_spy_c2 = patch.object(c2.color_button, "set_color")

    with set_color_spy_c1 as spy1, set_color_spy_c2 as spy2:
        curve_tree.renormalize_colors()
        spy1.assert_called_once()
        spy2.assert_called_once()


def test_expand_collapse(curve_tree_fixture):
    """
    Test expand_all_daps() and collapse_all_daps() calls expand/collapse on every top-level item.
    """
    curve_tree, wf = curve_tree_fixture
    c1 = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    curve_tree.tree.expandAll()
    expand_spy = patch.object(curve_tree.tree, "expandItem")
    collapse_spy = patch.object(curve_tree.tree, "collapseItem")

    with expand_spy as e_spy:
        curve_tree.expand_all_daps()
        e_spy.assert_called_once_with(c1)

    with collapse_spy as c_spy:
        curve_tree.collapse_all_daps()
        c_spy.assert_called_once_with(c1)


def test_send_curve_json(curve_tree_fixture, monkeypatch):
    """
    Test that send_curve_json sets the waveform's color_palette and curve_json
    to the exported config from the tree.
    """
    curve_tree, wf = curve_tree_fixture
    # Add multiple curves
    curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    curve_tree.add_new_curve(device="bpm3a", signal="bpm3a")

    curve_tree.color_palette = "viridis"
    curve_tree.send_curve_json()

    assert wf.color_palette == "viridis"
    data = json.loads(wf.curve_json)
    assert len(data) == 2
    labels = [d["label"] for d in data]
    assert "bpm4i-bpm4i" in labels
    assert "bpm3a-bpm3a" in labels


def test_refresh_from_waveform(qtbot, mocked_client_with_dap, monkeypatch):
    """
    Test that refresh_from_waveform() rebuilds the tree from the waveform's curve_json
    """
    patched_models = {"GaussianModel": {}, "LorentzModel": {}, "SineModel": {}}
    monkeypatch.setattr(mocked_client_with_dap.dap, "_available_dap_plugins", patched_models)

    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
    wf.x_mode = "auto"
    curve_tree = create_widget(qtbot, CurveTree, parent=None, waveform=wf)

    wf.plot(arg1="bpm4i", dap="GaussianModel")
    wf.plot(arg1="bpm3a", dap="GaussianModel")

    # Clear the tree to simulate a fresh rebuild.
    curve_tree.tree.clear()
    curve_tree.all_items.clear()
    assert curve_tree.tree.topLevelItemCount() == 0

    # For DAP rows
    curve_tree.refresh_from_waveform()
    assert curve_tree.tree.topLevelItemCount() == 2


def test_refresh_from_waveform_preserves_linear_model_toggle(
    qtbot, mocked_client_with_dap, monkeypatch
):
    """
    Test that a composite DAP model with LinearModel is restored into the curve-tree UI.
    """
    patched_models = {"GaussianModel": {}, "LorentzModel": {}, "SineModel": {}, "LinearModel": {}}
    monkeypatch.setattr(mocked_client_with_dap.dap, "_available_dap_plugins", patched_models)

    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
    curve_tree = create_widget(qtbot, CurveTree, parent=None, waveform=wf)

    curve = wf.plot(arg1="bpm4i")
    wf.add_dap_curve(device_label=curve.name(), dap_name=["GaussianModel", "LinearModel"])

    curve_tree.refresh_from_waveform()

    device_row = curve_tree.tree.topLevelItem(0)
    dap_child = device_row.child(0)
    assert dap_child.dap_combo.fit_model_combobox.currentText() == "GaussianModel"
    assert dap_child.composite_models == ["LinearModel"]
    assert dap_child.composite_button.text() == "Composite (2)"


def test_add_dap_row(curve_tree_fixture):
    """
    Test that add_dap_row creates a new DAP curve as a child of a device curve,
    with the correct configuration and parent-child relationship.
    """
    curve_tree, wf = curve_tree_fixture

    # Add a device curve first
    device_row = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    assert device_row.source == "device"
    assert curve_tree.tree.topLevelItemCount() == 1
    assert device_row.childCount() == 0

    # Now add a DAP row to it
    device_row.add_dap_row()

    # Check that child was added
    assert device_row.childCount() == 1
    dap_child = device_row.child(0)

    # Verify the DAP child has the correct configuration
    assert dap_child.source == "dap"
    assert dap_child.config.parent_label == device_row.config.label

    # Check that the DAP inherits device name/entry from parent
    assert dap_child.config.signal.device == "bpm4i"
    assert dap_child.config.signal.signal == "bpm4i"

    # Check that the item is in the curve_tree's all_items list
    assert dap_child in curve_tree.all_items


def test_remove_self_top_level(curve_tree_fixture):
    """
    Test that remove_self removes a top-level device row from the tree.
    """
    curve_tree, wf = curve_tree_fixture

    # Add two device curves
    row1 = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    row2 = curve_tree.add_new_curve(device="bpm3a", signal="bpm3a")
    assert curve_tree.tree.topLevelItemCount() == 2
    assert len(curve_tree.all_items) == 2

    # Remove the first row
    row1.remove_self()

    # Check that only one row remains and it's the correct one
    assert curve_tree.tree.topLevelItemCount() == 1
    assert curve_tree.tree.topLevelItem(0) == row2
    assert len(curve_tree.all_items) == 1
    assert curve_tree.all_items[0] == row2


def test_remove_self_child(curve_tree_fixture):
    """
    Test that remove_self removes a child DAP row while preserving the parent device row.
    """
    curve_tree, wf = curve_tree_fixture

    # Add a device curve and a DAP child
    device_row = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    device_row.add_dap_row()
    dap_child = device_row.child(0)

    assert curve_tree.tree.topLevelItemCount() == 1
    assert device_row.childCount() == 1
    assert len(curve_tree.all_items) == 2

    # Remove the DAP child
    dap_child.remove_self()

    # Check that the parent device row still exists but has no children
    assert curve_tree.tree.topLevelItemCount() == 1
    assert device_row.childCount() == 0
    assert len(curve_tree.all_items) == 1
    assert curve_tree.all_items[0] == device_row


def test_export_data_dap(curve_tree_fixture):
    """
    Test that export_data from a DAP row correctly includes parent relationship and DAP model.
    """
    curve_tree, wf = curve_tree_fixture

    # Add a device curve with specific parameters
    device_row = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")

    # Add a DAP child
    device_row.add_dap_row()
    dap_child = device_row.child(0)

    # Set a specific model in the DAP combobox
    dap_child.dap_combo.fit_model_combobox.setCurrentText("GaussianModel")

    # Export data from the DAP row
    exported = dap_child.export_data()

    # Check the exported data
    assert exported["source"] == "dap"
    assert exported["parent_label"] == "bpm4i-bpm4i"
    assert exported["signal"]["device"] == "bpm4i"
    assert exported["signal"]["signal"] == "bpm4i"
    assert exported["signal"]["dap"] == "GaussianModel"
    assert exported["label"] == "bpm4i-bpm4i-GaussianModel"


def test_export_data_dap_with_linear_model(curve_tree_fixture):
    """
    Test that export_data serializes the optional LinearModel composite selection.
    """
    curve_tree, wf = curve_tree_fixture

    device_row = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    device_row.add_dap_row()
    dap_child = device_row.child(0)

    dap_child.dap_combo.fit_model_combobox.setCurrentText("GaussianModel")
    dap_child.composite_models = ["LinearModel"]
    dap_child._update_composite_button_text()

    exported = dap_child.export_data()

    assert exported["signal"]["dap"] == ["GaussianModel", "LinearModel"]
    assert exported["label"] == "bpm4i-bpm4i-GaussianModel+LinearModel"
    assert dap_child.composite_button.text() == "Composite (2)"


def test_primary_model_change_removes_duplicate_composite_model(curve_tree_fixture):
    """
    Test that changing the primary model removes the same model from the composite extras.
    """
    curve_tree, wf = curve_tree_fixture

    device_row = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    device_row.add_dap_row()
    dap_child = device_row.child(0)

    dap_child.composite_models = ["LorentzModel", "LinearModel"]
    dap_child.dap_combo.fit_model_combobox.setCurrentText("LorentzModel")

    assert dap_child.composite_models == ["LinearModel"]
    assert dap_child.composite_button.text() == "Composite (2)"


def test_export_data_history_curve(curve_tree_fixture, scan_history_factory):
    """
    Test that export_data for a history curve row correctly serializes scan_number
    and scan_id when a numeric scan is typed into the scan index combobox.
    """
    curve_tree, wf = curve_tree_fixture
    # Inject two history scans into the waveform client
    msgs = [
        scan_history_factory(scan_id="hid1", scan_number=1),
        scan_history_factory(scan_id="hid2", scan_number=2),
    ]
    wf.client.history = ScanHistory(wf.client, False)
    for m in msgs:
        wf.client.history._scan_data[m.scan_id] = m
        wf.client.history._scan_ids.append(m.scan_id)
        wf.client.history._scan_numbers.append(m.scan_number)
    wf.client.queue.scan_storage.current_scan = None

    # Create a device row and select scan index "2"
    device_row = curve_tree.add_new_curve(device="bpm4i", signal="bpm4i")
    device_row.scan_index_combo.setCurrentText("2")

    exported = device_row.export_data()
    assert exported["source"] == "history"
    assert exported["scan_number"] == 2
    assert exported["scan_id"] == "hid2"
    assert exported["label"] == "bpm4i-bpm4i-scan-2"


##################################################
# Staging of edits (issue #1314)
##################################################


@pytest.fixture
def history_curve_setting(qtbot, mocked_client, scan_history_factory):
    """
    A Waveform showing one history curve of scan 1 (two scans in history) and a CurveSetting
    targeting it.
    """
    wf = create_widget(qtbot, Waveform, client=mocked_client)
    history_msgs = inject_scan_history(wf, scan_history_factory, ("hist1", 1), ("hist2", 2))
    # The scan index combobox lists scans by number, so register those as well.
    for msg in history_msgs:
        wf.client.history._scan_numbers.append(msg.scan_number)
    wf.x_mode = "index"
    history_curve = wf.plot(device_y="bpm4i", signal_y="bpm4i", scan_id="hist1")
    curve_setting = create_widget(qtbot, CurveSetting, parent=None, target_widget=wf)
    curve_setting.mode_combo.setCurrentText("index")
    return curve_setting, wf, history_curve


def test_switch_history_curve_to_live_keeps_history_refresh_working(history_curve_setting):
    """
    Regression for #1314: switching a history curve to 'live' and applying must not leave a
    device-source curve behind in the Waveform's history tracking. Previously the settings row
    edited the live curve's config in place, so the removed history curve turned into a
    device curve that the next history refresh tried to fill, raising
    "Source device do not allow custom data setting."
    """
    curve_setting, wf, history_curve = history_curve_setting
    row = curve_setting.curve_manager.tree.topLevelItem(0)
    assert row.scan_index_combo.scan_id == "hist1"

    row.scan_index_combo.set_scan_id(None)  # select 'live'
    curve_setting.accept_changes()

    # The next history refresh (x-axis change, Apply again, ...) must not raise.
    wf._refresh_history_curves()

    assert [curve.config.source for curve in wf.curves] == ["device"]
    assert wf._history_curves == []
    # The removed history curve kept its own configuration.
    assert history_curve.config.source == "history"
    assert history_curve.config.scan_id == "hist1"


def test_curve_settings_edits_do_not_touch_live_curve_config(history_curve_setting):
    """
    Edits in the curve tree are staged on a copy of each curve's config: changing the color,
    renormalizing colors or exporting the rows must leave the plotted curve untouched until
    the Waveform applies the exported configuration (also #1304, Cancel keeps color edits).
    """
    curve_setting, wf, history_curve = history_curve_setting
    curve_manager = curve_setting.curve_manager
    original_config = history_curve.config.model_dump()

    row = curve_manager.tree.topLevelItem(0)
    curve_manager.renormalize_colors()
    row.color_button.set_color("#123456")
    row.scan_index_combo.set_scan_id(None)  # select 'live'
    curve_manager.export_all_curves()

    assert history_curve.config.model_dump() == original_config


##################################################
# Apply keeps custom curves and DAP settings (issues #1302, #1303)
##################################################


def test_apply_keeps_custom_curves(qtbot, mocked_client_with_dap):
    """
    Regression for #1302: custom curves are not listed in the curve tree, but applying the
    curve settings must keep them with their configuration, their data and their DAP fit.
    """
    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
    wf.plot(arg1="bpm4i")
    custom = wf.plot(
        x=[1, 2, 3],
        y=[4, 5, 6],
        label="custom",
        color="#ff0000",
        pen_width=2,
        pen_style="dash",
        dap="GaussianModel",
    )
    custom_config = custom.config.model_dump()
    curve_setting = create_widget(qtbot, CurveSetting, parent=None, target_widget=wf)

    curve_setting.accept_changes()

    assert custom in wf.curves
    assert custom.config.model_dump() == custom_config
    x_data, y_data = custom.get_data()
    assert list(x_data) == [1, 2, 3]
    assert list(y_data) == [4, 5, 6]
    assert sorted(c.name() for c in wf.curves) == ["bpm4i-bpm4i", "custom", "custom-GaussianModel"]
    custom_dap = wf.get_curve("custom-GaussianModel")
    assert custom_dap.config.parent_label == "custom"
    assert custom_dap in wf._dap_curves


def test_apply_keeps_dap_oversample_and_parameters(qtbot, mocked_client_with_dap):
    """
    Regression for #1303: applying an unrelated change in the curve settings must keep the
    oversampling factor and the fit parameter overrides of a DAP curve.
    """
    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
    wf.plot(arg1="bpm4i")
    wf.add_dap_curve(
        device_label="bpm4i-bpm4i",
        dap_name="GaussianModel",
        dap_oversample=3,
        dap_parameters={"amplitude": 1.0},
    )
    dap_parameters = wf.get_curve("bpm4i-bpm4i-GaussianModel").config.signal.dap_parameters
    assert dap_parameters
    curve_setting = create_widget(qtbot, CurveSetting, parent=None, target_widget=wf)

    dap_row = curve_setting.curve_manager.tree.topLevelItem(0).child(0)
    dap_row.color_button.set_color("#123456")  # an unrelated change
    curve_setting.accept_changes()

    dap_curve = wf.get_curve("bpm4i-bpm4i-GaussianModel")
    assert dap_curve.config.color == "#123456"
    assert dap_curve.dap_oversample == 3
    assert dap_curve.config.signal.dap_parameters == dap_parameters


def test_dap_model_change_drops_parameters_but_keeps_oversample(qtbot, mocked_client_with_dap):
    """
    Fit parameter overrides belong to the selected model, so choosing another model drops them;
    the oversampling factor does not depend on the model and is kept.
    """
    wf = create_widget(qtbot, Waveform, client=mocked_client_with_dap)
    wf.plot(arg1="bpm4i")
    wf.add_dap_curve(
        device_label="bpm4i-bpm4i",
        dap_name="GaussianModel",
        dap_oversample=3,
        dap_parameters={"amplitude": 1.0},
    )
    curve_tree = create_widget(qtbot, CurveTree, parent=None, waveform=wf)
    dap_row = curve_tree.tree.topLevelItem(0).child(0)

    dap_row.dap_combo.fit_model_combobox.setCurrentText("LorentzModel")
    exported = dap_row.export_data()

    assert exported["signal"]["dap"] == "LorentzModel"
    assert exported["signal"]["dap_oversample"] == 3
    assert exported["signal"]["dap_parameters"] is None