nmreval/src/gui_qt/fit/fit_forms.py

from __future__ import annotations

from nmreval.utils.text import convert

from ..Qt import QtCore, QtWidgets, QtGui
from .._py.fitmodelwidget import Ui_FitParameter
from .._py.save_fitmodel_dialog import Ui_SaveDialog
from ..lib import get_icon


class FitModelWidget(QtWidgets.QWidget, Ui_FitParameter):
    value_requested = QtCore.pyqtSignal(object)
    value_changed = QtCore.pyqtSignal(str)
    state_changed = QtCore.pyqtSignal()

    def __init__(self, label: str = 'Fitparameter', parent=None, fixed: bool = False):
        super().__init__(parent)
        self.setupUi(self)

        self.parametername.setText(label + ' ')

        validator = QtGui.QDoubleValidator()
        validator.setDecimals(9)
        self.parameter_line.setValidator(validator)
        self.parameter_line.setText('1')
        self.parameter_line.setMaximumWidth(60)
        self.lineEdit.setMaximumWidth(60)
        self.lineEdit_2.setMaximumWidth(60)

        self.label_3.setText(f'< {label} <')

        self.checkBox.stateChanged.connect(self.enableBounds)
        self.global_checkbox.stateChanged.connect(lambda: self.state_changed.emit())
        self.parameter_line.values_requested.connect(lambda: self.value_requested.emit(self))
        self.parameter_line.editingFinished.connect(lambda: self.value_changed.emit(self.parameter_line.text()))
        self.fixed_check.toggled.connect(self.set_fixed)

        if fixed:
            self.fixed_check.hide()

        self.menu = QtWidgets.QMenu(self)
        self.add_links()

        self.is_linked = None
        self.parameter_pos = None
        self.func_idx = None

        self._linetext = '1'

    @property
    def name(self):
        return convert(self.parametername.text().strip(), old='html', new='str')

    def set_parameter_string(self, p: str):
        self.parameter_line.setText(p)
        self.parameter_line.setToolTip(p)

    def set_bounds(self, lb: float, ub: float, cbox: bool = True):
        self.checkBox.setCheckState(QtCore.Qt.Checked if cbox else QtCore.Qt.Unchecked)
        for val, bds_line in [(lb, self.lineEdit), (ub, self.lineEdit_2)]:
            if val is not None:
                bds_line.setText(str(val))
            else:
                bds_line.setText('')

    def enableBounds(self, value: int):
        self.lineEdit.setEnabled(value == 2)
        self.lineEdit_2.setEnabled(value == 2)

    def set_parameter(self, p: float | None, bds: tuple[float, float, bool] = None,
                      fixed: bool = None, glob: bool = None):
        if p is None:
            # bad hack: linked parameter return (None, linked parameter)
            # if p is None -> parameter is linked to argument given by bds
            self.link_parameter(linkto=bds)
        else:
            ptext = f'{p:.4g}'

            self.set_parameter_string(ptext)

            if bds is not None:
                self.set_bounds(*bds)

            if fixed is not None:
                self.fixed_check.setCheckState(QtCore.Qt.Unchecked if fixed else QtCore.Qt.Checked)

            if glob is not None:
                self.global_checkbox.setCheckState(QtCore.Qt.Checked if glob else QtCore.Qt.Unchecked)

    def get_parameter(self):
        if self.is_linked:
            try:
                p = float(self._linetext)
            except ValueError:
                p = 1.0
        else:
            try:
                p = float(self.parameter_line.text().replace(',', '.'))
            except ValueError:
                _ = QtWidgets.QMessageBox().warning(self, 'Invalid value',
                                                    f'{self.parametername.text()} contains invalid values',
                                                    QtWidgets.QMessageBox.Cancel)
                return None

        if self.checkBox.isChecked():
            try:
                lb = float(self.lineEdit.text().replace(',', '.'))
            except ValueError:
                lb = None

            try:
                rb = float(self.lineEdit_2.text().replace(',', '.'))
            except ValueError:
                rb = None
        else:
            lb = rb = None

        bounds = (lb, rb)

        return p, bounds, not self.fixed_check.isChecked(), self.global_checkbox.isChecked(), self.is_linked

    @QtCore.pyqtSlot(bool)
    def set_fixed(self, state: bool):
        # self.global_checkbox.setVisible(not state)
        self.frame.setVisible(not state)

    def add_links(self, parameter: dict = None):
        if parameter is None:
            parameter = {}
        self.menu.clear()

        ac = QtWidgets.QAction('Link to...', self)
        ac.triggered.connect(self.link_parameter)
        self.menu.addAction(ac)

        for model_key, model_funcs in parameter.items():
            m = QtWidgets.QMenu('Model ' + model_key, self)
            for func_name, func_params in model_funcs.items():
                m2 = QtWidgets.QMenu(func_name, m)
                for p_name, idx in func_params:
                    ac = QtWidgets.QAction(p_name, m2)
                    ac.setData((model_key, *idx))
                    ac.triggered.connect(self.link_parameter)
                    m2.addAction(ac)
                m.addMenu(m2)
            self.menu.addMenu(m)

        self.toolButton.setMenu(self.menu)

    @QtCore.pyqtSlot()
    def link_parameter(self, linkto=None):
        if linkto is None:
            action = self.sender()
        else:
            action = False
            for m in self.menu.actions():
                if m.menu():
                    for a in m.menu().actions():
                        if a.data() == linkto:
                            action = a
                            break
                if action:
                    break

        if (self.func_idx, self.parameter_pos) == action.data():
            return

        try:
            new_text = f'Linked to {action.parentWidget().title()}.{action.text()}'
            self._linetext = self.parameter_line.text()
            self.parameter_line.setText(new_text)
            self.parameter_line.setEnabled(False)
            self.global_checkbox.hide()
            self.global_checkbox.blockSignals(True)
            self.global_checkbox.setCheckState(QtCore.Qt.Checked)
            self.global_checkbox.blockSignals(False)
            self.frame.hide()
            self.is_linked = action.data()

        except AttributeError:
            self.parameter_line.setText(self._linetext)
            self.parameter_line.setEnabled(True)
            if self.fixed_check.isEnabled():
                self.global_checkbox.show()
                self.frame.show()
            self.is_linked = None

        self.state_changed.emit()


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.UserRole + 1
    operatorRole = QtCore.Qt.UserRole + 2

    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setHeaderHidden(True)
        self.setDragEnabled(True)
        self.setDragDropMode(QtWidgets.QTreeWidget.InternalMove)
        self.setDefaultDropAction(QtCore.Qt.MoveAction)

        self.itemSelectionChanged.connect(lambda: self.treeChanged.emit())

    def keyPressEvent(self, evt):
        operators = [QtCore.Qt.Key_Plus, QtCore.Qt.Key_Asterisk,
                     QtCore.Qt.Key_Minus, QtCore.Qt.Key_Slash]

        if evt.key() == QtCore.Qt.Key_Delete:
            for item in self.selectedItems():
                self.remove_function(item)

        elif evt.key() == QtCore.Qt.Key_Space:
            for item in self.treeWidget.selectedItems():
                item.setCheckState(0, QtCore.Qt.Checked) if item.checkState(
                    0) == QtCore.Qt.Unchecked else item.setCheckState(0, QtCore.Qt.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, **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.UserRole, idx)
        it.setData(0, self.counterRole, cnt)
        it.setData(0, self.operatorRole, op)
        it.setText(0, name)
        it.setForeground(0, QtGui.QBrush(color))

        it.setIcon(0, get_icon(self.icons[op]))
        it.setCheckState(0, QtCore.Qt.Checked if active else QtCore.Qt.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.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.UserRole),
                'op': it.data(0, self.operatorRole),
                'pos': pos,
                'active': (it.checkState(0) == QtCore.Qt.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(QtWidgets.QTableWidget):
    def __init__(self, parent=None):
        super().__init__(parent=parent)

        self.horizontalHeader().hide()
        self.verticalHeader().hide()
        self.setColumnCount(2)
        self.setSelectionBehavior(QtWidgets.QTableWidget.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.Checked)
            item.setData(QtCore.Qt.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.Checked:
                mod = self.cellWidget(i, 1).currentData()
                if mod is None:
                    mod = default

                if include_name:
                    arg = (item.data(QtCore.Qt.UserRole+1), item.text())
                else:
                    arg = item.data(QtCore.Qt.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.UserRole+1), item.text()))
            else:
                ret_val.append(item.data(QtCore.Qt.UserRole+1))

        return ret_val