import re from aarecommon.models.models import SampleShortInfo, SampleShortInfoList from PySide6.QtCore import QAbstractTableModel, QMimeData, Qt from PySide6.QtGui import QBrush, QColor def get_entry(sample: SampleShortInfo, column: int): if column == 0: return sample.sample_name elif column == 1: return sample.puck_name elif column == 2: return sample.dewar_name elif column == 3: return sample.loc_str() elif column == 4: return sample.priority elif column == 5: return sample.user elif column == 6: return sample.mount_count elif column == 7: return sample.raster_count elif column == 8: return sample.rotation_count elif column == 9: return sample.screening_count elif column == 10: return sample.comment class UserSampleSpreadsheet(QAbstractTableModel): def __init__( self, parent=None, samples: list[SampleShortInfo] | None = None, current_puck: str | None = None, current_sample: int | None = None, ): super().__init__(parent) if samples is None: samples = [] self.samples: list[SampleShortInfo] = samples self.header = [ "Sample name", "Puck", "Dewar", "Location", "Priority", "User", "Mount count", "Raster count", "Rotation count", "Screening count", "Comment", ] self.current_sample = current_sample self.current_puck = current_puck self.__sort_col = 3 self.__sort_order = Qt.SortOrder.AscendingOrder self.__filters: dict[int, str] = {} self.__filter_col: int | None = 5 self.__filter_value: str | None = None self.current_pgroup: str | None = None self.show_all_pgroups: bool = False self._sort() def to_list(self) -> list[dict]: return [sample.to_dict() for sample in self.samples] def from_list(self, data: list[dict]): self.beginResetModel() self.samples = [SampleShortInfo.from_dict(d) for d in data] self.endResetModel() def rowCount(self, parent=None): return len(self.__sorted_samples) def columnCount(self, parent=None): return len(self.header) def data(self, index, role=None): if role == Qt.ItemDataRole.DisplayRole: return get_entry(self.__sorted_samples[index.row()], index.column()) elif role == Qt.ItemDataRole.TextAlignmentRole: # Align text to center return Qt.AlignmentFlag.AlignCenter elif role == Qt.ItemDataRole.BackgroundRole: if self.__sorted_samples[index.row()].db_id == self.current_sample: return QBrush(QColor(114, 159, 207)) # darker blue if self.__sorted_samples[index.row()].puck_name == self.current_puck: return QBrush(QColor(216, 228, 253)) # light blue return QBrush(QColor(255, 255, 255)) # White return None # For other roles, return None def headerData(self, section, orientation, role=None): if role == Qt.ItemDataRole.DisplayRole: if orientation == Qt.Orientation.Horizontal: # Column header return self.header[section] if self.header else f"Column {section + 1}" if orientation == Qt.Orientation.Vertical: # Row header return str(section + 1) # Row numbers start from 1 return None def updateCurrentSample( self, current_puck: str | None = None, current_sample: int | None = None ): self.current_puck = current_puck self.current_sample = current_sample def updateData(self, samples: list[SampleShortInfo]): if samples != self.samples: self.beginResetModel() self.samples = samples self._sort() self.endResetModel() def sort(self, column, order): self.layoutAboutToBeChanged.emit() self.__sort_order = order self.__sort_col = column self._sort() self.layoutChanged.emit() def _sort(self): filtered = self._apply_filter(self.samples) if self.__sort_col == 3: self.__sorted_samples = sorted( filtered, key=lambda row: row.loc_str_sort(), reverse=( self.__sort_order == Qt.SortOrder.DescendingOrder ), # Reverse for descending order ) else: # Sort the samples based on the specified column and order self.__sorted_samples = sorted( filtered, key=lambda row: get_entry(row, self.__sort_col), reverse=( self.__sort_order == Qt.SortOrder.DescendingOrder ), # Reverse for descending order ) def _apply_filter(self, rows: list[SampleShortInfo]) -> list[SampleShortInfo]: # Default filter by User using current p-group if no explicit filter set filters: dict[int, str] = { col: v for col, v in (self.__filters or {}).items() if (v or "").strip() } if self.__filter_col is not None and (self.__filter_value or "").strip(): filters[self.__filter_col] = self.__filter_value if 5 not in filters and self.current_pgroup and not self.show_all_pgroups: filters[5] = self.current_pgroup if not filters: return rows def matches(sample: SampleShortInfo) -> bool: for col, txt in filters.items(): val = get_entry(sample, col) s = "" if val is None else str(val) if str(txt).strip().lower() not in s.strip().lower(): return False return True return [r for r in rows if matches(r)] def flags(self, index): default_flags = super().flags(index) if index.isValid(): return Qt.ItemFlag.ItemIsDragEnabled | default_flags return default_flags def mimeTypes(self): return ["text/plain"] def mimeData(self, indexes): mime_data = QMimeData() sample_data = SampleShortInfoList(s=[]) for i in sorted(set(index.row() for index in indexes)): sample_data.s.append(self.__sorted_samples[i]) mime_data.setText(sample_data.model_dump_json()) return mime_data def get_id(self, row: int) -> SampleShortInfo: return self.__sorted_samples[row] def set_filter(self, field: str, text: str | None): try: col = self.header.index(field) except ValueError: col = None self.layoutAboutToBeChanged.emit() self.__filter_col = col if text and text.strip() and col is not None else None self.__filter_value = text.strip() if text else None # mirror into multi-filter map if col is not None: if text and text.strip(): self.__filters[col] = text.strip() else: self.__filters.pop(col, None) self._sort() self.layoutChanged.emit() def clear_filter(self): self.layoutAboutToBeChanged.emit() self.__filter_col = None self.__filter_value = None self.__filters.clear() self._sort() self.layoutChanged.emit() def set_default_user_filter(self, pgroup: str | None): self.layoutAboutToBeChanged.emit() self.current_pgroup = (pgroup or "").strip() or None self._sort() self.layoutChanged.emit() def set_column_filter(self, column: int, text: str | None): self.layoutAboutToBeChanged.emit() if text and text.strip(): self.__filters[column] = text.strip() else: self.__filters.pop(column, None) self._sort() self.layoutChanged.emit() def clear_all_column_filters(self): self.clear_filter() def set_show_all_pgroups(self, show_all: bool): self.layoutAboutToBeChanged.emit() self.show_all_pgroups = bool(show_all) self._sort() self.layoutChanged.emit() def unique_values_for_column(self, column: int, limit: int = 200) -> list[str]: """Get unique values for a column from currently filtered samples (excluding this column's filter).""" # Get currently filtered samples, but exclude the filter for this column temp_filter = self.__filters.pop(column, None) filtered_samples = self._apply_filter(self.samples) # Restore the filter if temp_filter is not None: self.__filters[column] = temp_filter seen: set[str] = set() out: list[str] = [] for s in filtered_samples: v = get_entry(s, column) if v is None: continue txt = str(v).strip() if not txt: continue if txt not in seen: seen.add(txt) out.append(txt) if len(out) >= limit: break # Sort appropriately if column == 5: # User/pgroup column try: out.sort( key=lambda x: ( int(x[1:]) if x and x[0].lower() == "p" and x[1:].isdigit() else float("inf") ) ) except Exception: out.sort() else: out.sort() return out def suggested_prefixes_for_sample_name(self, limit: int = 200) -> list[str]: """Get sample name prefixes from currently filtered samples (excluding column 0 filter).""" # Get currently filtered samples, excluding the sample name filter temp_filter = self.__filters.pop(0, None) filtered_samples = self._apply_filter(self.samples) if temp_filter is not None: self.__filters[0] = temp_filter rx = re.compile(r"^([A-Za-z]+)") counts: dict[str, int] = {} for s in filtered_samples: name = (s.sample_name or "").strip() if not name: continue m = rx.match(name) pref = m.group(1) if m else name.split("_")[0] pref = pref.strip() if not pref: continue counts[pref] = counts.get(pref, 0) + 1 # Sort by frequency desc, then alpha items = sorted(counts.items(), key=lambda kv: (-kv[1], kv[0])) return [k for k, _ in items[:limit]] def suggested_prefixes_for_location(self, limit: int = 200) -> tuple[list[str], list[str]]: """Get location prefixes from currently filtered samples (excluding column 3 filter).""" # Get currently filtered samples, excluding the location filter temp_filter = self.__filters.pop(3, None) filtered_samples = self._apply_filter(self.samples) if temp_filter is not None: self.__filters[3] = temp_filter seg_seen: set[str] = set() segpos_seen: set[str] = set() for s in filtered_samples: loc = s.loc_str() # "-" or like "B3-10" if loc and loc != "-" and isinstance(loc, str) and "-" in loc: left = loc.split("-", 1)[0] # e.g., "B3" if left: segpos_seen.add(left) seg = left[0] if seg: seg_seen.add(seg) segs = sorted(seg_seen) segpos = sorted(segpos_seen, key=lambda x: (x[0], int(x[1:]) if x[1:].isdigit() else 0)) return (segs[:limit], segpos[:limit])