TraceDemo/src/gui/widgets/plot.py

from typing import Optional, Any, NamedTuple

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

from base.base import BasePlotWidget


class ProcessStage(NamedTuple):
    mean_value: floating[Any]
    start_index: int
    finish_index: int


class PlotWidget(BasePlotWidget):

    def _create_stage_region(self,
                             stage: str,
                             times: pd.Series,
                             dataframe: pd.DataFrame) -> Optional[pg.LinearRegionItem]:
        # TODO: сделать фабрику регионов: вертикальный, горизонтальный...
        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 get_stage_info(self,
                       stage: str,
                       dataframe: pd.DataFrame,
                       signal_name: str) -> Optional[ProcessStage]:
        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 ProcessStage(mean_value=mean, start_index=int(start), finish_index=int(finish))
            return None

    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:
                stages = [self.get_stage_info("Welding",
                                              dataframe,
                                              signal["name"]) for signal in description["Signals"]]
                if stages:
                    means_raw = [stage.mean_value for stage in stages]
                    mean = max(means_raw)
                    start = time_axis[stages[0].start_index]
                    finish = time_axis[stages[0].finish_index]

                    overshoot = pg.BarGraphItem(x0=0,
                                                y0=mean - mean * 0.05,
                                                height=mean * 0.05 * 2,
                                                width=start,
                                                brush=pg.mkBrush([0, 250, 0, 100]))
                    plot_widget.addItem(overshoot)

                    stable = pg.BarGraphItem(x0=start,
                                             y0=mean - mean * 0.015,
                                             height=mean * 0.015 * 2,
                                             width=finish - start,
                                             brush=pg.mkBrush([0, 250, 0, 100]))
                    plot_widget.addItem(stable)

                    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)