Source code for prsctrl_gui.gui.widgets.data.prsdata_plot

from PyQt6.QtGui import QPalette
from PyQt6.QtCore import Qt
from PyQt6.QtWidgets import QCheckBox, QLabel, QScrollArea, QWidget, QHBoxLayout, QVBoxLayout, QSplitter
from pyqtgraph import PlotWidget, PlotDataItem, GraphicsLayoutWidget
from prsctrl.data import PrsData
from typing import Optional

import logging
log = logging.getLogger(__name__)



class PrsDataWidget(QWidget):
    """Widget for showing a :class:`PrsData` instance.

    The widget shows a list of y-column checkboxes on the left side. The plot for each checked column is shown on 
    the right on a shared x-axis.
    """
    DEFAULT_CHECK_COLUMNS = ["dI_I", "I"]
    """Default columns to enabled the plots for."""



    def __init__(self, data: PrsData, default_check_columns: Optional[list[str]]=None, **kwargs):
        """
        :param data: :class:`.PrsData` instance to display.
        :param default_check_columns: The default plots to show after the widget is created. Defaults to :member:`.DEFAULT_CHECK_COLUMNS`.
        """
        QWidget.__init__(self, **kwargs)
        self._data = data
        self.l_main = QHBoxLayout()
        self.setLayout(self.l_main)
        self.l_left = QVBoxLayout()
        self.l_main.addLayout(self.l_left, 1)
        self.l_checkboxes = QVBoxLayout()
        self.l_left.addLayout(self.l_checkboxes)

        self._default_check_columns = default_check_columns if default_check_columns is not None else self.DEFAULT_CHECK_COLUMNS

        self.checkboxes: list[tuple[str, QCheckBox]] = []
        self._update_column_selector()

        self.w_plots = GraphicsLayoutWidget()
        self.l_main.addWidget(self.w_plots, 4)
        self.update_plots()




    def _update_column_selector(self):
        """Update the checkbox list of data columns depending on which are available."""
        while self.l_checkboxes.count() > 0:
            item = self.l_checkboxes.itemAt(0)
            self.l_checkboxes.removeItem(item)
            w = item.widget()
            if w:
                w.setParent(None)
                w.deleteLater()

        if not self._data.has_wavelengths():
            self.l_checkboxes.addWidget(QLabel(self.tr("No data available", "No data points in data object")))
            self.l_checkboxes.addStretch(1)
            return
        ycolumns = []
        if self._data.has_AC():
            ycolumns += self._data.COLUMN_KEYS_AC
        if self._data.has_DC():
            ycolumns += self._data.COLUMN_KEYS_DC
        if len(ycolumns) == 0:
            self.l_checkboxes.addWidget(QLabel(self.tr("No data columns available", "Data points exist, but no columns available")))
            self.l_checkboxes.addStretch(1)
            return
        for ycol in ycolumns:
            name = self._data.get_column_label(ycol)
            w_box = QCheckBox(name)
            if ycol in self._default_check_columns:
                w_box.setChecked(True)
            w_box.stateChanged.connect(lambda _: self.update_plots())
            self.checkboxes.append((ycol, w_box))
            self.l_checkboxes.addWidget(w_box)
        self.l_checkboxes.addStretch()




    def _get_selected_columns(self) -> list[str]:
        selected = []
        for ycol, w_box in self.checkboxes:
            if w_box.isChecked():
                selected.append(ycol)
        return selected




    def update_plots(self):
        """Recreate plots for checked y-labels."""
        self.w_plots.clear()
        # Update columns if the first data point was added (required for live measurements, where the initial data object is empty.)
        # if len(self._data.data) == 1:
        # Update columns every time there appear to be too few columns. First data point might offset and thus no data available
        if self.l_checkboxes.count() < 3:
            self._update_column_selector()
        ycols = self._get_selected_columns()
        log.debug(f"Updating plot: Checked columns are {ycols}")
        if len(ycols) == 0:
            return

        get_spectrum_kw = {}

        columns = ["E", "wl"] + ycols
        sdata, columns = self._data.get_spectrum_data(columns=columns, **get_spectrum_kw)

        xkey = "E"
        x = sdata[:,columns.index(xkey)]

        name1 = ycols[0]

        for i, ycol in enumerate(ycols):
            y = sdata[:,columns.index(ycol)]
            plotItem = self.w_plots.addPlot(x=x, y=y, name=ycol, row=i, col=0)
            if i != 0:
                plotItem.setXLink(name1)

            # labels
            if i == len(ycols) - 1:
                xlabel = self._data.get_column_label(xkey)
                xunit = self._data.get_column_unit(xkey)
                plotItem.setLabel("bottom", xlabel, xunit)
            ylabel = self._data.get_column_label(ycol)
            yunit = self._data.get_column_unit(ycol)
            plotItem.setLabel("left", ylabel, yunit)
            plotItem.addLine(y=0)