slsDetectorPackage/pyctbgui/pyctbgui/services/Signals.py

from functools import partial
from pathlib import Path

import numpy as np
from PyQt5 import QtWidgets, uic
import pyqtgraph as pg
from pyqtgraph import LegendItem

from pyctbgui.utils.bit_utils import bit_is_set, manipulate_bit
from pyctbgui.utils.defines import Defines
from pyctbgui.utils.recordOrApplyPedestal import recordOrApplyPedestal


class SignalsTab(QtWidgets.QWidget):

    def __init__(self, parent):
        super().__init__(parent)
        pg.setConfigOptions(imageAxisOrder="row-major")  
        uic.loadUi(Path(__file__).parent.parent / 'ui' / "signals.ui", parent)
        self.view = parent
        self.mainWindow = None
        self.det = None
        self.plotTab = None
        self.legend: LegendItem | None = None
        self.rx_dbitoffset = None
        self.rx_dbitreorder = None
        self.rx_dbitlist = None

    def refresh(self):
        self.updateSignalNames()
        self.updateDigitalBitEnable()
        self.updateIOOut()
        self.getDBitOffset()
        self.getDBitReorder()

    def connect_ui(self):
        for i in range(Defines.signals.count):
            getattr(self.view, f"checkBoxBIT{i}DB").stateChanged.connect(partial(self.setDigitalBitEnable, i))
            getattr(self.view, f"checkBoxBIT{i}Out").stateChanged.connect(partial(self.setIOOut, i))
            getattr(self.view, f"checkBoxBIT{i}Plot").stateChanged.connect(partial(self.setBitPlot, i))
            getattr(self.view, f"pushButtonBIT{i}").clicked.connect(partial(self.selectBitColor, i))
        self.view.checkBoxBIT0_31DB.stateChanged.connect(
            partial(self.setDigitalBitEnableRange, 0, Defines.signals.half))
        self.view.checkBoxBIT32_63DB.stateChanged.connect(
            partial(self.setDigitalBitEnableRange, Defines.signals.half, Defines.signals.count))
        self.view.checkBoxBIT0_31Plot.stateChanged.connect(partial(self.setBitPlotRange, 0,
                                                                   Defines.signals.half))
        self.view.checkBoxBIT32_63Plot.stateChanged.connect(
            partial(self.setBitPlotRange, Defines.signals.half, Defines.signals.count))
        self.view.checkBoxBIT0_31Out.stateChanged.connect(partial(self.setIOOutRange, 0, Defines.signals.half))
        self.view.checkBoxBIT32_63Out.stateChanged.connect(
            partial(self.setIOOutRange, Defines.signals.half, Defines.signals.count))
        self.view.lineEditPatIOCtrl.editingFinished.connect(self.setIOOutReg)
        self.view.spinBoxDBitOffset.editingFinished.connect(self.setDbitOffset)
        self.view.checkBoxDBitReorder.stateChanged.connect(self.setDbitReorder)

    def setup_ui(self):
        self.plotTab = self.mainWindow.plotTab

        for i in range(Defines.signals.count):
            self.setDBitButtonColor(i, self.plotTab.getRandomColor())

        self.initializeAllDigitalPlots()

        self.legend = self.mainWindow.plotDigitalWaveform.getPlotItem().legend
        self.legend.clear()
        # subscribe to toggle legend
        self.plotTab.subscribeToggleLegend(self.updateLegend)

    def getEnabledPlots(self):
        """
        return plots that are shown (checkBoxTransceiver{i}Plot is checked)
        """
        enabledPlots = []
        self.legend.clear()
        for i in range(Defines.signals.count):
            if getattr(self.view, f'checkBoxBIT{i}Plot').isChecked():
                plotName = getattr(self.view, f"labelBIT{i}").text()
                enabledPlots.append((self.mainWindow.digitalPlots[i], plotName))
        return enabledPlots

    def updateLegend(self):
        """
        update the legend for the signals waveform plot
        should be called after checking or unchecking plot checkbox
        """
        if not self.mainWindow.showLegend:
            self.legend.clear()
        else:
            for plot, name in self.getEnabledPlots():
                self.legend.addItem(plot, name)

    def updatePlotRange(self): 
        vb = self.mainWindow.plotDigitalWaveform.getViewBox()
        vb.enableAutoRange(enable=True)  # Enable auto-range
        vb.updateAutoRange()  # Force immediate update
        
    @recordOrApplyPedestal
    def _processWaveformData(self, data, aSamples, dSamples, rx_dbitreorder, rx_dbitlist, romode,
                             nADCEnabled):
        
        #transform raw waveform data into a processed numpy array
        #@param data:  raw waveform data
    
        start_digital_data = 0 
        if romode == 2:
            start_digital_data += nADCEnabled * 2 * aSamples
        digital_array = np.array(np.frombuffer(data, offset=start_digital_data, dtype=np.uint8)) 
        if rx_dbitreorder:
            samples_per_bit = np.empty((len(rx_dbitlist), dSamples), dtype=np.uint8) #stored per row all the corresponding signals of all samples 
            nbitsPerDBit = dSamples
            if nbitsPerDBit % 8 != 0:
                nbitsPerDBit += (8 - (dSamples % 8))
            bit_index = 0
            for idx, i in enumerate(rx_dbitlist):
                # where numbits * numsamples is not a multiple of 8
                if bit_index % 8 != 0:
                    bit_index += (8 - (bit_index % 8))
            
                for iSample in range(dSamples):
                    # all samples for digital bit together from slsReceiver
                    index = int(bit_index / 8)
                    iBit = bit_index % 8
                    bit = (digital_array[index] >> iBit) & 1
                    samples_per_bit[idx,iSample] = bit
                    bit_index += 1
            return samples_per_bit
        else: 
            nbitsPerSample = len(rx_dbitlist) if len(rx_dbitlist) % 8 == 0 else len(rx_dbitlist) + (8 - (len(rx_dbitlist) % 8))
            bits_per_sample = np.empty((dSamples, len(rx_dbitlist)), dtype=np.uint8) #store per row all selected bits of a sample
            for iSample in range(dSamples): 
                bit_index = nbitsPerSample * iSample
                for idx, i in enumerate(rx_dbitlist): 
                    index = int(bit_index/8)
                    iBit = idx % 8
                    bit = (digital_array[index] >> iBit) & 1
                    bits_per_sample[iSample, idx] = bit
                    bit_index += 1
            return bits_per_sample.T.copy() 


                
    def processWaveformData(self, data, aSamples, dSamples):
        
        #view function
        #plots processed waveform data
        #data: raw waveform data
        #dsamples: digital samples
        #asamples: analog samples

        self.refresh()
        
        waveforms = {}

        digital_array = self._processWaveformData(data, aSamples, dSamples, self.rx_dbitreorder, self.rx_dbitlist, 
                                                  self.mainWindow.romode.value,
                                                  self.mainWindow.nADCEnabled)

        irow = 0
        for idx, i in enumerate(self.rx_dbitlist): 
            waveform = digital_array[idx, :]
            self.mainWindow.digitalPlots[i].setData(waveform)
            plotName = getattr(self.view, f"labelBIT{i}").text()
            waveforms[plotName] = waveform
            # TODO: left axis does not show 0 to 1, but keeps increasing
            if self.plotTab.view.radioButtonStripe.isChecked():
                self.mainWindow.digitalPlots[i].setY(irow * 2)
                irow += 1
            else:
                self.mainWindow.digitalPlots[i].setY(0)

        self.updatePlotRange()  # Call after all data is set
        
        return waveforms


    def initializeAllDigitalPlots(self):
        self.mainWindow.plotDigitalWaveform = pg.plot()
        self.mainWindow.plotDigitalWaveform.addLegend(colCount=Defines.colCount)
        self.mainWindow.verticalLayoutPlot.addWidget(self.mainWindow.plotDigitalWaveform, 3)
        self.mainWindow.digitalPlots = {}
        waveform = np.zeros(1000)
        for i in range(Defines.signals.count):
            pen = pg.mkPen(color=self.getDBitButtonColor(i), width=1)
            legendName = getattr(self.view, f"labelBIT{i}").text()
            self.mainWindow.digitalPlots[i] = self.mainWindow.plotDigitalWaveform.plot(waveform,
                                                                                       pen=pen,
                                                                                       name=legendName,
                                                                                       stepMode="left")
            self.mainWindow.digitalPlots[i].hide()

        self.mainWindow.plotDigitalImage = pg.ImageView()
        self.mainWindow.nDigitalRows = 0
        self.mainWindow.nDigitalCols = 0
        self.mainWindow.digital_frame = np.zeros((self.mainWindow.nDigitalRows, self.mainWindow.nDigitalCols))
        self.mainWindow.plotDigitalImage.setImage(self.mainWindow.digital_frame)
        self.mainWindow.verticalLayoutPlot.addWidget(self.mainWindow.plotDigitalImage, 4)

    def updateSignalNames(self):
        for i, name in enumerate(self.det.getSignalNames()):
            getattr(self.view, f"labelBIT{i}").setText(name)

    def getDigitalBitEnable(self, i, dbitList):
        checkBox = getattr(self.view, f"checkBoxBIT{i}DB")
        checkBox.stateChanged.disconnect()
        checkBox.setChecked(i in list(dbitList))
        checkBox.stateChanged.connect(partial(self.setDigitalBitEnable, i))

    def updateDigitalBitEnable(self):
        retval = self.det.rx_dbitlist
        self.rx_dbitlist = list(retval)
        self.mainWindow.nDBitEnabled = len(list(retval))
        for i in range(Defines.signals.count):
            self.getDigitalBitEnable(i, retval)
            self.EnableBitPlot(i)
            self.getEnableBitColor(i)
            self.plotTab.addSelectedDigitalPlots(i)
        self.getDigitalBitEnableRange(retval)
        self.EnableBitPlotRange()

    def setDigitalBitEnable(self, i):
        bitList = self.det.rx_dbitlist
        checkBox = getattr(self.view, f"checkBoxBIT{i}DB")
        if checkBox.isChecked():
            bitList.append(i)
        else:
            bitList.remove(i)
        self.det.rx_dbitlist = bitList

        self.updateDigitalBitEnable()

    def getDigitalBitEnableRange(self, dbitList):
        self.view.checkBoxBIT0_31DB.stateChanged.disconnect()
        self.view.checkBoxBIT32_63DB.stateChanged.disconnect()
        self.view.checkBoxBIT0_31DB.setChecked(all(x in list(dbitList) for x in range(Defines.signals.half)))
        self.view.checkBoxBIT32_63DB.setChecked(
            all(x in list(dbitList) for x in range(Defines.signals.half, Defines.signals.count)))
        self.view.checkBoxBIT0_31DB.stateChanged.connect(
            partial(self.setDigitalBitEnableRange, 0, Defines.signals.half))
        self.view.checkBoxBIT32_63DB.stateChanged.connect(
            partial(self.setDigitalBitEnableRange, Defines.signals.half, Defines.signals.count))

    def setDigitalBitEnableRange(self, start_nr, end_nr):
        bitList = self.det.rx_dbitlist
        checkBox = getattr(self.view, f"checkBoxBIT{start_nr}_{end_nr - 1}DB")
        for i in range(start_nr, end_nr):
            if checkBox.isChecked():
                if i not in list(bitList):
                    bitList.append(i)
            else:
                if i in list(bitList):
                    bitList.remove(i)
        self.det.rx_dbitlist = bitList

        self.updateDigitalBitEnable()

    def EnableBitPlot(self, i):
        """ enables plot check box if bit is enabled, otherwise unchecks and disables plot check box """

        checkBox = getattr(self.view, f"checkBoxBIT{i}DB")
        checkBoxPlot = getattr(self.view, f"checkBoxBIT{i}Plot")
        if(checkBoxPlot.isChecked()): 
            checkBoxPlot.setChecked(checkBox.isChecked())

        checkBoxPlot.setEnabled(checkBox.isChecked())

    def setBitPlot(self, i): 
        """ sets plot check box e.g. adds plots to plot tab """ 
        pushButton = getattr(self.view, f"pushButtonBIT{i}")
        checkBox = getattr(self.view, f"checkBoxBIT{i}Plot")
        pushButton.setEnabled(checkBox.isChecked())

        self.EnableBitPlotRange()
        self.plotTab.addSelectedDigitalPlots(i)
        self.updateLegend()
        self.updatePlotRange()
        

    def EnableBitPlotRange(self):
        self.view.checkBoxBIT0_31Plot.stateChanged.disconnect()
        self.view.checkBoxBIT32_63Plot.stateChanged.disconnect()
        self.view.checkBoxBIT0_31Plot.setEnabled(
            all(getattr(self.view, f"checkBoxBIT{i}Plot").isEnabled() for i in range(Defines.signals.half)))
        self.view.checkBoxBIT32_63Plot.setEnabled(
            all(
                getattr(self.view, f"checkBoxBIT{i}Plot").isEnabled()
                for i in range(Defines.signals.half, Defines.signals.count)))
        self.view.checkBoxBIT0_31Plot.setChecked(
            all(getattr(self.view, f"checkBoxBIT{i}Plot").isChecked() for i in range(Defines.signals.half)))
        self.view.checkBoxBIT32_63Plot.setChecked(
            all(
                getattr(self.view, f"checkBoxBIT{i}Plot").isChecked()
                for i in range(Defines.signals.half, Defines.signals.count)))
        self.view.checkBoxBIT0_31Plot.stateChanged.connect(partial(self.setBitPlotRange, 0,
                                                                   Defines.signals.half))
        self.view.checkBoxBIT32_63Plot.stateChanged.connect(
            partial(self.setBitPlotRange, Defines.signals.half, Defines.signals.count))

    def setBitPlotRange(self, start_nr, end_nr):
        checkBox = getattr(self.view, f"checkBoxBIT{start_nr}_{end_nr - 1}Plot")
        enable = checkBox.isChecked()
        for i in range(start_nr, end_nr):
            checkBox = getattr(self.view, f"checkBoxBIT{i}Plot")
            checkBox.setChecked(enable)
        self.plotTab.addAllSelectedDigitalPlots()

    def getEnableBitColor(self, i):
        checkBox = getattr(self.view, f"checkBoxBIT{i}Plot")
        pushButton = getattr(self.view, f"pushButtonBIT{i}")
        pushButton.setEnabled(checkBox.isEnabled() and checkBox.isChecked())

    def selectBitColor(self, i):
        pushButton = getattr(self.view, f"pushButtonBIT{i}")
        self.plotTab.showPalette(pushButton)
        pen = pg.mkPen(color=self.getDBitButtonColor(i), width=1)
        self.mainWindow.digitalPlots[i].setPen(pen)

    def getDBitButtonColor(self, i):
        pushButton = getattr(self.view, f"pushButtonBIT{i}")
        return self.plotTab.getActiveColor(pushButton)

    def setDBitButtonColor(self, i, color):
        pushButton = getattr(self.view, f"pushButtonBIT{i}")
        return self.plotTab.setActiveColor(pushButton, color)

    def getIOOutReg(self):
        retval = self.det.patioctrl
        self.view.lineEditPatIOCtrl.editingFinished.disconnect()
        self.view.lineEditPatIOCtrl.setText("0x{:016x}".format(retval))
        self.view.lineEditPatIOCtrl.editingFinished.connect(self.setIOOutReg)
        return retval

    def setIOOutReg(self):
        self.view.lineEditPatIOCtrl.editingFinished.disconnect()
        try:
            self.det.patioctrl = int(self.view.lineEditPatIOCtrl.text(), 16)
        except Exception as e:
            QtWidgets.QMessageBox.warning(self.mainWindow, "IO Out Fail", str(e), QtWidgets.QMessageBox.Ok)
            pass
        # TODO: handling double event exceptions
        self.view.lineEditPatIOCtrl.editingFinished.connect(self.setIOOutReg)
        self.updateIOOut()

    def updateCheckBoxIOOut(self, i, out):
        checkBox = getattr(self.view, f"checkBoxBIT{i}Out")
        checkBox.stateChanged.disconnect()
        checkBox.setChecked(bit_is_set(out, i))
        checkBox.stateChanged.connect(partial(self.setIOOut, i))

    def updateIOOut(self):
        retval = self.getIOOutReg()
        for i in range(Defines.signals.count):
            self.updateCheckBoxIOOut(i, retval)
        self.getIOoutRange(retval)

    def setIOOut(self, i):
        out = self.det.patioctrl
        checkBox = getattr(self.view, f"checkBoxBIT{i}Out")
        mask = manipulate_bit(checkBox.isChecked(), out, i)
        self.det.patioctrl = mask

        retval = self.getIOOutReg()
        self.updateCheckBoxIOOut(i, retval)
        self.getIOoutRange(retval)

    def getIOoutRange(self, out):
        self.view.checkBoxBIT0_31Out.stateChanged.disconnect()
        self.view.checkBoxBIT32_63Out.stateChanged.disconnect()
        self.view.checkBoxBIT0_31Out.setChecked((out & Defines.signals.BIT0_31_MASK) == Defines.signals.BIT0_31_MASK)
        self.view.checkBoxBIT32_63Out.setChecked((out
                                                  & Defines.signals.BIT32_63_MASK) == Defines.signals.BIT32_63_MASK)
        self.view.checkBoxBIT0_31Out.stateChanged.connect(partial(self.setIOOutRange, 0, Defines.signals.half))
        self.view.checkBoxBIT32_63Out.stateChanged.connect(
            partial(self.setIOOutRange, Defines.signals.half, Defines.signals.count))

    def setIOOutRange(self, start_nr, end_nr):
        out = self.det.patioctrl
        checkBox = getattr(self.view, f"checkBoxBIT{start_nr}_{end_nr - 1}Out")
        mask = getattr(Defines.signals, f"BIT{start_nr}_{end_nr - 1}_MASK")
        if checkBox.isChecked():
            self.det.patioctrl = out | mask
        else:
            self.det.patioctrl = out & ~mask
        self.updateIOOut()

    def getDBitOffset(self):
        self.view.spinBoxDBitOffset.editingFinished.disconnect()
        self.rx_dbitoffset = self.det.rx_dbitoffset
        self.view.spinBoxDBitOffset.setValue(self.rx_dbitoffset)
        self.view.spinBoxDBitOffset.editingFinished.connect(self.setDbitOffset)

    def getDBitReorder(self):
        self.view.checkBoxDBitReorder.stateChanged.disconnect()
        self.rx_dbitreorder = self.det.rx_dbitreorder
        self.view.checkBoxDBitReorder.setChecked(self.rx_dbitreorder)
        self.view.checkBoxDBitReorder.stateChanged.connect(self.setDbitReorder)


    def setDbitOffset(self):
        self.det.rx_dbitoffset = self.view.spinBoxDBitOffset.value()

    def setDbitReorder(self):
        self.det.rx_dbitreorder = self.view.checkBoxDBitReorder.isChecked()

    def saveParameters(self) -> list:
        commands = []
        dblist = [str(i) for i in range(Defines.signals.count) if getattr(self.view, f"checkBoxBIT{i}DB").isChecked()]
        if len(dblist) > 0:
            commands.append(f"rx_dbitlist {', '.join(dblist)}")
        commands.append(f"rx_dbitoffset {self.view.spinBoxDBitOffset.value()}")
        commands.append(f"rx_dbitreorder {self.view.checkBoxDBitReorder.isChecked()}")
        commands.append(f"patioctrl {self.view.lineEditPatIOCtrl.text()}")
        return commands