nmreval/src/gui_qt/fit/fit_parameter.py

from __future__ import annotations

from nmreval.fit.parameter import Parameter
from nmreval.utils.text import convert

from ..Qt import QtWidgets, QtCore, QtGui
from .._py.fitfuncwidget import Ui_FormFit
from .._py.fitmodelwidget import Ui_FitParameter
from ..lib.forms import SelectionWidget


class QFitParameterWidget(QtWidgets.QWidget, Ui_FormFit):
    value_requested = QtCore.pyqtSignal(int)

    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)

        self.func = None
        self.func_idx = None
        self.max_width = QtCore.QSize(0, 0)
        self.global_parameter = []
        self.data_parameter = []
        self.glob_values = None
        self.data_values = {}

        self.scrollwidget.setLayout(QtWidgets.QVBoxLayout())
        self.scrollwidget2.setLayout(QtWidgets.QVBoxLayout())

    def eventFilter(self, src: QtCore.QObject, evt: QtCore.QEvent):
        modifiers = QtCore.Qt.KeyboardModifier.ControlModifier | QtCore.Qt.KeyboardModifier.ShiftModifier
        if isinstance(evt, QtGui.QKeyEvent):
            if (evt.key() == QtCore.Qt.Key.Key_Right) and (evt.modifiers() == modifiers):
                self.change_single_parameter(src.value, sender=src)
                self.select_next_preview(1)

                return True

            elif (evt.key() == QtCore.Qt.Key.Key_Left) and (evt.modifiers() == modifiers):
                self.change_single_parameter(src.value, sender=src)
                self.select_next_preview(-1)

                return True

        return super().eventFilter(src, evt)

    def load(self, data):
        self.comboBox.blockSignals(True)
        while self.comboBox.count():
            self.comboBox.removeItem(0)

        for sid, name in data:
            self.comboBox.addItem(name, userData=sid)
            self._make_parameter(sid)
        self.comboBox.blockSignals(False)
        self.change_data(0)

    def set_function(self, func, idx):
        self.func = func
        self.func_idx = idx

        self.glob_values = [1] * len(func.params)

        for k, v in enumerate(func.params):
            widgt = ParameterGlobalWidget(name=v, parent=self.scrollwidget)
            widgt.parameter_pos = k
            widgt.func_idx = idx
            try:
                widgt.set_bounds(*func.bounds[k], False)
            except (AttributeError, IndexError):
                pass

            size = widgt.parametername.sizeHint()
            if self.max_width.width() < size.width():
                self.max_width = size

            widgt.state_changed.connect(self.make_global)
            widgt.value_requested.connect(self.look_for_value)
            widgt.value_changed.connect(self.change_global_parameter)
            widgt.replace_single_value.connect(self.delete_single_parameter)

            self.global_parameter.append(widgt)
            self.scrollwidget.layout().addWidget(widgt)

            widgt2 = ParameterSingleWidget(name=v, parent=self.scrollwidget2)
            widgt2.valueChanged.connect(self.change_single_parameter)
            widgt2.removeSingleValue.connect(self.change_single_parameter)
            widgt2.installEventFilter(self)
            self.scrollwidget2.layout().addWidget(widgt2)
            self.data_parameter.append(widgt2)

        for w1, w2 in zip(self.global_parameter, self.data_parameter):
            w1.parametername.setFixedSize(self.max_width)
            w1.checkBox.setFixedSize(self.max_width)
            w2.parametername.setFixedSize(self.max_width)

        if hasattr(func, 'choices') and func.choices is not None:
            cbox = func.choices
            for c in cbox:
                widgt = SelectionWidget(*c)
                widgt.selectionChanged.connect(self.change_global_choice)
                self.global_parameter.append(widgt)
                self.glob_values.append(widgt.value)
                self.scrollwidget.layout().addWidget(widgt)

                widgt2 = SelectionWidget(*c)
                widgt2.selectionChanged.connect(self.change_single_choice)
                self.data_parameter.append(widgt2)
                self.scrollwidget2.layout().addWidget(widgt2)

        for i in range(self.comboBox.count()):
            self._make_parameter(self.comboBox.itemData(i))

        self.scrollwidget.layout().addStretch(1)
        self.scrollwidget2.layout().addStretch(1)

    # def set_links(self, parameter):
    #     for w in self.global_parameter:
    #         if isinstance(w, FitModelWidget):
    #             w.add_links(parameter)

    @QtCore.pyqtSlot(str)
    def change_global_parameter(self, value: str, idx: int = None):
        if idx is None:
            idx = self.global_parameter.index(self.sender())

        # self.glob_values[idx] = float(value)
        self.glob_values[idx] = value
        if self.data_values[self.comboBox.currentData()][idx] is None:
            self.data_parameter[idx].blockSignals(True)
            # self.data_parameter[idx].value = float(value)
            self.data_parameter[idx].value = value
            self.data_parameter[idx].blockSignals(False)

    @QtCore.pyqtSlot(str, object)
    def change_global_choice(self, _, value):
        idx = self.global_parameter.index(self.sender())
        self.glob_values[idx] = value
        if self.data_values[self.comboBox.currentData()][idx] is None:
            self.data_parameter[idx].blockSignals(True)
            self.data_parameter[idx].value = value
            self.data_parameter[idx].blockSignals(False)

    def change_single_parameter(self, value: float = None, sender=None):
        if sender is None:
            sender = self.sender()
        idx = self.data_parameter.index(sender)
        self.data_values[self.comboBox.currentData()][idx] = value
        # look for global parameter values if value is reset, ie None
        if value is None:
            self.change_data(self.comboBox.currentIndex())

    def delete_single_parameter(self):
        idx = self.global_parameter.index(self.sender())
        for i in range(self.comboBox.count()):
            set_id = self.comboBox.itemData(i)
            self.data_values[set_id][idx] = None
        self.change_data(self.comboBox.currentIndex())

    def change_single_choice(self, _, value, sender=None):
        if sender is None:
            sender = self.sender()
        idx = self.data_parameter.index(sender)
        self.data_values[self.comboBox.currentData()][idx] = value

    @QtCore.pyqtSlot(object)
    def look_for_value(self, sender):
        self.value_requested.emit(self.global_parameter.index(sender))

    @QtCore.pyqtSlot()
    def make_global(self):
        # disable single parameter if it is set global, enable if global is unset
        widget = self.sender()
        idx = self.global_parameter.index(widget)
        enable = (widget.global_checkbox.checkState() == QtCore.Qt.CheckState.Unchecked)
        self.data_parameter[idx].setEnabled(enable)

    def select_next_preview(self, direction):
        curr_idx = self.comboBox.currentIndex()
        next_idx = (curr_idx + direction) % self.comboBox.count()
        self.comboBox.setCurrentIndex(next_idx)

    @QtCore.pyqtSlot(int, name='on_comboBox_currentIndexChanged')
    def change_data(self, idx: int):
        # new dataset is selected, look for locally set parameter else use global values
        sid = self.comboBox.itemData(idx)
        if sid not in self.data_values:
            self._make_parameter(sid)

        for i, value in enumerate(self.data_values[sid]):
            w = self.data_parameter[i]
            w.blockSignals(True)
            try:
                w.show_as_local_parameter(value is not None)
            except AttributeError:
                pass

            if value is None:
                w.value = self.glob_values[i]
            else:
                w.value = value
            w.blockSignals(False)

    def _make_parameter(self, sid):
        if sid not in self.data_values:
            self.data_values[sid] = [None] * len(self.data_parameter)

    def get_parameter(self, use_func=None) -> tuple[dict, list]:
        bds = []
        is_global = []
        is_fixed = []
        param_general = []

        for g in self.global_parameter:
            if isinstance(g, ParameterGlobalWidget):
                p_i, bds_i, fixed_i, global_i = g.get_parameter()
                parameter_i = Parameter(name=g.name, value=p_i, lb=bds_i[0], ub=bds_i[1], var=fixed_i)
                param_general.append(parameter_i)

                bds.append(bds_i)
                is_fixed.append(fixed_i)
                is_global.append(global_i)

        data_parameter = {}
        if use_func is None:
            use_func = list(self.data_values.keys())

        for sid, parameter in self.data_values.items():
            if sid not in use_func:
                continue

            kw_p = {}
            p = []

            for i, (p_i, g) in enumerate(zip(parameter, self.global_parameter)):
                if isinstance(g, ParameterGlobalWidget):
                    if (p_i is None) or is_global[i]:
                        # set has no oen value
                        p.append(param_general[i].copy())
                    else:
                        lb, ub = bds[i]
                        try:
                            if not ((lb < p_i < ub) or (not is_fixed[i])):
                                raise ValueError(f'Parameter {g.name} is outside bounds ({lb}, {ub})')
                        except TypeError:
                            pass

                        # create Parameter
                        p.append(
                            Parameter(name=g.name, value=p_i, lb=lb, ub=ub, var=is_fixed[i])
                        )

                else:
                    if p_i is None:
                        kw_p.update(g.value)
                    elif isinstance(p_i, dict):
                        kw_p.update(p_i)
                    else:
                        kw_p[g.argname] = p_i

            data_parameter[sid] = (p, kw_p)

        global_parameter = []
        for param, global_flag in zip(param_general, is_global):
            if global_flag:
                param.is_global = True
                global_parameter.append(param)
            else:
                global_parameter.append(None)

        return data_parameter, global_parameter

    def set_parameter(self, set_id: str | None, parameter: list[float]) -> int:
        num_parameter = list(filter(lambda g: not isinstance(g, SelectionWidget), self.global_parameter))
        param_len = len(num_parameter)
        if set_id is None:
            for i, g in enumerate(num_parameter):
                val = parameter[i]
                g.set_parameter(val)
                self.glob_values[i] = val

        else:
            new_param = self.data_values[set_id]
            min_len = min(param_len, len(new_param))
            for i in range(min_len):
                new_param[i] = parameter[i]

            self.change_data(self.comboBox.currentIndex())

        return param_len


