AareDAQ/src/aare/gui/panels/reference_tools_panel.py

# reference_tools_panel.py

from aarecommon.config.logger import setup_logger
from aarecommon.models.models import DAQStatusModel, SampleShortInfo, SampleShortInfoList
from PySide6.QtCore import QAbstractTableModel, QModelIndex, Qt, Signal, Slot
from PySide6.QtGui import QBrush
from PySide6.QtWidgets import QAbstractItemView, QFrame, QGridLayout, QHeaderView, QMenu, QTableView

from aare.gui.constants import LOGGER_NAME
from aare.gui.styles import SAMPLE_ROW_QUEUED_BG, SAMPLE_STATUS_TEXT, qcolor
from aare.gui.widgets.title_label import TitleLabel

logger = setup_logger(LOGGER_NAME)


def get_entry(sample: SampleShortInfo, column: int):
    if column == 0:
        return sample.loc_str()
    elif column == 1:
        return sample.sample_name
    elif column == 2:
        return sample.mount_count
    elif column == 3:
        return sample.rotation_count
    elif column == 4:
        return sample.raster_count
    elif column == 5:
        return sample.screening_count
    return ""


class ReferenceToolsModel(QAbstractTableModel):
    def __init__(
        self,
        rows: list[SampleShortInfo] | None = None,
        parent=None,
        current_reference: int | None = None,
    ):
        super().__init__(parent)

        self.samples: list[SampleShortInfo] = rows or []
        self.current_reference = current_reference
        # Column 0 is display-only: the row position ("#"), matching the
        # Dewar samples table; the vertical header is hidden in the panel.
        self.header = [
            "#",
            "Position",
            "Sample name",
            "Mount count",
            "Raster count",
            "Rotation count",
            "Screening count",
        ]
        self._sort_col = 1
        self._sort_order = Qt.SortOrder.AscendingOrder
        self._sorted_samples: list[SampleShortInfo] = []
        if self.samples:
            self._sort()
        else:
            self._sorted_samples = []

    def rowCount(self, parent: QModelIndex | None = None) -> int:
        return len(self._sorted_samples)

    def columnCount(self, parent: QModelIndex | None = None) -> int:
        return len(self.header)

    def data(self, index: QModelIndex, role=None):
        if not index.isValid():
            return None

        if index.row() >= len(self._sorted_samples):
            return None

        if role == Qt.ItemDataRole.DisplayRole:
            if index.column() == 0:
                return str(index.row() + 1)
            return get_entry(self._sorted_samples[index.row()], index.column() - 1)
        elif role == Qt.ItemDataRole.TextAlignmentRole:
            return Qt.AlignmentFlag.AlignCenter
        elif role == Qt.ItemDataRole.BackgroundRole:
            # Tint only the current-reference row; plain rows return None so
            # the theme QSS paints them (a hardcoded WHITE fill here was the
            # big white table in dark mode and fought the light theme's
            # alternating stripes).
            if self._sorted_samples[index.row()].db_id == self.current_reference:
                return QBrush(qcolor(SAMPLE_ROW_QUEUED_BG))
        elif (
            role == Qt.ItemDataRole.ForegroundRole
            and self._sorted_samples[index.row()].db_id == self.current_reference
        ):
            # Fixed dark ink on the tint — the tint stays pale in BOTH themes,
            # so the dark theme's near-white text would vanish on it.
            return QBrush(qcolor(SAMPLE_STATUS_TEXT))

        return None

    def headerData(self, section: int, orientation, role=None):
        if role != Qt.ItemDataRole.DisplayRole:
            return None
        if orientation == Qt.Orientation.Horizontal:
            return self.header[section] if section < len(self.header) else f"Column {section + 1}"
        return str(section + 1)

    def update_rows(self, rows: list[SampleShortInfo]):
        self.beginResetModel()
        try:
            self.samples = rows
            if self.samples:  # Only sort if we have samples
                self._sort()
            else:
                self._sorted_samples = []  # Clear sorted samples if no data
        finally:
            self.endResetModel()

    def sort(self, column, order):
        # The "#" column is display-only — nothing to sort by.
        if column == 0:
            return
        self.layoutAboutToBeChanged.emit()
        self._sort_order = order
        self._sort_col = column
        self._sort()
        self.layoutChanged.emit()

    def _sort(self):
        """Sort samples, using special handling for location column."""
        if not self.samples:
            self._sorted_samples = []
            return

        if self._sort_col == 1:
            # Special sorting for location (Position column)
            self._sorted_samples = sorted(
                self.samples,
                key=lambda row: row.loc_str_sort(),
                reverse=(self._sort_order == Qt.SortOrder.DescendingOrder),
            )
        elif self._sort_col == 3:
            # Numeric sort for Mount count; place None last on ascending, first on descending
            none_sentinel = (
                float("inf") if self._sort_order == Qt.SortOrder.AscendingOrder else float("-inf")
            )
            self._sorted_samples = sorted(
                self.samples,
                key=lambda row: (
                    row.mount_count if isinstance(row.mount_count, (int, float)) else none_sentinel
                ),
                reverse=(self._sort_order == Qt.SortOrder.DescendingOrder),
            )
        else:
            # String sort with empty fallback (get_entry columns sit one left
            # of the view columns because of the display-only "#").
            self._sorted_samples = sorted(
                self.samples,
                key=lambda row: get_entry(row, self._sort_col - 1) or "",
                reverse=(self._sort_order == Qt.SortOrder.DescendingOrder),
            )

    def get_item(self, row: int) -> SampleShortInfo | None:
        """Get sample at the given row index."""
        if 0 <= row < len(self._sorted_samples):
            return self._sorted_samples[row]
        return None

    def update_current_reference(self, reference_id: int | None):
        """Update which sample is highlighted as current reference."""
        self.current_reference = reference_id
        # Trigger repaint of background colors
        if self.rowCount() > 0:
            self.dataChanged.emit(
                self.index(0, 0),
                self.index(self.rowCount() - 1, self.columnCount() - 1),
                [Qt.ItemDataRole.BackgroundRole, Qt.ItemDataRole.ForegroundRole],
            )


