nmreval/src/gui_qt/fit/fitfunction.py

from __future__ import annotations

from itertools import cycle, count

from nmreval.configs import config_paths
from nmreval import models
from nmreval.lib.importer import find_models
from nmreval.lib.colors import BaseColor, Tab10
from nmreval.lib.logger import logger
from nmreval.utils.text import convert

from ..lib.iconloading import get_icon
from .._py.fitfunctionwidget import Ui_Form
from ..Qt import QtWidgets, QtCore, QtGui


class QFunctionWidget(QtWidgets.QWidget, Ui_Form):
    func_cnt = count()
    func_colors = cycle(Tab10)
    op_names = ['plus_icon', 'mal_icon', 'minus_icon', 'geteilt_icon']

    newFunction = QtCore.pyqtSignal(int, int)
    treeChanged = QtCore.pyqtSignal()
    itemRemoved = QtCore.pyqtSignal(int)
    showFunction = QtCore.pyqtSignal(int)

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

        self.setupUi(self)

        self._types = []
        self.user_functions = {}
        self.functions = find_models(models)
        self.read_and_load_functions()

        self.functree.treeChanged.connect(lambda: self.treeChanged.emit())
        self.functree.itemRemoved.connect(self.remove_function)

        self.iscomplex = False
        self.complex_widget.hide()

        for i, op_icon in enumerate(self.op_names):
            self.operator_combobox.setItemIcon(i, get_icon(op_icon))

    def read_and_load_functions(self):
        """
        Reads user-defined fit functions from usermodels.py and adds them to list of fit function.
        Previous class definitions are replaced.
        """
        user_defined = []
        try:
            user_defined = find_models(config_paths() / 'usermodels.py')
        except Exception as e:
            import traceback
            _ = QtWidgets.QMessageBox.warning(
                self,
                'No user functions',
                f'Loading user-defined function failed with exception:\n'
                f'{"".join(traceback.format_exception(e, limit=1))}')
            logger.exception("Invalid usermodels.py", exc_info=e)

        for model in user_defined:
            name = model.__name__
            if name not in self.user_functions:
                self.functions.append(model)
                self.user_functions[name] = model
            else:
                idx = self.functions.index(self.user_functions[name])
                self.functions[idx] = model
                self.user_functions[name] = model

        self._types = []
        for m in self.functions:
            try:
                m.type
            except AttributeError:
                m.type = 'Other'

            if m.type not in self._types:
                self._types.append(m.type)

        self.typecomboBox.blockSignals(True)
        self.typecomboBox.clear()
        self.typecomboBox.addItems(sorted(self._types))
        self.typecomboBox.blockSignals(False)

        self.change_group(0)

    def __len__(self) -> int:
        num = 0
        iterator = QtWidgets.QTreeWidgetItemIterator(self.functree)
        while iterator.value():
            num += 1
            iterator += 1

        return num

    @QtCore.pyqtSlot(int, name='on_typecomboBox_currentIndexChanged')
    def change_group(self, idx: int):
        """
        Change items in fitcombobox to new entries
        """
        self.fitcomboBox.blockSignals(True)
        while self.fitcomboBox.count():
            self.fitcomboBox.removeItem(0)

        selected_type = self.typecomboBox.itemText(idx)
        for m in self.functions:
            if m.type == selected_type:
                self.fitcomboBox.addItem(m.name, userData=self.functions.index(m))
        self.fitcomboBox.blockSignals(False)
        self.on_fitcomboBox_currentIndexChanged(0)

    @QtCore.pyqtSlot(int, name='on_fitcomboBox_currentIndexChanged')
    def change_function(self, idx: int):
        """
        Display new equation on changing function
        """
        index = self.fitcomboBox.itemData(idx)
        if self.functions:
            fitfunc = self.functions[index]
            try:
                self.fitequation.setText(convert(fitfunc.equation))
            except AttributeError:
                self.fitequation.setText('')

    @QtCore.pyqtSlot(name='on_use_function_button_clicked')
    def new_function(self):
        idx = self.fitcomboBox.itemData(self.fitcomboBox.currentIndex())
        cnt = next(self.func_cnt)
        op = self.operator_combobox.currentIndex()
        name = self.functions[idx].name
        col = next(self.func_colors)

        self.newFunction.emit(idx, cnt)

        self.add_function(idx, cnt, op, name, col, param_names=self.functions[idx].params)

    def add_function(self, idx: int, cnt: int, op: int,
                     name: str, color: str | tuple[float, float, float] | BaseColor, **kwargs):
        """
        Add function to tree and dictionary of functions.
        """
        if isinstance(color, BaseColor):
            qcolor = QtGui.QColor.fromRgbF(*color.rgb(normed=True))
        elif isinstance(color, tuple):
            qcolor = QtGui.QColor.fromRgbF(*color)
        else:
            qcolor = QtGui.QColor(color)

        self.functree.add_function(idx, cnt, op, name, qcolor, **kwargs)

        f = self.functions[idx]
        if hasattr(f, 'iscomplex') and f.iscomplex:
            self.iscomplex = True
            self.complex_widget.show()

    @QtCore.pyqtSlot(int)
    def remove_function(self, idx: int):
        iterator = QtWidgets.QTreeWidgetItemIterator(self.functree)
        self.iscomplex = False
        while iterator.value():
            item = iterator.value()
            f = self.functions[item.data(0, QtCore.Qt.UserRole)]
            if hasattr(f, 'iscomplex') and f.iscomplex:
                self.iscomplex = True
                break

            iterator += 1
        self.complex_widget.setVisible(self.iscomplex)

        self.itemRemoved.emit(idx)

    @QtCore.pyqtSlot(QtWidgets.QTreeWidgetItem, int, name='on_functree_itemClicked')
    def show_parameter(self, item: QtWidgets.QTreeWidgetItem, _: int):
        self.showFunction.emit(item.data(0, self.functree.counterRole))

    def get_selected(self):
        function_nr, idx = self.functree.get_selected()

        if function_nr is not None:
            return self.functions[function_nr], idx
        else:
            return None, None

    def get_functions(self, full: bool = True, clsname=False, include_all: bool = True):
        """
        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.
        """

        used_functions = self.functree.get_functions(full=full)
        self._prepare_function_for_model(used_functions, full=full, clsname=clsname, include_all=include_all)

        return used_functions

    def _prepare_function_for_model(self, func_list: list[dict],
                                    full: bool = True, clsname: bool = False, include_all: bool = True):

        for func_args in func_list:
            is_active = func_args.get('active')
            if (not is_active) and (not include_all):
                continue

            if not clsname:
                func_args['func'] = self.functions[func_args['idx']]
            else:
                func_args['func'] = self.functions[func_args['idx']].name

            if not full:
                func_args.pop('active')
                func_args.pop('idx')

            if func_args['children']:
                self._prepare_function_for_model(func_args['children'],
                                                 full=full, clsname=clsname,
                                                 include_all=include_all)

    def get_parameter_list(self):
        all_parameters = {}

        iterator = QtWidgets.QTreeWidgetItemIterator(self.functree)
        while iterator.value():
            item = iterator.value()
            f = self.functions[item.data(0, QtCore.Qt.UserRole)]
            cnt = item.data(0, self.functree.counterRole)
            all_parameters[f'{f.name}_{cnt}'] = [(convert(pp, new='str'), (cnt, i)) for i, pp in enumerate(f.params)]

            iterator += 1

        return all_parameters

    def get_complex_state(self):
        return self.complex_comboBox.currentIndex() if self.iscomplex else None

    def set_complex_state(self, state):
        if state is not None:
            self.complex_comboBox.setCurrentIndex(state)

    def clear(self):
        self.functree.blockSignals(True)
        self.functree.clear()
        self.functree.blockSignals(False)

        self.complex_comboBox.setCurrentIndex(0)
        self.complex_widget.hide()