Skip to content
Snippets Groups Projects
mainui.py 17.74 KiB
import os
import platform
import subprocess
import logging
import pyqtgraph as pg

from PyQt5.QtWidgets import QStatusBar, QLineEdit
from pymeasure.display.browser import BrowserItem
from pymeasure.display.manager import Manager, Experiment
from pymeasure.display.Qt import QtCore, QtWidgets, QtGui
from pymeasure.display.widgets import (
    PlotWidget,
    BrowserWidget,
    InputsWidget,
    LogWidget,
    ResultsDialog,
    DirectoryLineEdit,
)
from pymeasure.experiment import Results, Procedure

from ..classes.logger import log, StatusBarHandler

class UIWindowBase(QtWidgets.QMainWindow):
    def __init__(self,
                 procedure_class,
                 widget_list=(),
                 inputs=(),
                 displays=(),
                 parent=None,
                 directory_input=False,
                 hide_groups=True,
                 ):

        super().__init__(parent)
        app = QtCore.QCoreApplication.instance()
        app.aboutToQuit.connect(self.quit)
        self.procedure_class = procedure_class
        self.inputs = inputs
        self.hide_groups = hide_groups
        self.displays = displays
        self.directory_input = directory_input
        self.log_level = logging.INFO
        self.widget_list = widget_list
        log.setLevel(logging.INFO)

        self._setup_ui()
        self._layout()

    def _setup_ui(self):
        if self.directory_input:
            self.directory_label = QtWidgets.QLabel(self)
            self.directory_label.setText('Directory')
            self.directory_line = DirectoryLineEdit(parent=self)

        self.repetitions_label = QtWidgets.QLabel(self)
        self.repetitions_label.setText("Number of experiment to perform:")
        self.repetitions_line = QLineEdit(self)
        self.repetitions_line.setText("1")

        self.queue_button = QtWidgets.QPushButton('Queue', self)
        self.queue_button.clicked.connect(self._queue)

        self.abort_button = QtWidgets.QPushButton('Abort', self)
        self.abort_button.setEnabled(False)
        self.abort_button.clicked.connect(self.abort)

        self.browser_widget = BrowserWidget(
            self.procedure_class[0],
            self.displays,
            [],
            parent=self
        )
        self.browser_widget.show_button.clicked.connect(self.show_experiments)
        self.browser_widget.hide_button.clicked.connect(self.hide_experiments)
        self.browser_widget.clear_button.clicked.connect(self.clear_experiments)
        self.browser_widget.open_button.clicked.connect(self.open_experiment)
        self.browser = self.browser_widget.browser

        self.browser.setContextMenuPolicy(QtCore.Qt.ContextMenuPolicy.CustomContextMenu)
        self.browser.customContextMenuRequested.connect(self.browser_item_menu)
        self.browser.itemChanged.connect(self.browser_item_changed)

        self.inputs = [InputsWidget(
            self.procedure_class[i],
            self.inputs[i],
            parent=self,
            hide_groups=self.hide_groups,
        ) for i in range(len(self.procedure_class))]

        self.manager = Manager(self.widget_list,
                               self.browser,
                               log_level=self.log_level,
                               parent=self)
        self.manager.abort_returned.connect(self.abort_returned)
        self.manager.queued.connect(self.queued)
        self.manager.running.connect(self.running)
        self.manager.finished.connect(self.finished)
        self.manager.log.connect(log.handle)


    def _layout(self):
        self.main = QtWidgets.QWidget(self)

        self.statusBar = QStatusBar()
        self.setStatusBar(self.statusBar)

        self.handler = StatusBarHandler(self.statusBar)
        self.handler.setLevel(logging.INFO)
        log.addHandler(self.handler)

        inputs_widget = QtWidgets.QWidget()
        inputs_vbox = QtWidgets.QVBoxLayout(inputs_widget)

        self.tabs_exp = QtWidgets.QTabWidget()

        for i in range(0, len(self.procedure_class)):
            tab_content = QtWidgets.QWidget()
            tab_content_layout = QtWidgets.QVBoxLayout(tab_content)
            self.inputs[i].setSizePolicy(QtWidgets.QSizePolicy.Policy.Minimum,
                            QtWidgets.QSizePolicy.Policy.Fixed)
            tab_content_layout.addWidget(self.inputs[i])
            self.tabs_exp.addTab(tab_content, self.procedure_class[i].name)

        inputs_vbox.addWidget(self.tabs_exp)
        inputs_vbox.addStretch(0)

        hbox = QtWidgets.QHBoxLayout()
        hbox.addWidget(self.queue_button)
        hbox.addWidget(self.abort_button)

        vbox = QtWidgets.QVBoxLayout()
        vbox.addWidget(self.repetitions_label)
        vbox.addWidget(self.repetitions_line)
        vbox.addWidget(self.directory_label)
        vbox.addWidget(self.directory_line)
        vbox.addLayout(hbox)
        inputs_vbox.addLayout(vbox)


        self.tabs = QtWidgets.QTabWidget(self.main)
        for wdg in self.widget_list:
            self.tabs.addTab(wdg, wdg.name)

        splitter = QtWidgets.QSplitter(QtCore.Qt.Orientation.Vertical)
        splitter.addWidget(self.tabs)
        splitter.addWidget(self.browser_widget)

        splitter_main = QtWidgets.QSplitter(QtCore.Qt.Orientation.Horizontal)
        splitter_main.addWidget(inputs_widget)
        splitter_main.addWidget(splitter)

        main_layout = QtWidgets.QHBoxLayout(self.main)
        main_layout.addWidget(splitter_main)

        self.main.setLayout(main_layout)
        self.setCentralWidget(self.main)
        self.main.show()
        self.resize(1000, 800)

    def get_selected_tab_index(self):
        if self.tabs_exp is not None:
            return self.tabs_exp.currentIndex()
        return 0

    def quit(self, evt=None):
        if self.manager.is_running():
            self.abort()

        self.close()

    def browser_item_changed(self, item, column):
        if column == 0:
            state = item.checkState(0)
            experiment = self.manager.experiments.with_browser_item(item)
            if state == QtCore.Qt.CheckState.Unchecked:
                for wdg, curve in zip(self.widget_list, experiment.curve_list):
                    wdg.remove(curve)
            else:
                for wdg, curve in zip(self.widget_list, experiment.curve_list):
                    wdg.load(curve)

    def browser_item_menu(self, position):
        item = self.browser.itemAt(position)

        if item is not None:
            experiment = self.manager.experiments.with_browser_item(item)

            menu = QtWidgets.QMenu(self)

            # Open
            action_open = QtGui.QAction(menu)
            action_open.setText("Open Data Externally")
            action_open.triggered.connect(
                lambda: self.open_file_externally(experiment.results.data_filename))
            menu.addAction(action_open)

            # Change Color
            action_change_color = QtGui.QAction(menu)
            action_change_color.setText("Change Color")
            action_change_color.triggered.connect(
                lambda: self.change_color(experiment))
            menu.addAction(action_change_color)

            # Remove
            action_remove = QtGui.QAction(menu)
            action_remove.setText("Remove Graph")
            if self.manager.is_running():
                if self.manager.running_experiment() == experiment:  # Experiment running
                    action_remove.setEnabled(False)
            action_remove.triggered.connect(lambda: self.remove_experiment(experiment))
            menu.addAction(action_remove)

            # Delete
            action_delete = QtGui.QAction(menu)
            action_delete.setText("Delete Data File")
            if self.manager.is_running():
                if self.manager.running_experiment() == experiment:  # Experiment running
                    action_delete.setEnabled(False)
            action_delete.triggered.connect(lambda: self.delete_experiment_data(experiment))
            menu.addAction(action_delete)

            # Use parameters
            action_use = QtGui.QAction(menu)
            action_use.setText("Use These Parameters")
            action_use.triggered.connect(
                lambda: self.set_parameters(experiment.procedure.parameter_objects()))
            menu.addAction(action_use)
            menu.exec(self.browser.viewport().mapToGlobal(position))

    def remove_experiment(self, experiment):
        reply = QtWidgets.QMessageBox.question(self, 'Remove Graph',
                                               "Are you sure you want to remove the graph?",
                                               QtWidgets.QMessageBox.StandardButton.Yes |
                                               QtWidgets.QMessageBox.StandardButton.No,
                                               QtWidgets.QMessageBox.StandardButton.No)
        if reply == QtWidgets.QMessageBox.StandardButton.Yes:
            self.manager.remove(experiment)

    def delete_experiment_data(self, experiment):
        reply = QtWidgets.QMessageBox.question(self, 'Delete Data',
                                               "Are you sure you want to delete this data file?",
                                               QtWidgets.QMessageBox.StandardButton.Yes |
                                               QtWidgets.QMessageBox.StandardButton.No,
                                               QtWidgets.QMessageBox.StandardButton.No)
        if reply == QtWidgets.QMessageBox.StandardButton.Yes:
            self.manager.remove(experiment)
            os.unlink(experiment.data_filename)

    def show_experiments(self):
        root = self.browser.invisibleRootItem()
        for i in range(root.childCount()):
            item = root.child(i)
            item.setCheckState(0, QtCore.Qt.CheckState.Checked)

    def hide_experiments(self):
        root = self.browser.invisibleRootItem()
        for i in range(root.childCount()):
            item = root.child(i)
            item.setCheckState(0, QtCore.Qt.CheckState.Unchecked)

    def clear_experiments(self):
        self.manager.clear()

    def open_experiment(self):
        dialog = ResultsDialog(self.procedure_class[0].DATA_COLUMNS, self.x_axis, self.y_axis)
        if dialog.exec():
            filenames = dialog.selectedFiles()
            for filename in map(str, filenames):
                if filename in self.manager.experiments:
                    QtWidgets.QMessageBox.warning(
                        self, "Load Error",
                        "The file %s cannot be opened twice." % os.path.basename(filename)
                    )
                elif filename == '':
                    return
                else:
                    results = Results.load(filename)
                    experiment = self.new_experiment(results)
                    for curve in experiment.curve_list:
                        if curve:
                            curve.update_data()
                    experiment.browser_item.progressbar.setValue(100)
                    self.manager.load(experiment)
                    log.info('Opened data file %s' % filename)

    def change_color(self, experiment):
        color = QtWidgets.QColorDialog.getColor(
            parent=self)
        if color.isValid():
            pixelmap = QtGui.QPixmap(24, 24)
            pixelmap.fill(color)
            experiment.browser_item.setIcon(0, QtGui.QIcon(pixelmap))
            for wdg, curve in zip(self.widget_list, experiment.curve_list):
                wdg.set_color(curve, color=color)

    def open_file_externally(self, filename):
        """ Method to open the datafile using an external editor or viewer. Uses the default
        application to open a datafile of this filetype, but can be overridden by the child
        class in order to open the file in another application of choice.
        """
        system = platform.system()
        if (system == 'Windows'):
            # The empty argument after the start is needed to be able to cope
            # correctly with filenames with spaces
            _ = subprocess.Popen(['start', '', filename], shell=True)
        elif (system == 'Linux'):
            _ = subprocess.Popen(['xdg-open', filename])
        elif (system == 'Darwin'):
            _ = subprocess.Popen(['open', filename])
        else:
            raise Exception("{cls} method open_file_externally does not support {system} OS".format(
                cls=type(self).__name__, system=system))

    def make_procedure(self):
        if not isinstance(self.inputs[self.get_selected_tab_index()], InputsWidget):
            raise Exception("ManagedWindow can not make a Procedure"
                            " without a InputsWidget type")
        return self.inputs[self.get_selected_tab_index()].get_procedure()

    def new_curve(self, wdg, results, color=None, **kwargs):
        if color is None:
            color = pg.intColor(self.browser.topLevelItemCount() % 8)
        return wdg.new_curve(results, color=color, **kwargs)

    def new_experiment(self, results, curve=None):
        if curve is None:
            curve_list = []
            for wdg in self.widget_list:
                curve_list.append(self.new_curve(wdg, results))
        else:
            curve_list = curve[:]

        curve_color = pg.intColor(0)
        for wdg, curve in zip(self.widget_list, curve_list):
            if isinstance(wdg, PlotWidget):
                curve_color = curve.opts['pen'].color()
                break

        browser_item = BrowserItem(results, curve_color)
        return Experiment(results, curve_list, browser_item)

    def set_parameters(self, parameters):
        """ This method should be overwritten by the child class. The
        parameters argument is a dictionary of Parameter objects.
        The Parameters should overwrite the GUI values so that a user
        can click "Queue" to capture the same parameters.
        """
        if not isinstance(self.inputs[0], InputsWidget):
            raise Exception("ManagedWindow can not set parameters"
                            " without a InputsWidget")
        self.inputs[0].set_parameters(parameters)

    def _queue(self, checked):
        """ This method is a wrapper for the `self.queue` method to be connected
        to the `queue` button. It catches the positional argument that is passed
        when it is called by the button and calls the `self.queue` method without
        any arguments.
        """
        self.widget_list[0].close_connections() # Not clean because works only for this really specific case
        for i in range(self.number_repetitions):
            self.queue()

    def queue(self, procedure=None):
        raise NotImplementedError(
            "Abstract method ManagedWindow.queue not implemented")

    def abort(self):
        self.abort_button.setEnabled(False)
        self.abort_button.setText("Resume")
        self.abort_button.clicked.disconnect()
        self.abort_button.clicked.connect(self.resume)
        try:
            self.manager.abort()
        except:  # noqa
            log.error('Failed to abort experiment', exc_info=True)
            self.abort_button.setText("Abort")
            self.abort_button.clicked.disconnect()
            self.abort_button.clicked.connect(self.abort)

    def resume(self):
        self.abort_button.setText("Abort")
        self.abort_button.clicked.disconnect()
        self.abort_button.clicked.connect(self.abort)
        if self.manager.experiments.has_next():
            self.manager.resume()
        else:
            self.abort_button.setEnabled(False)

    def queued(self, experiment):
        self.abort_button.setEnabled(True)
        self.browser_widget.show_button.setEnabled(True)
        self.browser_widget.hide_button.setEnabled(True)
        self.browser_widget.clear_button.setEnabled(True)

    def running(self, experiment):
        self.browser_widget.clear_button.setEnabled(False)

    def abort_returned(self, experiment):
        if self.manager.experiments.has_next():
            self.abort_button.setText("Resume")
            self.abort_button.setEnabled(True)
        else:
            self.browser_widget.clear_button.setEnabled(True)

    def finished(self, experiment):
        if not self.manager.experiments.has_next():
            self.abort_button.setEnabled(False)
            self.browser_widget.clear_button.setEnabled(True)

    @property
    def directory(self):
        if not self.directory_input:
            raise ValueError("No directory input in the ManagedWindow")
        return self.directory_line.text()
    
    @property
    def number_repetitions(self) -> int:
        try:
            return int(self.repetitions_line.text())
        except:
            log.info("The number of repetitions should be an integer, currently using 1")
            return 1

    @directory.setter
    def directory(self, value):
        if not self.directory_input:
            raise ValueError("No directory input in the ManagedWindow")

        self.directory_line.setText(str(value))


class UIWindown(UIWindowBase):
    def __init__(self, procedure_class, x_axis=None, y_axis=None, linewidth=1, **kwargs):
        self.x_axis = x_axis
        self.y_axis = y_axis
        self.log_widget = LogWidget("Experiment Log")
        self.plot_widget = PlotWidget("Results Graph", procedure_class[0].DATA_COLUMNS, self.x_axis,
                                      self.y_axis, linewidth=linewidth)
        self.plot_widget.setMinimumSize(100, 200)

        if "widget_list" not in kwargs:
            kwargs["widget_list"] = ()
        kwargs["widget_list"] = kwargs["widget_list"] + (self.plot_widget, self.log_widget)

        super().__init__(procedure_class, **kwargs)

        # Setup measured_quantities once we know x_axis and y_axis
        self.browser_widget.browser.measured_quantities = [self.x_axis, self.y_axis]

        logging.getLogger().addHandler(self.log_widget.handler)  # needs to be in Qt context?
        log.setLevel(self.log_level)
        log.info("miniMOKE ready for action")