class ReferenceToolsPanel(QFrame):
    mount = Signal(SampleShortInfo, bool)
    unmount = Signal()

    def __init__(
        self,
        samples: SampleShortInfoList | None = None,
        parent=None,
        refresh_interval_ms: int = 5000,
        model: ReferenceToolsModel | None = None,
    ):
        """
        :param samples: optional initial SampleShortInfoList to populate the table
        :param parent: Qt parent
        :param refresh_interval_ms: how often to call request_refresh (panel doesn't implement the request itself)
        :param model: share an existing model instead of owning one (pop-out window)
        """
        super().__init__(parent)

        if samples is None:
            samples = SampleShortInfoList(s=[])

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

        layout = QGridLayout(self)
        self.setLayout(layout)
        # Flush layout, matching TellSamplePanel: full-width banner, no
        # padding ring, no gap to the dock's tab row.
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)

        layout.addWidget(TitleLabel("Reference tools", parent=self), 0, 0, 1, 4)

        self.table_view = QTableView(parent=self)
        # Row colors carry the separation — no grid lines. Alternating rows
        # come from the theme QSS (alternate-background-color), same as the
        # Dewar sample list.
        self.table_view.setShowGrid(False)
        self.table_view.setAlternatingRowColors(True)
        layout.addWidget(self.table_view, 1, 0, 1, 4)

        # initialize model with provided samples (or adopt the shared one)
        self.table_model = model if model is not None else ReferenceToolsModel(rows=samples.s)
        self.table_view.setModel(self.table_model)

        self.table_view.setEditTriggers(QTableView.EditTrigger.NoEditTriggers)
        logger.debug("Setting up table header")
        header = self.table_view.horizontalHeader()
        logger.debug("Setting up table header")
        header.setSectionResizeMode(QHeaderView.ResizeMode.Interactive)
        logger.debug("Setting up table header")
        header.setStretchLastSection(True)
        # No bold column titles when cells are selected.
        header.setHighlightSections(False)
        # Row numbers live in the display-only "#" column (like the Dewar
        # table), not the vertical header.
        self.table_view.verticalHeader().setVisible(False)
        logger.debug("Setting up table view sorting")
        # Adopt the model's current order first — a second panel on a shared
        # model must not re-sort it on open.
        header.setSortIndicator(self.table_model._sort_col, self.table_model._sort_order)
        self.table_view.setSortingEnabled(True)
        logger.debug("Setting up table view context menu")
        self.table_view.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.table_view.customContextMenuRequested.connect(self._context_menu)
        self.table_view.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows)
        self.table_view.setSelectionMode(QTableView.SelectionMode.SingleSelection)

        # Columns at full content width, sized once (see TellSamplePanel).
        self._columns_autosized = False
        if self.table_model.rowCount() > 0:
            self.table_view.resizeColumnsToContents()
            self._columns_autosized = True

    def _selected_item(self) -> SampleShortInfo | None:
        idx = self.table_view.currentIndex()
        if not idx.isValid():
            return None
        return self.table_model.get_item(idx.row())

    def _context_menu(self, position):
        idx = self.table_view.indexAt(position)

        if not idx.isValid():
            return

        item = self.table_model.get_item(idx.row())

        if item is None:
            return

        menu = QMenu()
        act_mount = menu.addAction("Mount")

        act_unmount = menu.addAction("Unmount")

        action = menu.exec_(self.table_view.viewport().mapToGlobal(position))

        if action == act_mount:
            self.mount.emit(item, True)
        elif action == act_unmount:
            self.unmount.emit()

    @Slot(SampleShortInfoList)
    def new_list(self, samples: SampleShortInfoList):
        # signal from DAQWorker will call this with the model
        self.table_model.update_rows(rows=samples.s)
        if not self._columns_autosized and self.table_model.rowCount() > 0:
            self.table_view.resizeColumnsToContents()
            self._columns_autosized = True

    @Slot(DAQStatusModel)
    def update_daq_status(self, status: DAQStatusModel):
        # The "No sample mounted" label and Unmount button were dropped —
        # mounted state now shows as the existing row highlight instead.
        sample = status.sample
        self.table_model.update_current_reference(sample.db_id if sample is not None else None)