AareDAQ/gui/src/aaregui/panels/sample_queue_panel.py

from PySide6.QtCore import Signal, Slot, Qt
from PySide6.QtWidgets import QFrame, QVBoxLayout, QHBoxLayout, QPushButton, QHeaderView, QSizePolicy, \
    QTableView, QMessageBox
from PySide6.QtGui import QKeySequence, QShortcut
from aaredaqlib.models import SampleShortInfoList, SampleShortInfo

from aaregui.models.sample_queue_model import SampleQueueSpreadsheet
from aaregui.widgets.title_label import TitleLabel
from aaredaqlib.models import DAQStatusModel

from aaredaqlib.logger_config import setup_logger

logger = setup_logger("aareGUI")

class SampleQueuePanel(QFrame):
    auto_scan = Signal(SampleShortInfo)
    unmount = Signal()
    viewer_track_online = Signal()

    def __init__(self, parent=None, samples: SampleShortInfoList | None = None):
        super().__init__(parent)
        self.__pause = True
        self.__set_to_pause = False
        self._current_db_id: int | None = None
        self.ring_current = None

        self.setFrameShape(QFrame.Shape.StyledPanel)
        self.setFrameShadow(QFrame.Shadow.Raised)

        layout = QVBoxLayout(self)
        self.setLayout(layout)

        layout.addWidget(TitleLabel("Sample queue", self))
        # Create the table
        self.table_view = QTableView(self)

        # Create the custom table model
        self.table_model = SampleQueueSpreadsheet()
        # Create the QTableView and set the model
        self.table_view.setModel(self.table_model)
        # Enable automatic column resizing
        header = self.table_view.horizontalHeader()
        header.setSectionResizeMode(QHeaderView.ResizeMode.Stretch)  # Columns stretch to fill the table

        # Set the table's size policy to expand
        self.table_view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        self.table_view.setAcceptDrops(True)
        self.table_view.setDropIndicatorShown(True)

        # Enable selection of multiple rows
        self.table_view.setSelectionBehavior(QTableView.SelectionBehavior.SelectRows)

        # Set up keyboard shortcut for delete
        self.delete_shortcut = QShortcut(QKeySequence.StandardKey.Delete, self.table_view)
        self.delete_shortcut.activated.connect(self.remove_selected_samples)

        layout.addWidget(self.table_view)  # Add table to the layout

        # Create a horizontal layout for the buttons
        button_layout = QHBoxLayout()
        self.play_button = QPushButton("▶ Run", self)
        self.play_button.clicked.connect(self.run)

        self.remove_button = QPushButton("🗑 Remove selected", self)
        self.remove_button.clicked.connect(self.remove_selected_samples)

        self.clear_button = QPushButton("✖ Clear list", self)
        self.clear_button.clicked.connect(self.clear)

        self.clear_button.clicked.connect(self.table_model.clearSamples)

        button_layout.addWidget(self.play_button)
        button_layout.addWidget(self.remove_button)
        button_layout.addWidget(self.clear_button)

        # Add the button layout to the main layout
        layout.addLayout(button_layout)

    def remove_selected_samples(self):
        """Remove selected samples from the queue."""
        selected_indexes = self.table_view.selectionModel().selectedRows()
        if not selected_indexes:
            return

        # Get the row numbers and sort them in descending order
        # This ensures we remove from bottom to top to avoid index shifting
        rows = sorted([index.row() for index in selected_indexes], reverse=True)

        # Remove samples by their database IDs
        for row in rows:
            if 0 <= row < len(self.table_model.samples):
                sample = self.table_model.samples[row]
                self.table_model.remove_sample(sample.db_id)

    def __ring_current_low_check(self) -> bool:
        if self.ring_current is not None and self.ring_current < 100:
            logger.debug(f"Ring current too low {self.ring_current}")
            self.table_model.set_running(False)
            self.set_to_pause = True
            self.__pause = True
            self.play_button.setText("▶ Run")
            reply = QMessageBox.question(self, "Ring current too low", f"Ring current is too low: {self.ring_current} mA. Do you wish to continue?",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No)
            if reply == QMessageBox.StandardButton.Yes:
                return True
            else:
                return False
        logger.debug(f"Ring current: {self.ring_current}")
        return True

    def run(self):
        if not self.__ring_current_low_check():
            logger.debug("low ring current, skipping")
            return

        elif self.__pause:
            if len(self.table_model.samples) > 0:
                self.table_model.set_running(True)
                self.__set_to_pause = False
                self.__pause = False
                self.play_button.setText("⏸ Pause")
                current = self.table_model.samples[0]
                self._current_db_id = current.db_id
                self.auto_scan.emit(current)
                self.viewer_track_online.emit()
        else:
            self.__set_to_pause = True
            self.__pause = True
            self.play_button.setText("▶ Run")

    def clear(self):
        self.table_model.clearSamples()

    @Slot(DAQStatusModel)
    def update_daq_status(self, s: DAQStatusModel):
        self.ring_current = s.bl.ring_current_mA

    @Slot(int, bool)
    def automated_scan_done(self, db_id: int, success: bool):
        if self._current_db_id is not None and db_id != self._current_db_id:
            return

        if not self.__ring_current_low_check():
            self.unmount.emit()
            return

        if success:
            self.table_model.remove_sample(db_id)
            self._current_db_id = None
            if not self.__pause and len(self.table_model.samples) > 0:
                next_item = self.table_model.samples[0]
                self._current_db_id = next_item.db_id
                self.auto_scan.emit(next_item)
            else:
                self.table_model.set_running(False)
                self.__pause = True
                self.play_button.setText("▶ Run")
                self.unmount.emit()
        else:
            self.table_model.set_running(False)
            self.__pause = True
            self.play_button.setText("▶ Run")
            self._current_db_id = None