artiq/artiq/gui/explorer.py

import asyncio
import logging

from quamash import QtGui, QtCore
from pyqtgraph import dockarea
from pyqtgraph import LayoutWidget

from artiq.protocols import pyon
from artiq.gui.models import DictSyncModel
from artiq.gui.scan import ScanController
from artiq.gui.shortcuts import ShortcutManager


class Model(DictSyncModel):
    def __init__(self, init):
        DictSyncModel.__init__(self,
            ["Experiment"],
            init)
        self.explorer = None

    def sort_key(self, k, v):
        return k

    def convert(self, k, v, column):
        return k

    def __setitem__(self, k, v):
        DictSyncModel.__setitem__(self, k, v)
        if self.explorer is not None:
            if k == self.explorer.selected_key:
                self.explorer.update_selection(k, k)


class _FreeValueEntry(QtGui.QLineEdit):
    def __init__(self, procdesc):
        QtGui.QLineEdit.__init__(self)
        if "default" in procdesc:
            self.set_argument_value(procdesc["default"])

    def get_argument_value(self):
        return pyon.decode(self.text())

    def set_argument_value(self, value):
        self.setText(pyon.encode(value))


class _BooleanEntry(QtGui.QCheckBox):
    def __init__(self, procdesc):
        QtGui.QCheckBox.__init__(self)
        if "default" in procdesc:
            self.set_argument_value(procdesc["default"])

    def get_argument_value(self):
        return self.isChecked()

    def set_argument_value(self, value):
        self.setChecked(value)


class _EnumerationEntry(QtGui.QComboBox):
    def __init__(self, procdesc):
        QtGui.QComboBox.__init__(self)
        self.choices = procdesc["choices"]
        self.addItems(self.choices)
        if "default" in procdesc:
            self.set_argument_value(procdesc["default"])

    def get_argument_value(self):
        return self.choices[self.currentIndex()]

    def set_argument_value(self, value):
        idx = self.choices.index(value)
        self.setCurrentIndex(idx)


class _NumberEntry(QtGui.QDoubleSpinBox):
    def __init__(self, procdesc):
        QtGui.QDoubleSpinBox.__init__(self)
        self.scale = procdesc["scale"]
        self.setDecimals(procdesc["ndecimals"])
        self.setSingleStep(procdesc["step"]/self.scale)
        if procdesc["min"] is not None:
            self.setMinimum(procdesc["min"]/self.scale)
        else:
            self.setMinimum(float("-inf"))
        if procdesc["max"] is not None:
            self.setMaximum(procdesc["max"]/self.scale)
        else:
            self.setMaximum(float("inf"))
        if procdesc["unit"]:
            self.setSuffix(" " + procdesc["unit"])
        if "default" in procdesc:
            self.set_argument_value(procdesc["default"])

    def get_argument_value(self):
        return self.value()*self.scale

    def set_argument_value(self, value):
        self.setValue(value/self.scale)


class _StringEntry(QtGui.QLineEdit):
    def __init__(self, procdesc):
        QtGui.QLineEdit.__init__(self)
        if "default" in procdesc:
            self.set_argument_value(procdesc["default"])

    def get_argument_value(self):
        return self.text()

    def set_argument_value(self, value):
        self.setText(value)


_procty_to_entry = {
    "FreeValue": _FreeValueEntry,
    "BooleanValue": _BooleanEntry,
    "EnumerationValue": _EnumerationEntry,
    "NumberValue": _NumberEntry,
    "StringValue": _StringEntry,
    "Scannable": ScanController
}


