nmreval/src/gui_qt/data/signaledit/phase_dialog.py

from __future__ import annotations

import numpy as np
from pyqtgraph import mkPen
from numpy import inf, linspace
from numpy.fft import fft, fftfreq, fftshift

from ...lib.pg_objects import PlotItem, LogInfiniteLine
from nmreval.lib.importer import find_models
from nmreval.math import apodization as apodization
from nmreval.utils.text import convert

from ...Qt import QtCore, QtWidgets
from ..._py.apod_dialog import Ui_ApodEdit
from ..._py.phase_corr_dialog import Ui_SignalEdit
from ...lib.forms import FormWidget


class QPreviewDialogs(QtWidgets.QDialog):
    finished = QtCore.pyqtSignal(str, tuple)

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

        self.data = []
        self.graphs = []

        self.mode = ''

    def setRange(self, xlim: list, ylim: list, logmode: list[bool]):
        self.graphicsView.getPlotItem().setLogMode(x=logmode[0], y=logmode[1])
        if logmode[0]:
            xlim = [np.log10(x) for x in xlim]
        if logmode[1]:
            ylim = [np.log10(y) for y in ylim]

        self.graphicsView.setRange(xRange=xlim, yRange=ylim, padding=0, disableAutoRange=True)

    def add_data(self, x, y):
        self.data.append((x, y))
        real_plt = PlotItem(x=x, y=y.real, pen=mkPen('b'), )
        imag_plt = PlotItem(x=x, y=y.imag, pen=mkPen('r'))
        self.graphs.append((real_plt, imag_plt))
        self.graphicsView.addItem(real_plt)
        self.graphicsView.addItem(imag_plt)

    def done(self, val):
        self.cleanup()
        super().done(val)

    def close(self):
        self.cleanup()
        super().close()

    def accept(self):
        self.finished.emit(self.mode, self.get_value())
        super().accept()

    def get_value(self):
        raise NotImplementedError

    def cleanup(self):
        self.blockSignals(True)

        for line in self.graphs:
            for g in line:
                self.graphicsView.removeItem(g)
                del g

        self.graphicsView.clear()

        self.data = []
        self.graphs = []

        self.blockSignals(False)


class QPhasedialog(QPreviewDialogs, Ui_SignalEdit):
    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)

        self.mode = 'ph'

        self.pvt_line = LogInfiniteLine(pos=0, movable=True)
        self.graphicsView.addItem(self.pvt_line)
        self.pvt_line.sigPositionChanged.connect(self.move_line)

    @QtCore.pyqtSlot(float, name='on_ph1slider_valueChanged')
    @QtCore.pyqtSlot(float, name='on_ph0slider_valueChanged')
    def _temp_phase(self, *args):
        ph0, ph1, pvt = self.get_value()
        self.pvt_line.setValue(pvt)

        for i, (x, y) in enumerate(self.data):
            phasecorr = np.exp(-1j * (ph0 + ph1*(x-pvt)/np.max(x))*np.pi/180.)
            _y = y * phasecorr

            self.graphs[i][0].setData(x=x, y=_y.real)
            self.graphs[i][1].setData(x=x, y=_y.imag)

    def get_value(self):
        return float(self.ph0slider.text()), float(self.ph1slider.text()), float(self.pivot_lineedit.text())

    def move_line(self, evt):
        self.pivot_lineedit.setText(f'{evt.value():.5g}')


class QApodDialog(QPreviewDialogs, Ui_ApodEdit):
    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)

        self._limits = (-inf, inf), -inf

        self.apods = []
        self.apods = find_models(apodization)

        self.apodcombobox.blockSignals(True)
        for ap in self.apods:
            self.apodcombobox.addItem(ap().name)
        self.apodcombobox.blockSignals(False)

        self.apod_graph = PlotItem(x=[], y=[])
        self.graphicsView.addItem(self.apod_graph)

        self.mode = 'ap'

        self.change_apodization(0)

    def add_data(self, x, y):
        real_plt = PlotItem(x=x, y=y.real, pen=mkPen('b'))
        # imag_plt = (x=x, y=y.imag, pen=pg.mkPen('r'))
        self.graphicsView.addItem(real_plt)
        # self.graphicsView.addItem(imag_plt)

        y_fft = fftshift(fft(y))
        x_fft = fftshift(fftfreq(len(x), d=x[1]-x[0]))
        real_plt_fft = PlotItem(x=x_fft, y=y_fft.real, pen=mkPen('b'))
        # imag_plt_fft = pg.PlotDataItem(x=x_fft, y=y_fft.imag, pen=pg.mkPen('b'))
        self.graphicsView_2.addItem(real_plt_fft)
        # self.graphicsView_2.addItem(imag_plt_fft)

        self.graphs.append((real_plt, real_plt_fft))
        self.data.append((x, y, x_fft))

        xlimits = (max(x.min(), self._limits[0][0]), min(x.max(), self._limits[0][1]))
        ylimit = max(self._limits[1], y.real.max())
        self._limits = xlimits, ylimit

    @QtCore.pyqtSlot(int, name='on_apodcombobox_currentIndexChanged')
    def change_apodization(self, index):
        # delete old widgets
        self.eqn_label.setText(convert(self.apods[index].equation))
        while self.widget_layout.count():
            item = self.widget_layout.takeAt(0)
            try:
                item.widget().deleteLater()
            except AttributeError:
                pass

        # set up parameter widgets for new model
        for k, v in enumerate(self.apods[index]().params):
            widgt = FormWidget(name=v)
            widgt.valueChanged.connect(self._temp_apod)
            self.widget_layout.addWidget(widgt)

        self.widget_layout.addStretch()
        self._temp_apod()

    def _temp_apod(self):
        apodmodel = self.apods[self.apodcombobox.currentIndex()]
        p = self._get_parameter()

        if self.data:
            for i, (x, y, x_fft) in enumerate(self.data):
                y2 = apodmodel.apod(x, *p)
                _y = y2 * y
                self.graphs[i][0].setData(x=x, y=_y.real)
                # self.graphs[i][1].setData(y=_y.imag)
                y_fft = fftshift(fft(_y))
                self.graphs[i][1].setData(x=x_fft, y=y_fft.real)
                # self.graphs[i][3].setData(y=y_fft.imag)

        _x_apod = linspace(self._limits[0][0], self._limits[0][1])
        try:
            _y_apod = apodmodel.apod(_x_apod, *p)
            self.apod_graph.setData(x=_x_apod, y=self._limits[1]*_y_apod)
        except IndexError:
            pass

    def _get_parameter(self):
        p = []
        for i in range(self.widget_layout.count()):
            item = self.widget_layout.itemAt(i)
            w = item.widget()
            try:
                p.append(w.value)
            except AttributeError:
                continue

        return p

    def get_value(self):
        apodmodel = self.apods[self.apodcombobox.currentIndex()]
        p = self._get_parameter()

        return p, apodmodel