nmreval/src/gui_qt/data/shift_graphs.py

import numpy as np
from itertools import cycle

from pyqtgraph import mkColor, mkPen

from nmreval.lib.colors import Tab10

from ..Qt import QtGui, QtCore, QtWidgets
from .._py.shift_scale_dialog import Ui_shift_dialog
from ..lib.pg_objects import PlotItem
from ..lib.spinboxes import SciSpinBox


class QShift(QtWidgets.QDialog, Ui_shift_dialog):
    valuesChanged = QtCore.pyqtSignal(dict, tuple)

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

        self.graphicsView.setMenuEnabled(False)
        self.splitter.setSizes([int(self.width()/4), int(self.width()*2/3)])

        self.movements = {}
        self.data = {}
        self._colors = cycle(Tab10)

        delegate = SpinBoxDelegate()
        delegate.valueChanged.connect(self.shift)
        self.shift_table.setItemDelegate(delegate)
        self.x_shift_spinbox.valueChanged.connect(lambda: self.glob_shift('h'))
        self.y_shift_spinbox.valueChanged.connect(lambda: self.glob_shift('v'))

        delegate = SpinBoxDelegate()
        delegate.valueChanged.connect(self.scale)
        self.scale_table.setItemDelegate(delegate)
        self.x_scale_spinbox.valueChanged.connect(lambda: self.glob_scale('h'))
        self.y_scale_spinbox.valueChanged.connect(lambda: self.glob_scale('v'))

    def add_item(self, idx, name, x, y):
        color = mkColor(next(self._colors).rgb())
        if np.iscomplexobj(y):
            pl = [PlotItem(x, y.real, name=name, pen=mkPen(color=color)),
                  PlotItem(x, y.imag, name=name, pen=mkPen(color=color))]
        else:
            pl = [PlotItem(x, y, name=name, pen=mkPen(color=color))]

        self.data[idx] = (pl, x, y)

        # [[horizontal shift, vertical shift], [horizontal scale, vertical scale]]
        self.movements[idx] = [[0, 0], [1, 1]]

        for i, tw in enumerate([self.shift_table, self.scale_table]):
            tw.blockSignals(True)
            row = tw.rowCount()
            tw.insertRow(row)

            item = QtWidgets.QTableWidgetItem(name)
            item.setForeground(QtGui.QBrush(color))
            item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsUserCheckable | QtCore.Qt.ItemIsEnabled)
            item.setCheckState(QtCore.Qt.Checked)
            item.setData(QtCore.Qt.UserRole, idx)
            tw.setItem(row, 0, item)

            tw.setItem(row, 1, QtWidgets.QTableWidgetItem(str(i)))
            tw.setItem(row, 2, QtWidgets.QTableWidgetItem(str(i)))

            tw.blockSignals(False)
        for i in pl:
            self.graphicsView.addItem(i)

    def set_graphs(self, graphs: list):
        for key, name in graphs:
            self.data_combobox.addItem(name, userData=key)
            self.values_combobox.addItem(name, userData=key)

    def glob_shift_scale(self, widget: QtWidgets.QTableWidget, mode: int, col: int, value: float):
        for row in range(widget.rowCount()):
            if widget.item(row, 0).checkState() == QtCore.Qt.Checked:
                item = widget.item(row, col)
                item.setText(str(value))
                self.shift_scale(widget, mode, row, col-1, value)

    def glob_shift(self, direction: str):
        if direction == 'h':
            val = self.x_shift_spinbox.value()
            self.glob_shift_scale(self.shift_table, 0, 1, val)
        else:
            val = self.y_shift_spinbox.value()
            self.glob_shift_scale(self.shift_table, 0, 2, val)

    def glob_scale(self, direction: str):
        if direction == 'h':
            val = self.x_scale_spinbox.value()
            self.glob_shift_scale(self.scale_table, 1, 1, val)
        else:
            val = self.y_scale_spinbox.value()
            self.glob_shift_scale(self.scale_table, 1, 2, val)

    def shift_scale(self, widget: QtWidgets.QTableWidget, mode: int,
                    row: int, col: int, value: float):
        item = widget.item(row, 0)
        key = item.data(QtCore.Qt.UserRole)
        self.movements[key][mode][col] = value

        (x_off, y_off), (x_scale, y_scale) = self.movements[key]

        pl, x, y = self.data[key]
        y_part = [np.real, np.imag]
        for i, item in enumerate(pl):
            item.setData(x=x*x_scale+x_off, y=y_part[i](y) * y_scale + y_off)

    @QtCore.pyqtSlot(int, int, float)
    def shift(self, row: int, column: int, value: float):
        self.shift_scale(self.shift_table, 0, row, column-1, value)

    @QtCore.pyqtSlot(int, int, float)
    def scale(self, row: int, column: int, value: float):
        self.shift_scale(self.scale_table, 1, row, column-1, value)

    @QtCore.pyqtSlot(int, name='on_xlog_checkbox_stateChanged')
    @QtCore.pyqtSlot(int, name='on_ylog_checkbox_stateChanged')
    def set_log(self, state: int):
        if self.sender() == self.xlog_checkbox:
            log_state = self.graphicsView.plotItem.ctrl.logXCheck
        else:
            log_state = self.graphicsView.plotItem.ctrl.logYCheck
        log_state.setCheckState(state)
        self.graphicsView.plotItem.updateLogMode()

    def on_overwrite_checkbox_stateChanged(self, state: int):
        self.data_newgraph.setVisible(state != QtCore.Qt.Checked)
        self.data_combobox.setVisible(state != QtCore.Qt.Checked)

    def on_value_checkbox_stateChanged(self, state: int):
        self.values_newgraph.setVisible(state == QtCore.Qt.Checked)
        self.values_combobox.setVisible(state == QtCore.Qt.Checked)

    def on_data_newgraph_stateChanged(self, state: int):
        self.data_combobox.setEnabled(state != QtCore.Qt.Checked)

    def on_values_newgraph_stateChanged(self, state: int):
        self.values_combobox.setEnabled(state != QtCore.Qt.Checked)

    def accept(self):
        data_saving = None
        if not self.overwrite_checkbox.isChecked():
            if self.data_newgraph.isChecked():
                data_saving = ''
            else:
                data_saving = self.data_combobox.currentData()

        value_saving = None
        if self.value_checkbox.isChecked():
            if self.values_newgraph.isChecked():
                value_saving = ''
            else:
                value_saving = self.values_combobox.currentData()

        self.valuesChanged.emit(self.movements, (data_saving, value_saving))
        self.close()


class SpinBoxDelegate(QtWidgets.QStyledItemDelegate):
    valueChanged = QtCore.pyqtSignal(int, int, float)

    def createEditor(self, parent: QtWidgets.QWidget, option: QtWidgets.QStyleOptionViewItem,
                     idx: QtCore.QModelIndex) -> QtWidgets.QWidget:
        editor = SciSpinBox(parent)
        editor.valueChanged.connect(self.new_value)

        return editor

    def new_value(self, val):
        # geht bestimmt besser...
        table = self.sender().parent().parent()
        self.valueChanged.emit(table.currentRow(), table.currentColumn(), val)