nmreval/src/gui_qt/lib/namespace.py

import inspect
import re
from collections import namedtuple

import numpy as np

from nmreval import models
from nmreval.configs import config_paths
from nmreval.lib.importer import find_models, import_
from nmreval.utils import constants as constants
from nmreval.utils.text import convert

from ..Qt import QtWidgets, QtCore
from .._py.namespace_widget import Ui_Form


class Namespace:
    def __init__(self, fitfuncs=False, const=False, basic=False):
        self.namespace = {}
        self.groupings = {}
        self.top_levels = {}

        if basic:
            self.add_namespace({'x': (None, 'x values'), 'y': (None, 'x values'), 'y_err': (None, 'y error values'),
                                'fit': (None, 'dictionary of fit parameter', 'fit["PIKA"]'), 'np': (np, 'numpy module')},
                               parents=('Basic', 'General'))

            self.add_namespace({'sin': (np.sin, 'Sine', 'sin(PIKA)'), 'cos': (np.cos, 'Cosine', 'cos(PIKA)'),
                                'tan': (np.tan, 'Tangens', 'tan(PIKA)'), 'ln': (np.log, 'Natural Logarithm', 'ln(PIKA)'),
                                'log': (np.log10, 'Logarithm (base 10)', 'log(PIKA)'),
                                'exp': (np.exp, 'Exponential', 'exp(PIKA)'), 'sqrt': (np.sqrt, 'Root', 'sqrt(PIKA)'),
                                'lin_range': (np.linspace, 'N evenly spaced over interval [start, stop]',
                                              'lin_range(start, stop, N)'),
                                'log_range': (np.geomspace, 'N evenly spaced (log-scale) over interval [start, stop]',
                                              'lin_range(start, stop, N)')},
                               parents=('Basic', 'Functions'))
            
            self.add_namespace({'max': (np.max, 'Maximum value', 'max(PIKA)'),
                                'min': (np.min, 'Minimum value', 'min(PIKA)'),
                                'argmax': (np.argmax, 'Index of maximum value', 'argmax(PIKA)'),
                                'argmin': (np.argmax, 'Index of minimum value', 'argmin(PIKA)')},
                               parents=('Basic', 'Values'))

        if const:
            self.add_namespace({'e': (constants.e, 'e / As'), 'eps0': (constants.epsilon0, 'epsilon0 / As/Vm'),
                                'Eu': (constants.Eu,), 'h': (constants.h, 'h / eVs'),
                                'hbar': (constants.hbar, 'hbar / eVs'), 'kB': (constants.kB, 'kB / eV/K'),
                                'mu0': (constants.mu0, 'mu0 / Vs/Am'), 'NA': (constants.NA, 'NA / 1/mol'),
                                'pi': (constants.pi,), 'R': (constants.R, 'R / eV')},
                               parents=('Constants', 'Maybe useful'))

            self.add_namespace({f'gamma["{k}"]': (v, k, f'gamma["{k}"]') for k, v in constants.gamma.items()},
                               parents=('Constants', 'Magnetogyric ratios (in 1/(sT))'))

        if fitfuncs:
            self.make_dict_from_fitmodule(models)
        try:
            usermodels = import_(config_paths() / 'usermodels.py')
            self.make_dict_from_fitmodule(usermodels, parent='User-defined')
        except FileNotFoundError:
            pass

    def __str__(self):
        ret = '\n'.join([f'{k}: {v}' for k, v in self.groupings.items()])
        return ret

    def make_dict_from_fitmodule(self, module, parent='Fit function'):
        fitfunc_dict = {}
        for funcs in find_models(module):
            name = funcs.name
            if funcs.type not in fitfunc_dict:
                fitfunc_dict[funcs.type] = {}

            func_list = fitfunc_dict[funcs.type]

            func_args = 'x, ' + ', '.join(convert(p, old='latex', new='str', brackets=False) for p in funcs.params)

            # keywords arguments need names in function
            sig = inspect.signature(funcs.func)
            for p in sig.parameters.values():
                if p.default != p.empty:
                    func_args += ', ' + str(p)

            func_list[f"{funcs.__name__}"] = (funcs.func, name, f"{funcs.__name__}({func_args})")

        for k, v in fitfunc_dict.items():
            self.add_namespace(v, parents=(parent, k))

    def add_namespace(self, namespace, parents):
        if parents[0] in self.top_levels:
            if parents[1] not in self.top_levels[parents[0]]:
                self.top_levels[parents[0]].append(parents[1])
        else:
            self.top_levels[parents[0]] = [parents[1]]
        if (parents[0], parents[1]) not in self.groupings:
            self.groupings[(parents[0], parents[1])] = list(namespace.keys())
        else:
            self.groupings[(parents[0], parents[1])].extend(list(namespace.keys()))

        self.namespace.update(namespace)

    def flatten(self):
        ret_dic = {}
        graph = namedtuple('graphs', ['s'])
        sets = namedtuple('sets', ['x', 'y', 'y_err', 'value', 'fit'], defaults=(None,))

        gs = re.compile(r'g\[(\d+)].s\[(\d+)].(x|y(?:_err)*|value|fit)')

        for k, v in self.namespace.items():
            m = gs.match(k)
            if m:
                if 'g' not in ret_dic:
                    ret_dic['g'] = {}

                g_num, s_num, pos = m.groups()
                if int(g_num) not in ret_dic['g']:
                    ret_dic['g'][int(g_num)] = graph({})

                gg = ret_dic['g'][int(g_num)]
                if int(s_num) not in gg.s:
                    gg.s[int(s_num)] = sets('', '', '', '')

                ss = gg.s[int(s_num)]
                if pos == 'fit':
                    if ss.fit is None:
                        gg.s[int(s_num)] = ss._replace(**{pos: {}})
                        ss = gg.s[int(s_num)]
                    ss.fit[k[m.end()+2:-2]] = v[0]

                else:
                    ret_dic['g'][int(g_num)].s[int(s_num)] = ss._replace(**{pos: v[0]})

            else:
                ret_dic[k] = v[0]

        return ret_dic


