1import itertools 2import random 3 4from PyQt5.QtCore import QAbstractTableModel, pyqtSignal, QModelIndex, Qt, QItemSelection 5 6from urh import settings 7from urh.signalprocessing.Participant import Participant 8 9 10class ParticipantTableModel(QAbstractTableModel): 11 INITIAL_NAMES = ["Alice", "Bob", "Carl", "Dave", "Eve", "Frank", "Grace", "Heidi", "Judy", "Mallory", "Oscar", 12 "Peggy", "Sybil", "Trudy", "Victor", "Walter"] 13 14 updated = pyqtSignal() 15 participant_edited = pyqtSignal() 16 17 def __init__(self, participants): 18 super().__init__() 19 self.participants = participants 20 self.header_labels = ["Name", "Shortname", "Color", "Relative RSSI", "Address (hex)"] 21 22 def update(self): 23 self.beginResetModel() 24 self.endResetModel() 25 self.updated.emit() 26 27 def columnCount(self, parent: QModelIndex = None, *args, **kwargs): 28 return len(self.header_labels) 29 30 def rowCount(self, parent: QModelIndex = None, *args, **kwargs): 31 return len(self.participants) 32 33 def headerData(self, section, orientation, role=Qt.DisplayRole): 34 if role == Qt.DisplayRole and orientation == Qt.Horizontal: 35 return self.header_labels[section] 36 return super().headerData(section, orientation, role) 37 38 def data(self, index: QModelIndex, role=Qt.DisplayRole): 39 if role == Qt.DisplayRole or role == Qt.EditRole: 40 i = index.row() 41 j = index.column() 42 part = self.participants[i] 43 if j == 0: 44 return part.name 45 elif j == 1: 46 return part.shortname 47 elif j == 2: 48 return part.color_index 49 elif j == 3: 50 return part.relative_rssi 51 elif j == 4: 52 return part.address_hex 53 54 def setData(self, index: QModelIndex, value, role=Qt.DisplayRole): 55 i = index.row() 56 j = index.column() 57 if i >= len(self.participants): 58 return False 59 60 participant = self.participants[i] 61 62 if j == 0: 63 participant.name = value 64 elif j == 1: 65 participant.shortname = value 66 elif j == 2: 67 participant.color_index = int(value) 68 elif j == 3: 69 for other in self.participants: 70 if other.relative_rssi == int(value): 71 other.relative_rssi = participant.relative_rssi 72 break 73 participant.relative_rssi = int(value) 74 elif j == 4: 75 participant.address_hex = value 76 77 self.update() 78 self.participant_edited.emit() 79 80 return True 81 82 def flags(self, index: QModelIndex): 83 if not index.isValid(): 84 return Qt.NoItemFlags 85 86 return Qt.ItemIsEditable | Qt.ItemIsEnabled | Qt.ItemIsSelectable 87 88 def __get_initial_name(self) -> (str, str): 89 given_names = set(p.name for p in self.participants) 90 name = next((name for name in self.INITIAL_NAMES if name not in given_names), None) 91 if name is not None: 92 return name, name[0] 93 94 name = next(("P" + str(i) for i in itertools.count() if "P" + str(i) not in given_names), None) 95 if name is not None: 96 return name, name[1:] 97 98 return "Participant X", "X" 99 100 def add_participant(self): 101 used_colors = set(p.color_index for p in self.participants) 102 avail_colors = set(range(0, len(settings.PARTICIPANT_COLORS))) - used_colors 103 if len(avail_colors) > 0: 104 color_index = avail_colors.pop() 105 else: 106 color_index = random.choice(range(len(settings.PARTICIPANT_COLORS))) 107 108 name, shortname = self.__get_initial_name() 109 participant = Participant(name, shortname=shortname, color_index=color_index) 110 self.participants.append(participant) 111 participant.relative_rssi = len(self.participants) - 1 112 113 self.update() 114 self.participant_edited.emit() 115 116 def remove_participants(self, selection: QItemSelection): 117 if len(self.participants) < 1: 118 return 119 120 if selection.isEmpty(): 121 start, end = len(self.participants) - 1, len(self.participants) - 1 # delete last element 122 else: 123 start, end = min([rng.top() for rng in selection]), max([rng.bottom() for rng in selection]) 124 125 del self.participants[start:end + 1] 126 num_removed = (end + 1) - start 127 for participant in self.participants: 128 if participant.relative_rssi > len(self.participants) - 1: 129 participant.relative_rssi -= num_removed 130 131 # fix duplicates 132 n = len(self.participants) 133 for p1, p2 in itertools.combinations(self.participants, 2): 134 if p1.relative_rssi == p2.relative_rssi: 135 p1.relative_rssi = next((i for i in range(n) 136 if i not in set(p.relative_rssi for p in self.participants)), 137 0) 138 139 self.update() 140 self.participant_edited.emit() 141 142 def move_up(self, selection: QItemSelection): 143 if selection.isEmpty() or len(self.participants) < 1: 144 return None, None 145 146 start, end = min([rng.top() for rng in selection]), max([rng.bottom() for rng in selection]) 147 if start == 0: 148 return None, None 149 150 for i in range(start, end + 1): 151 self.participants[i], self.participants[i - 1] = self.participants[i - 1], self.participants[i] 152 153 self.update() 154 self.participant_edited.emit() 155 156 return start, end 157 158 def move_down(self, selection: QItemSelection): 159 if selection.isEmpty() or len(self.participants) < 1: 160 return None, None 161 162 start, end = min([rng.top() for rng in selection]), max([rng.bottom() for rng in selection]) 163 if end >= len(self.participants) - 1: 164 return None, None 165 166 for i in reversed(range(start, end + 1)): 167 self.participants[i], self.participants[i + 1] = self.participants[i + 1], self.participants[i] 168 169 self.update() 170 self.participant_edited.emit() 171 172 return start, end 173