299 lines
9.7 KiB
Python
299 lines
9.7 KiB
Python
from __future__ import annotations
|
|
|
|
from ..Qt import QtCore, QtWidgets, QtGui
|
|
from .._py.save_fitmodel_dialog import Ui_SaveDialog
|
|
from ..lib.iconloading import get_icon
|
|
from ..lib.tables import TableWidget
|
|
|
|
|
|
class QSaveModelDialog(QtWidgets.QDialog, Ui_SaveDialog):
|
|
def __init__(self, types=None, parent=None):
|
|
super().__init__(parent=parent)
|
|
self.setupUi(self)
|
|
|
|
if types is None:
|
|
types = []
|
|
|
|
self.comboBox.blockSignals(True)
|
|
self.comboBox.addItems(types)
|
|
self.comboBox.addItem('New group...')
|
|
self.comboBox.blockSignals(False)
|
|
|
|
self.frame.hide()
|
|
|
|
@QtCore.pyqtSlot(int, name='on_comboBox_currentIndexChanged')
|
|
def new_group(self, idx: int):
|
|
if idx == self.comboBox.count() - 1:
|
|
self.frame.show()
|
|
else:
|
|
self.lineEdit_2.clear()
|
|
self.frame.hide()
|
|
|
|
@QtCore.pyqtSlot(name='on_toolButton_clicked')
|
|
def accept_group(self):
|
|
self.comboBox.insertItem(self.comboBox.count() - 1, self.lineEdit_2.text())
|
|
self.comboBox.setCurrentIndex(self.comboBox.count() - 2)
|
|
|
|
def accept(self):
|
|
if self.lineEdit.text():
|
|
self.close()
|
|
|
|
|
|
class FitModelTree(QtWidgets.QTreeWidget):
|
|
icons = ['plus', 'mal_icon', 'minus_icon', 'geteilt_icon']
|
|
|
|
treeChanged = QtCore.pyqtSignal()
|
|
itemRemoved = QtCore.pyqtSignal(int)
|
|
|
|
counterRole = QtCore.Qt.ItemDataRole.UserRole + 1
|
|
operatorRole = QtCore.Qt.ItemDataRole.UserRole + 2
|
|
|
|
def __init__(self, parent=None):
|
|
super().__init__(parent=parent)
|
|
self.setHeaderHidden(True)
|
|
self.setDragEnabled(True)
|
|
self.setDragDropMode(QtWidgets.QTreeWidget.DragDropMode.InternalMove)
|
|
self.setDefaultDropAction(QtCore.Qt.DropAction.MoveAction)
|
|
|
|
self.itemSelectionChanged.connect(lambda: self.treeChanged.emit())
|
|
|
|
def keyPressEvent(self, evt):
|
|
operators = [
|
|
QtCore.Qt.Key.Key_Plus,
|
|
QtCore.Qt.Key.Key_Asterisk,
|
|
QtCore.Qt.Key.Key_Minus,
|
|
QtCore.Qt.Key.Key_Slash
|
|
]
|
|
|
|
if evt.key() == QtCore.Qt.Key.Key_Delete:
|
|
for item in self.selectedItems():
|
|
self.remove_function(item)
|
|
|
|
elif evt.key() == QtCore.Qt.Key.Key_Space:
|
|
for item in self.selectedItems():
|
|
cs = item.checkState(0)
|
|
if cs == QtCore.Qt.CheckState.Unchecked:
|
|
item.setCheckState(0, QtCore.Qt.CheckState.Checked)
|
|
else:
|
|
item.setCheckState(0, QtCore.Qt.CheckState.Unchecked)
|
|
|
|
elif evt.key() in operators:
|
|
idx = operators.index(evt.key())
|
|
for item in self.selectedItems():
|
|
item.setData(0, self.operatorRole, idx)
|
|
item.setIcon(0, get_icon(self.icons[idx]))
|
|
|
|
else:
|
|
super().keyPressEvent(evt)
|
|
|
|
def dropEvent(self, evt: QtGui.QDropEvent):
|
|
super().dropEvent(evt)
|
|
self.treeChanged.emit()
|
|
|
|
def remove_function(self, item: QtWidgets.QTreeWidgetItem):
|
|
"""
|
|
Remove function and children from tree and dictionary
|
|
"""
|
|
while item.childCount():
|
|
self.remove_function(item.child(0))
|
|
|
|
if item.parent():
|
|
item.parent().removeChild(item)
|
|
else:
|
|
self.invisibleRootItem().removeChild(item)
|
|
|
|
idx = item.data(0, self.counterRole)
|
|
self.itemRemoved.emit(idx)
|
|
|
|
def add_function(self,
|
|
idx: int,
|
|
cnt: int,
|
|
op: int,
|
|
name: str,
|
|
color: QtGui.QColor | str | tuple,
|
|
parent: QtWidgets.QTreeWidgetItem = None,
|
|
children: list = None,
|
|
active: bool = True,
|
|
param_names: list[str] = None,
|
|
**kwargs):
|
|
"""
|
|
Add function to tree and dictionary of functions.
|
|
"""
|
|
if not isinstance(color, QtGui.QColor):
|
|
if isinstance(color, tuple):
|
|
color = QtGui.QColor.fromRgbF(*color)
|
|
else:
|
|
color = QtGui.QColor(color)
|
|
|
|
it = QtWidgets.QTreeWidgetItem()
|
|
it.setData(0, QtCore.Qt.ItemDataRole.UserRole, idx)
|
|
it.setData(0, self.counterRole, cnt)
|
|
it.setData(0, self.operatorRole, op)
|
|
it.setText(0, name)
|
|
if param_names is not None:
|
|
it.setToolTip(0,
|
|
'Parameter names:\n' +
|
|
'\n'.join(f'{pn}({cnt})' for pn in param_names))
|
|
it.setForeground(0, QtGui.QBrush(color))
|
|
|
|
it.setIcon(0, get_icon(self.icons[op]))
|
|
it.setCheckState(0, QtCore.Qt.CheckState.Checked if active else QtCore.Qt.CheckState.Unchecked)
|
|
|
|
if parent is None:
|
|
self.addTopLevelItem(it)
|
|
else:
|
|
parent.addChild(it)
|
|
|
|
if children is not None:
|
|
for c in children:
|
|
self.add_function(**c, parent=it)
|
|
|
|
self.setCurrentIndex(self.indexFromItem(it, 0))
|
|
|
|
def sizeHint(self):
|
|
w = super().sizeHint().width()
|
|
return QtCore.QSize(w, 100)
|
|
|
|
def get_selected(self):
|
|
try:
|
|
it = self.selectedItems()[0]
|
|
function_nr = it.data(0, QtCore.Qt.ItemDataRole.UserRole)
|
|
idx = it.data(0, self.counterRole)
|
|
|
|
except IndexError:
|
|
function_nr = None
|
|
idx = None
|
|
|
|
return function_nr, idx
|
|
|
|
def get_functions(self, full: bool = True, pos: int = -1, return_pos: bool = False, parent=None):
|
|
"""
|
|
Create nested list of functions in tree. Parameters saved are idx (Index of function in list of all functions),
|
|
cnt (counter of number to associate with functione values), ops (+, -, *, /), and maybe children.
|
|
"""
|
|
if parent is None:
|
|
parent = self.invisibleRootItem()
|
|
|
|
funcs = []
|
|
for i in range(parent.childCount()):
|
|
pos += 1
|
|
it = parent.child(i)
|
|
|
|
child = {
|
|
'idx': it.data(0, QtCore.Qt.ItemDataRole.UserRole),
|
|
'op': it.data(0, self.operatorRole),
|
|
'pos': pos,
|
|
'active': (it.checkState(0) == QtCore.Qt.CheckState.Checked),
|
|
'children': []
|
|
}
|
|
|
|
if full:
|
|
child['name'] = it.text(0)
|
|
child['cnt'] = it.data(0, self.counterRole)
|
|
child['color'] = it.foreground(0).color().getRgbF()
|
|
|
|
if it.childCount():
|
|
child['children'], pos = self.get_functions(full=full, parent=it, pos=pos, return_pos=True)
|
|
|
|
funcs.append(child)
|
|
|
|
if return_pos:
|
|
return funcs, pos
|
|
else:
|
|
return funcs
|
|
|
|
|
|
class FitTableWidget(TableWidget):
|
|
def __init__(self, parent=None):
|
|
super().__init__(parent=parent)
|
|
|
|
self.horizontalHeader().hide()
|
|
self.verticalHeader().hide()
|
|
self.setColumnCount(2)
|
|
self.setSelectionBehavior(QtWidgets.QTableWidget.SelectionBehavior.SelectRows)
|
|
self.horizontalHeader().setStretchLastSection(True)
|
|
self.hideColumn(1)
|
|
|
|
def add_model(self, idx: str):
|
|
model_count = 0
|
|
for r in range(self.rowCount()):
|
|
cb = self.cellWidget(r, 1)
|
|
cb.addItem('Model ' + str(idx), userData=idx)
|
|
model_count = cb.count()
|
|
|
|
if model_count > 2:
|
|
if self.isColumnHidden(1):
|
|
self.showColumn(1)
|
|
self.resizeColumnToContents(0)
|
|
self.setColumnWidth(1, self.columnWidth(0) - self.columnWidth(1))
|
|
|
|
def remove_model(self, idx: str):
|
|
model_count = 0
|
|
for r in range(self.rowCount()):
|
|
cb = self.cellWidget(r, 1)
|
|
if cb.currentData() == idx:
|
|
cb.setCurrentIndex(0)
|
|
cb.removeItem(cb.findData(idx))
|
|
model_count = cb.count()
|
|
|
|
if model_count == 2:
|
|
self.hideColumn(1)
|
|
self.resizeColumnToContents(0)
|
|
|
|
def load(self, set_ids: list[str]):
|
|
self.blockSignals(True)
|
|
|
|
while self.rowCount():
|
|
self.removeRow(0)
|
|
|
|
self.setColumnCount(2)
|
|
self.hideColumn(1)
|
|
|
|
for (sid, name) in set_ids:
|
|
item = QtWidgets.QTableWidgetItem(name)
|
|
item.setCheckState(QtCore.Qt.CheckState.Checked)
|
|
item.setData(QtCore.Qt.ItemDataRole.UserRole+1, sid)
|
|
row = self.rowCount()
|
|
self.setRowCount(row+1)
|
|
self.setItem(row, 0, item)
|
|
|
|
item2 = QtWidgets.QTableWidgetItem('')
|
|
self.setItem(row, 1, item2)
|
|
cb = QtWidgets.QComboBox(parent=self)
|
|
cb.addItem('Default')
|
|
self.setCellWidget(row, 1, cb)
|
|
|
|
self.blockSignals(False)
|
|
|
|
def collect_data(self, default: str = None, include_name: bool = False) -> dict:
|
|
data = {}
|
|
|
|
for i in range(self.rowCount()):
|
|
item = self.item(i, 0)
|
|
if item.checkState() == QtCore.Qt.CheckState.Checked:
|
|
mod = self.cellWidget(i, 1).currentData()
|
|
if mod is None:
|
|
mod = default
|
|
|
|
if include_name:
|
|
arg = (item.data(QtCore.Qt.ItemDataRole.UserRole+1), item.text())
|
|
else:
|
|
arg = item.data(QtCore.Qt.ItemDataRole.UserRole+1)
|
|
|
|
if mod not in data:
|
|
data[mod] = []
|
|
data[mod].append(arg)
|
|
|
|
return data
|
|
|
|
def data_list(self, include_name: bool = True) -> list:
|
|
ret_val = []
|
|
for i in range(self.rowCount()):
|
|
item = self.item(i, 0)
|
|
if include_name:
|
|
ret_val.append((item.data(QtCore.Qt.ItemDataRole.UserRole+1), item.text()))
|
|
else:
|
|
ret_val.append(item.data(QtCore.Qt.ItemDataRole.UserRole+1))
|
|
|
|
return ret_val
|