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