1# -*- coding: utf-8 -*-
2
3import numpy as np
4from PySide2.QtCore import QAbstractTableModel, QModelIndex, Qt
5from PySide2.QtGui import QColor
6
7
8class Model(QAbstractTableModel):
9    def __init__(self):
10        QAbstractTableModel.__init__(self)
11
12        self.ncols = 4
13        nrows = 1  # initial number of rows
14
15        self.data_x = np.zeros(nrows)
16        self.data_xerr = np.zeros(nrows)
17        self.data_y = np.zeros(nrows)
18        self.data_yerr = np.zeros(nrows)
19
20    def rowCount(self, parent=QModelIndex()):
21        return self.data_x.size
22
23    def columnCount(self, parent=QModelIndex()):
24        return self.ncols
25
26    def flags(self, index):
27        return Qt.ItemIsEditable | Qt.ItemIsEnabled | Qt.ItemIsSelectable
28
29    def setData(self, index, value, role):
30        if index.isValid() and role == Qt.EditRole:
31            column = index.column()
32            row = index.row()
33
34            try:
35                float_value = float(value)
36
37                if column == 0:
38                    self.data_x[row] = float_value
39
40                if column == 1:
41                    self.data_xerr[row] = float_value
42
43                if column == 2:
44                    self.data_y[row] = float_value
45
46                if column == 3:
47                    self.data_yerr[row] = float_value
48
49                self.dataChanged.emit(index, index)
50
51                return True
52            except ValueError:
53                return False
54
55        return False
56
57    def headerData(self, section, orientation, role):
58        if role != Qt.DisplayRole:
59            return None
60
61        if orientation == Qt.Horizontal:
62            return ("x", "xerr", "y", "yerr")[section]
63
64        return "{}".format(section)
65
66    def data(self, index, role):
67        if role == Qt.BackgroundRole:
68            return QColor(Qt.transparent)
69
70        elif role == Qt.TextAlignmentRole:
71            return Qt.AlignRight
72
73        elif role == Qt.DisplayRole:
74            column = index.column()
75            row = index.row()
76
77            if column == 0:
78                return "{0:.6e}".format(self.data_x[row])
79
80            if column == 1:
81                return "{0:.6e}".format(self.data_xerr[row])
82
83            if column == 2:
84                return "{0:.6e}".format(self.data_y[row])
85
86            if column == 3:
87                return "{0:.6e}".format(self.data_yerr[row])
88
89    def sort(self, column, order):
90        if column == 0:
91            sorted_idx = self.data_x.argsort()
92
93            if order == Qt.SortOrder.DescendingOrder:
94                sorted_idx = sorted_idx[::-1]
95
96            self.beginResetModel()
97
98            self.data_x = self.data_x[sorted_idx]
99            self.data_xerr = self.data_xerr[sorted_idx]
100            self.data_y = self.data_y[sorted_idx]
101            self.data_yerr = self.data_yerr[sorted_idx]
102
103            self.endResetModel()
104
105            first_index = self.index(0, 0)
106            last_index = self.index(self.rowCount() - 1, self.columnCount() - 1)
107
108            self.dataChanged.emit(first_index, last_index)
109
110    def append_row(self):
111        self.beginInsertRows(QModelIndex(), self.data_x.size, self.data_x.size)
112
113        self.data_x = np.append(self.data_x, 0)
114        self.data_xerr = np.append(self.data_xerr, 0)
115        self.data_y = np.append(self.data_y, 0)
116        self.data_yerr = np.append(self.data_yerr, 0)
117
118        self.endInsertRows()
119
120        first_index = self.index(0, 0)
121        last_index = self.index(self.rowCount() - 1, self.columnCount() - 1)
122
123        self.dataChanged.emit(first_index, last_index)
124
125    def remove_rows(self, index_list):
126        index_list.sort(reverse=True)
127
128        self.data_x = np.delete(self.data_x, index_list)
129        self.data_xerr = np.delete(self.data_xerr, index_list)
130        self.data_y = np.delete(self.data_y, index_list)
131        self.data_yerr = np.delete(self.data_yerr, index_list)
132
133        for index in index_list:
134            self.beginRemoveRows(QModelIndex(), index, index)
135            self.endRemoveRows()
136
137        first_index = self.index(index_list[-1], 0)
138        last_index = self.index(index_list[0], 3)
139
140        self.dataChanged.emit(first_index, last_index)
141