class _ArgumentEditor(QtGui.QTreeWidget):
    def __init__(self, main_window):
        QtGui.QTreeWidget.__init__(self)
        self.setColumnCount(2)
        self.header().setResizeMode(QtGui.QHeaderView.ResizeToContents)
        self.header().setVisible(False)
        self.setSelectionMode(QtGui.QAbstractItemView.NoSelection)

        self.main_window = main_window
        self._groups = dict()
        self.set_arguments([])

    def clear(self):
        QtGui.QTreeWidget.clear(self)
        self._groups.clear()

    def _get_group(self, name):
        if name in self._groups:
            return self._groups[name]
        group = QtGui.QTreeWidgetItem([name, ""])
        for c in 0, 1:
            group.setBackground(c, QtGui.QBrush(QtGui.QColor(100, 100, 100)))
            group.setForeground(c, QtGui.QBrush(QtGui.QColor(220, 220, 255)))
            font = group.font(c)
            font.setBold(True)
            group.setFont(c, font)
        self.addTopLevelItem(group)
        self._groups[name] = group
        return group

    def set_arguments(self, arguments):
        self.clear()

        if not arguments:
            self.addTopLevelItem(QtGui.QTreeWidgetItem(["No arguments", ""]))

        self._args_to_entries = dict()
        for n, (name, (procdesc, group)) in enumerate(arguments):
            entry = _procty_to_entry[procdesc["ty"]](procdesc)
            self._args_to_entries[name] = entry

            widget_item = QtGui.QTreeWidgetItem([name, ""])
            if group is None:
                self.addTopLevelItem(widget_item)
            else:
                self._get_group(group).addChild(widget_item)
            self.setItemWidget(widget_item, 1, entry)

    def get_argument_values(self, show_error_message):
        r = dict()
        for arg, entry in self._args_to_entries.items():
            try:
                r[arg] = entry.get_argument_value()
            except Exception as e:
                if show_error_message:
                    msgbox = QtGui.QMessageBox(self.main_window)
                    msgbox.setWindowTitle("Error")
                    msgbox.setText("Failed to obtain value for argument '{}':\n{}"
                                   .format(arg, str(e)))
                    msgbox.setStandardButtons(QtGui.QMessageBox.Ok)
                    msgbox.show()
                return None
        return r

    def set_argument_values(self, arguments, ignore_errors):
        for arg, value in arguments.items():
            try:
                entry = self._args_to_entries[arg]
                entry.set_argument_value(value)
            except:
                if not ignore_errors:
                    raise

    def save_state(self):
        expanded = []
        for k, v in self._groups.items():
            if v.isExpanded():
                expanded.append(k)
        argument_values = self.get_argument_values(False)
        return {
            "expanded": expanded,
            "argument_values": argument_values
        }

    def restore_state(self, state):
        self.set_argument_values(state["argument_values"], True)
        for e in state["expanded"]:
            try:
                self._groups[e].setExpanded(True)
            except KeyError:
                pass