class ParameterSingleWidget(QtWidgets.QWidget):
    valueChanged = QtCore.pyqtSignal(object)
    removeSingleValue = QtCore.pyqtSignal()

    def __init__(self, name: str, parent=None):
        super().__init__(parent=parent)

        self._init_ui()

        self.name = name
        self.parametername.setText(convert(name))
        self.parametername.setToolTip('If this is bold then this parameter is only for this data. '
                              'Otherwise, the general parameter is used and displayed')

        # self.value_line.setValidator(QtGui.QDoubleValidator())
        self.value_line.textChanged.connect(lambda: self.valueChanged.emit(self.value) if self.value is not None else 0)
        self.reset_button.clicked.connect(lambda x: self.removeSingleValue.emit())

    def _init_ui(self):
        layout = QtWidgets.QHBoxLayout(self)
        layout.setContentsMargins(2, 2, 2, 2)
        layout.setSpacing(2)

        self.parametername = QtWidgets.QLabel(self)
        layout.addWidget(self.parametername)

        layout.addSpacerItem(QtWidgets.QSpacerItem(0, 0, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum))

        self.value_line = QtWidgets.QLineEdit(self)
        self.value_line.textEdited.connect(lambda x: self.show_as_local_parameter(True))
        layout.addWidget(self.value_line)

        self.reset_button = QtWidgets.QToolButton(self)
        self.reset_button.setText('Use global')
        self.reset_button.clicked.connect(lambda: self.show_as_local_parameter(False))
        layout.addWidget(self.reset_button)

        self.setLayout(layout)

    @property
    def value(self) -> float:
        try:
            return float(self.value_line.text().replace(',', '.'))
        except ValueError:
            return 0.0

    @value.setter
    def value(self, val):
        # self.value_line.setText(f'{float(val):.5g}')
        self.value_line.setText(f'{val}')
        self.value_line.setCursorPosition(0)

    def show_as_local_parameter(self, is_local: bool):
        if is_local:
            self.parametername.setStyleSheet('font-weight: bold;')
        else:
            self.parametername.setStyleSheet('')


