dylab/HelperClasses/Plotting/MainPlotPanel.py

import pyqtgraph as pg
from pyqtgraph.Qt import QtCore, QtGui
import numpy as np
import os

from pyqtgraph.dockarea import *
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from .WidgetDataSelector import WidgetDataSelector
from .WidgetAnalysePanel import WidgetAnalysePanel


def write_pid():
    with open("living_plot.ini", 'w') as pid_file:
        pid = os.getpid()
        pid_file.write(str(pid))

    return pid


def clean_pid():
    with open("living_plot.ini", 'w') as pid_file:
        pid_file.truncate(0)


class MainPlotPanel(QtGui.QMainWindow):

    def __init__(self, h5_paths, n_rows=3, **kwargs):

        super().__init__(**kwargs)

        write_pid()

        self.h5_paths = np.array(h5_paths)

        pg.mkQApp()

        self.n_rows = n_rows

        self.setWindowFlag(QtCore.Qt.WindowCloseButtonHint, False)
        self.area = DockArea()

        self.setCentralWidget(self.area)
        self.resize(900, 500)

        self.dshotselector = Dock("Shot selector")
        self.shotselector = WidgetDataSelector()
        self.dshotselector.setFixedSize(400, 150)

        self.dshotselector.addWidget(self.shotselector)
        self.area.addDock(self.dshotselector, 'left')

        # self.qpg_dock = Dock("Quick Plot Generator")
        # self.qpg_dock.addWidget(QuickPlotGenerator(self))
        # self.qpg_dock.setMinimumSize(self.qpg_dock.minimumSizeHint())
        # self.area.addDock(self.qpg_dock)

        self.danalyse_panel = Dock("Analyse and Fitting")
        self.analyse_panel = WidgetAnalysePanel()
        self.danalyse_panel.addWidget(self.analyse_panel)
        self.area.addDock(self.danalyse_panel, 'bottom')

        self.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)
        self.showMaximized()

        self.plots = {}

        self.shotselector.valueChanged.connect(self.refresh)
        self.shotselector.selectionChanged.connect(self.refresh)

    def closeEvent(self, event):
        clean_pid()
        event.accept()

    def add_plot_dock(self, plot_name, plot_widget, overlap=True, **kwargs):
        if plot_name not in self.plots:
            plot_widget.plot_name = plot_name

            dock = Dock(plot_name, **kwargs)

            dock.sigClosed.connect(self.remove_plot)
            dock.addWidget(plot_widget)
            if overlap and len(self.plots):
                self.area.addDock(dock, 'above', relativeTo=self.area.docks[list(self.plots.keys())[-1]])
            elif len(self.plots):
                if len(self.plots) % self.n_rows == 2:
                    position = 'right'
                    self.area.addDock(dock, position)
                if len(self.plots) % (self.n_rows * 2) in [0, 1]:
                    position = 'bottom'
                    self.area.addDock(dock, position, relativeTo=self.area.docks[list(self.plots.keys())[-1]])
                if len(self.plots) % (self.n_rows * 2) in [3, 4]:
                    position = 'top'
                    self.area.addDock(dock, position, relativeTo=self.area.docks[list(self.plots.keys())[-1]])
                print(len(self.plots) % (self.n_rows * 2), len(self.plots) % self.n_rows)
            else:
                self.area.addDock(dock, 'right')

            self.plots[plot_name] = plot_widget
        else:
            print(f'Plot {plot_name} already exists. Please choose different name.')

    def remove_plot(self, dock):
        del self.plots[dock.title()]

    def update_h5_paths(self, h5_paths):
        self.h5_paths = np.array(h5_paths)

        self.shotselector.update_nshots(len(h5_paths))

        for plot_name in self.plots:
            self.plots[plot_name].data_extractor_manager.clean_memory(h5_paths)

    def add_h5_path(self, h5_path):
        self.h5_paths = np.append(self.h5_paths, h5_path)

        self.shotselector.update_nshots(len(self.h5_paths))

    def refresh(self, h5_path=None):
        if len(self.h5_paths):

            self.h5_paths_selected = self.h5_paths[np.array(self.shotselector.get_selected_indices())]

            if h5_path == None:
                i = self.shotselector.get_current_index()
                h5_path = self.h5_paths[np.array(i)]

            for plot_name in self.plots:
                self.plots[plot_name].data_extractor_manager.update_local_data(h5_path)

            for plot_name, plot in self.plots.items():
                plot.update_from_h5(h5_path)
        else:
            pass