class ExplorerDock(dockarea.Dock):
    def __init__(self, main_window, status_bar,
                 explist_sub, schedule_sub,
                 schedule_ctl, repository_ctl):
        dockarea.Dock.__init__(self, "Explorer", size=(1500, 500))

        self.main_window = main_window
        self.status_bar = status_bar
        self.schedule_sub = schedule_sub
        self.schedule_ctl = schedule_ctl

        self.splitter = QtGui.QSplitter(QtCore.Qt.Horizontal)
        self.addWidget(self.splitter)

        grid = LayoutWidget()
        self.splitter.addWidget(grid)

        self.el = QtGui.QListView()
        self.el.selectionChanged = self._selection_changed
        self.selected_key = None
        grid.addWidget(self.el, 0, 0, colspan=4)

        self.datetime = QtGui.QDateTimeEdit()
        self.datetime.setDisplayFormat("MMM d yyyy hh:mm:ss")
        self.datetime.setDate(QtCore.QDate.currentDate())
        self.datetime.dateTimeChanged.connect(self.enable_duedate)
        self.datetime_en = QtGui.QCheckBox("Due date:")
        grid.addWidget(self.datetime_en, 1, 0, colspan=2)
        grid.addWidget(self.datetime, 1, 2, colspan=2)

        self.pipeline = QtGui.QLineEdit()
        self.pipeline.setText("main")
        grid.addWidget(QtGui.QLabel("Pipeline:"), 2, 0, colspan=2)
        grid.addWidget(self.pipeline, 2, 2, colspan=2)

        self.priority = QtGui.QSpinBox()
        self.priority.setRange(-99, 99)
        grid.addWidget(QtGui.QLabel("Priority:"), 3, 0)
        grid.addWidget(self.priority, 3, 1)

        self.flush = QtGui.QCheckBox("Flush")
        self.flush.setToolTip("Flush the pipeline before starting the experiment")
        grid.addWidget(self.flush, 3, 2)

        self.log_level = QtGui.QComboBox()
        self.log_level.addItems(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"])
        self.log_level.setCurrentIndex(1)
        self.log_level.setToolTip("Minimum level for log entry production")
        grid.addWidget(self.log_level, 3, 3)

        submit = QtGui.QPushButton("Submit")
        submit.setToolTip("Schedule the selected experiment (Ctrl+Return)")
        grid.addWidget(submit, 4, 0, colspan=4)
        submit.clicked.connect(self.submit_clicked)

        self.argeditor = _ArgumentEditor(self.main_window)
        self.splitter.addWidget(self.argeditor)
        self.splitter.setSizes([grid.minimumSizeHint().width(), 1000])
        self.argeditor_states = dict()
        self.explist_model = Model(dict())
        explist_sub.add_setmodel_callback(self.set_model)

        self.shortcuts = ShortcutManager(self.main_window, self)

        self.el.setContextMenuPolicy(QtCore.Qt.ActionsContextMenu)
        submit_action = QtGui.QAction("Submit", self.el)
        submit_action.triggered.connect(self.submit_clicked)
        submit_action.setShortcut("CTRL+RETURN")
        self.el.addAction(submit_action)
        reqterm_action = QtGui.QAction("Request termination of instances", self.el)
        reqterm_action.triggered.connect(self.request_inst_term)
        reqterm_action.setShortcut("CTRL+BACKSPACE")
        self.el.addAction(reqterm_action)

        sep = QtGui.QAction(self.el)
        sep.setSeparator(True)
        self.el.addAction(sep)

        edit_shortcuts_action = QtGui.QAction("Edit shortcuts", self.el)
        edit_shortcuts_action.triggered.connect(self.edit_shortcuts)
        self.el.addAction(edit_shortcuts_action)
        scan_repository_action = QtGui.QAction("(Re)scan repository HEAD",
                                               self.el)
        def scan_repository():
            asyncio.ensure_future(repository_ctl.scan_async())
            self.status_bar.showMessage("Requested repository scan")
        scan_repository_action.triggered.connect(scan_repository)
        self.el.addAction(scan_repository_action)

    def set_model(self, model):
        model.explorer = self
        self.explist_model = model
        self.el.setModel(model)

    def update_selection(self, selected, deselected):
        if deselected:
            self.argeditor_states[deselected] = self.argeditor.save_state()

        if selected:
            expinfo = self.explist_model.backing_store[selected]
            self.argeditor.set_arguments(expinfo["arguments"])
            if selected in self.argeditor_states:
                self.argeditor.restore_state(self.argeditor_states[selected])
            self.splitter.insertWidget(1, self.argeditor)
        self.selected_key = selected

    def _sel_to_key(self, selection):
        selection = selection.indexes()
        if selection:
            row = selection[0].row()
            return self.explist_model.row_to_key[row]
        else:
            return None

    def _selection_changed(self, selected, deselected):
        self.update_selection(self._sel_to_key(selected),
                              self._sel_to_key(deselected))

    def save_state(self):
        idx = self.el.selectedIndexes()
        if idx:
            row = idx[0].row()
            key = self.explist_model.row_to_key[row]
            self.argeditor_states[key] = self.argeditor.save_state()
        return {
            "argeditor": self.argeditor_states,
            "shortcuts": self.shortcuts.save_state()
        }

    def restore_state(self, state):
        try:
            argeditor_states = state["argeditor"]
            shortcuts_state = state["shortcuts"]
        except KeyError:
            return
        self.argeditor_states = argeditor_states
        self.shortcuts.restore_state(shortcuts_state)

    def enable_duedate(self):
        self.datetime_en.setChecked(True)

    async def submit_task(self, pipeline_name, file, class_name, arguments,
                          priority, due_date, flush):
        expid = {
            "log_level": getattr(logging, self.log_level.currentText()),
            "repo_rev": None,
            "file": file,
            "class_name": class_name,
            "arguments": arguments,
        }
        rid = await self.schedule_ctl.submit(pipeline_name, expid,
                                             priority, due_date, flush)
        self.status_bar.showMessage("Submitted RID {}".format(rid))

    def submit(self, pipeline, key, priority, due_date, flush):
        # TODO: refactor explorer and cleanup.
        # Argument editors should immediately modify the global state.
        expinfo = self.explist_model.backing_store[key]
        if key == self.selected_key:
            arguments = self.argeditor.get_argument_values(True)
            if arguments is None:
                # There has been an error. Displaying the error message box
                # was done by argeditor.
                return
        else:
            try:
                arguments = self.argeditor_states[key]["argument_values"]
            except KeyError:
                arguments = dict()
        asyncio.ensure_future(self.submit_task(pipeline, expinfo["file"],
                                               expinfo["class_name"],
                                               arguments, priority, due_date,
                                               flush))

    def submit_clicked(self):
        if self.selected_key is not None:
            if self.datetime_en.isChecked():
                due_date = self.datetime.dateTime().toMSecsSinceEpoch()/1000
            else:
                due_date = None
            self.submit(self.pipeline.text(),
                        self.selected_key,
                        self.priority.value(),
                        due_date,
                        self.flush.isChecked())

    async def request_term_multiple(self, rids):
        for rid in rids:
            try:
                await self.schedule_ctl.request_termination(rid)
            except:
                pass

    def request_inst_term(self):
        if self.selected_key is not None:
            expinfo = self.explist_model.backing_store[self.selected_key]
            if self.schedule_sub.model is not None:
                current_schedule = self.schedule_sub.model.backing_store
                rids = []
                for rid, desc in current_schedule.items():
                    expid = desc["expid"]
                    if ("repo_rev" in expid  # only consider runs from repository
                            and expid["file"] == expinfo["file"]
                            and expid["class_name"] == expinfo["class_name"]):
                        rids.append(rid)
                asyncio.ensure_future(self.request_term_multiple(rids))

    def edit_shortcuts(self):
        experiments = sorted(self.explist_model.backing_store.keys())
        self.shortcuts.edit(experiments)