class ParameterGlobalWidget(QtWidgets.QWidget, Ui_FitParameter):
    """
    Widget to show a global parameter
    """

    value_requested = QtCore.pyqtSignal(object)
    value_changed = QtCore.pyqtSignal(str)
    state_changed = QtCore.pyqtSignal()
    replace_single_value = QtCore.pyqtSignal(object)

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

        self.name = name
        self.reset_button.setVisible(False)

        self.parametername.setText(convert(name) + ' ')

        self.parameter_line.setText('1')
        self.parameter_line.setMaximumWidth(160)
        self.lineEdit.setMaximumWidth(100)
        self.lineEdit_2.setMaximumWidth(100)

        self.label_3.setText(f'&lt; {convert(name)} &lt;')

        self.checkBox.stateChanged.connect(self.enableBounds)
        self.global_checkbox.stateChanged.connect(lambda: self.state_changed.emit())
        self.parameter_line.editingFinished.connect(self.update_parameter)
        self.parameter_line.values_requested.connect(lambda: self.value_requested.emit(self))
        self.parameter_line.replace_single_values.connect(lambda: self.replace_single_value.emit(None))
        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.reset_button.setVisible(False)

        self.parameter_pos = None
        self.func_idx = None

        self._linetext = '1'

        self.menu = QtWidgets.QMenu(self)

    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.CheckState.Checked if cbox else QtCore.Qt.CheckState.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):
        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.CheckState.Unchecked if fixed else QtCore.Qt.CheckState.Checked)

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

    def get_parameter(self):
        try:
            p = float(self.parameter_line.text().replace(',', '.'))
        except ValueError:
            p = self.parameter_line.text().replace(',', '.')

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

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

        bounds = (lb, rb)

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

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

    @QtCore.pyqtSlot()
    def update_parameter(self):
        new_value = self.parameter_line.text()
        if not new_value:
            self.parameter_line.setText('1')

        try:
            float(new_value)
            is_text = False
        except ValueError:
            is_text = True
            self.global_checkbox.setCheckState(False)

        self.set_fixed(is_text or self.fixed_check.isChecked())