class QNamespaceWidget(QtWidgets.QWidget, Ui_Form):
    selected = QtCore.pyqtSignal(str)
    sendKey = QtCore.pyqtSignal(str)

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

        self.setupUi(self)
        self.namespace = None

    def make_namespace(self):
        self.set_namespace(Namespace(fitfuncs=True, const=True, basic=True))

    @QtCore.pyqtSlot(int, name='on_groups_comboBox_currentIndexChanged')
    def show_subgroup(self, idx: int):
        name = self.groups_comboBox.itemText(idx)

        self.subgroups_comboBox.blockSignals(True)
        self.subgroups_comboBox.clear()
        for item in self.namespace.top_levels[name]:
            self.subgroups_comboBox.addItem(item)
        self.subgroups_comboBox.blockSignals(False)

        self.show_namespace(0)

    @QtCore.pyqtSlot(int, name='on_subgroups_comboBox_currentIndexChanged')
    def show_namespace(self, idx: int):
        subspace = self.namespace.groupings[(self.groups_comboBox.currentText(), self.subgroups_comboBox.itemText(idx))]

        self.namespace_table.clear()
        self.namespace_table.setRowCount(0)

        for entry in subspace:
            key_item = QtWidgets.QTableWidgetItem(entry)
            key_item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)

            vals = self.namespace.namespace[entry]

            if len(vals) < 3:
                alias = entry
            else:
                alias = vals[2]

            if len(vals) < 2:
                display = entry
            else:
                display = vals[1]

            value_item = QtWidgets.QTableWidgetItem(display)
            value_item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)

            key_item.setData(QtCore.Qt.UserRole, alias)
            key_item.setData(QtCore.Qt.UserRole+1, entry)
            value_item.setData(QtCore.Qt.UserRole, alias)
            value_item.setData(QtCore.Qt.UserRole+1, entry)

            row = self.namespace_table.rowCount()
            self.namespace_table.setRowCount(row+1)
            self.namespace_table.setItem(row, 0, key_item)
            self.namespace_table.setItem(row, 1, value_item)

        self.namespace_table.resizeColumnsToContents()

    def set_namespace(self, namespace):
        self.namespace = namespace

        self.groups_comboBox.blockSignals(True)
        self.groups_comboBox.clear()
        for k in self.namespace.top_levels.keys():
            self.groups_comboBox.addItem(k)
        self.groups_comboBox.blockSignals(False)

        self.show_subgroup(0)

    @QtCore.pyqtSlot(QtWidgets.QTableWidgetItem, name='on_namespace_table_itemDoubleClicked')
    def item_selected(self, item: QtWidgets.QTableWidgetItem):
        self.selected.emit(item.data(QtCore.Qt.UserRole))
        self.sendKey.emit(item.data(QtCore.Qt.UserRole+1))