TraceDemo/src/gui/widgets/plot.py

from typing import Optional

import pandas as pd
from PyQt5.QtWidgets import QWidget, QVBoxLayout
import pyqtgraph as pg
import numpy as np

from base.base import BasePlotWidget


class PlotWidget(BasePlotWidget):

    def _create_stage_region(self,
                             stage: str,
                             times: pd.Series,
                             dataframe: pd.DataFrame) -> Optional[pg.LinearRegionItem]:
        stage_diff = np.diff(dataframe[stage])
        start_index = np.where(stage_diff == 1)[0]
        finish_index = np.where(stage_diff == -1)[0]

        if start_index.size:
            start_timestamp = times[start_index[0]]
            finish_timestamp = times[finish_index[0]] if finish_index.size else times[len(times) - 1]
            region = pg.LinearRegionItem([start_timestamp, finish_timestamp], movable=False)
            region.setBrush(pg.mkBrush(self._stage_colors[stage]))
            return region
        return None

    def _find_mid_in_stage(self,
                           stage: str,
                           dataframe: pd.DataFrame,
                           signal_name: str) -> float:
        if stage in self._stages:
            stage_diff = np.diff(dataframe[stage])
            start_index = np.where(stage_diff == 1)[0]
            finish_index = np.where(stage_diff == -1)[0]

            data = dataframe[signal_name] if signal_name in dataframe.columns.tolist() else []

            if data.size and start_index.size:
                start = start_index[0]
                finish = finish_index[0] if finish_index.size else (len(data) - 1)
                data_slice = data[start:finish]
                mean = np.mean(data_slice)
                return mean
            return 0.0

    def _build_widget(self, dataframe: pd.DataFrame) -> QWidget:
        widget = QWidget()
        layout = QVBoxLayout()

        time_axis = dataframe["time"]
        dataframe_headers = dataframe.columns.tolist()

        for channel, description in self._plt_channels.items():
            plot_widget = pg.PlotWidget(title=channel)
            plot_widget.showGrid(x=True, y=True)
            legend = pg.LegendItem((80, 60), offset=(70, 20))
            legend.setParentItem(plot_widget.graphicsItem())
            settings = description["Settings"]

            if settings["stages"] and all([stage in dataframe_headers for stage in self._stages]):
                for stage in self._stages:
                    region = self._create_stage_region(stage, time_axis, dataframe)
                    if region:
                        plot_widget.addItem(region)

            if settings["zoom"] and max(time_axis) < 5.0:
                means_raw = [self._find_mid_in_stage("Welding",
                                                     dataframe,
                                                     signal["name"]) for signal in description["Signals"]]
                mean = max(means_raw)
                region_stable = pg.LinearRegionItem([mean - mean * 0.015,
                                                     mean + mean * 0.015],
                                                    orientation='horizontal')
                region_stable.setBrush(pg.mkBrush([250, 250, 0, 25]))
                plot_widget.addItem(region_stable)

                region_overshoot = pg.LinearRegionItem([mean - mean * 0.05,
                                                        mean + mean * 0.05],
                                                       orientation='horizontal')
                region_overshoot.setBrush(pg.mkBrush([0, 250, 0, 30]))
                plot_widget.addItem(region_overshoot)

                plot_widget.setYRange(mean - 260, mean + 260)
                plot_widget.setInteractive(False)

            for signal in description["Signals"]:
                if signal["name"] in dataframe_headers:
                    plot = plot_widget.plot(time_axis, dataframe[signal["name"]], pen=signal["pen"])
                    legend.addItem(plot, signal["name"])

            layout.addWidget(plot_widget)

        widget.setLayout(layout)
        return widget

    def build(self, data: list[pd.DataFrame]) -> None:
        widgets = [self._build_widget(data_sample) for data_sample in data]
        self._mediator.notify(self, widgets)