Source code for schrodinger.trajectory.trajectory_gui_dir.banners

from enum import Enum

from schrodinger import get_maestro
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtWidgets
from schrodinger.ui import picking
from schrodinger.ui.qt import basewidgets

from . import centroid_banner_ui
from . import planar_banner_ui
from . import stylesheet

maestro = get_maestro()


class MeasurementMode(Enum):
    """
    Type of measurement for the CentroidBanner to search for
    """
    DISTANCE = 1
    ANGLE = 2


class CentroidBanner(basewidgets.BaseWidget):
    """
    Banner for picking the centroid of angles or distances
    """

    SETS_REQUIRED = {MeasurementMode.DISTANCE: 2, MeasurementMode.ANGLE: 3}

    ui_module = centroid_banner_ui
    closeBanner = QtCore.pyqtSignal()

    # Signal to emit when all centroid atoms are picked
    centroidsPicked = QtCore.pyqtSignal(list)

    def __init__(self):
        super().__init__()
        self.picked_sets = []
        self.current_aids = []
        self._measurement_mode = MeasurementMode.DISTANCE

    def initSetUp(self):
        super().initSetUp()
        self.setStyleSheet(stylesheet.TRAJECTORY_PLOT_BANNERS)
        self.setObjectName('centroid_banner')
        self._setupMeasurementMenu()
        self.ui.cancel_btn.clicked.connect(self._handleCancel)
        self.ui.save_set_tb.clicked.connect(self._saveSet)
        self.ui.ok_btn.clicked.connect(self._handleOk)
        self.cb = QtWidgets.QCheckBox()
        self.atom_picker = picking.PickAtomsToggle(self.cb,
                                                   None,
                                                   self._pickCallback,
                                                   pick_text='')

    def startPicking(self):
        self.atom_picker.start()

    def stopPicking(self):
        self.atom_picker.stop()

    def _pickCallback(self, aids):
        required_sets = self.SETS_REQUIRED[self._measurement_mode]
        if len(self.picked_sets) < required_sets:
            self.current_aids = aids
            self._updateUI()

    def _saveSet(self):
        if len(self.picked_sets) < self.SETS_REQUIRED[self._measurement_mode]:
            self.picked_sets.append(self.current_aids)
            self.current_aids = set()
            self._updateUI()

            # Toggle Picking state to clear current selections
            self.stopPicking()
            self.startPicking()

    def _updateUI(self):
        """
        Updates UI to account for changes
        """
        self._updateCancelString()
        self._updatePickProgressString()
        self._updateDescriptionString()
        self._updateOKButton()

    def _updatePickProgressString(self):
        sets_needed = self.SETS_REQUIRED[self._measurement_mode]
        pick_string = f'{len(self.picked_sets)} / {sets_needed} sets defined'
        self.ui.number_selected_lbl.setText(pick_string)
        color = 'yellow' if len(self.picked_sets) >= 1 else 'white'
        formatted_color = 'QLabel {color: ' + color + ';}'
        self.ui.number_selected_lbl.setStyleSheet(formatted_color)

    def _updateOKButton(self):
        enable = len(
            self.picked_sets) == self.SETS_REQUIRED[self._measurement_mode]
        self.ui.ok_btn.setEnabled(enable)

    def _updateCancelString(self):
        text = 'Start Over' if len(self.picked_sets) >= 1 else 'Cancel'
        self.ui.cancel_btn.setText(text)

    def _updateDescriptionString(self):
        sets_needed = self.SETS_REQUIRED[self._measurement_mode]
        description = f"""<i>Define {sets_needed} sets of atoms, clicking Save Set after each set has been selected.</i>"""
        self.ui.description_lbl.setText(description)

    def _setupMeasurementMenu(self):
        self.measurement_menu = QtWidgets.QMenu(self)

        action = self.measurement_menu.addAction('Distance')
        action.triggered.connect(
            lambda: self._onMenuChange(MeasurementMode.DISTANCE))

        action = self.measurement_menu.addAction('Angle')
        action.triggered.connect(
            lambda: self._onMenuChange(MeasurementMode.ANGLE))

        self.ui.measurement_tb.setPopupMode(QtWidgets.QToolButton.InstantPopup)
        self.ui.measurement_tb.setMenu(self.measurement_menu)

    def _onMenuChange(self, mode):
        """
        Handle a change in the measurement menu
        """
        if mode != self._measurement_mode:
            self._measurement_mode = mode
            formatted_text = mode.name.title()
            self.ui.measurement_tb.setText(formatted_text)
            self.picked_sets.clear()
            self.current_aids = set()
            self._updateUI()

    def _handleCancel(self):
        # If we have more than one pick, reset all picks
        if len(self.picked_sets) >= 1:
            self.picked_sets.clear()
            self._updateUI()
        # Otherwise, quit the banner
        else:
            self._handleClose()

    def _handleClose(self):
        self.stopPicking()
        self.closeBanner.emit()

    def _handleOk(self):
        if len(self.picked_sets) == self.SETS_REQUIRED[self._measurement_mode]:
            self.centroidsPicked.emit(self.picked_sets)
            self._handleClose()


class PlanarBanner(basewidgets.BaseWidget):
    """
    Banner for picking the planar angle of six atoms
    """

    ui_module = planar_banner_ui
    closeBanner = QtCore.pyqtSignal()

    # Signal to emit when all planar atoms are picked
    planarAnglesPicked = QtCore.pyqtSignal(list)

    def __init__(self):
        super().__init__()
        self.setStyleSheet(stylesheet.TRAJECTORY_PLOT_BANNERS)

    def initSetUp(self):
        super().initSetUp()
        self.picked_atoms = []
        self.setObjectName('planar_banner')
        self.ui.cancel_btn.clicked.connect(self._handleCancel)
        self.ui.ok_btn.clicked.connect(self._handleClose)
        self._cb = QtWidgets.QCheckBox()
        self.atom_picker = picking.PickAtomsToggle(self._cb,
                                                   None,
                                                   self._pickCallback,
                                                   pick_text='')
        self._updateUI()

    def startPicking(self):
        self.atom_picker.start()

    def stopPicking(self):
        self.atom_picker.stop()
        self.picked_atoms.clear()
        self._updateUI()

    def _pickCallback(self, aids):
        self.picked_atoms = aids
        if len(self.picked_atoms) == 6:
            self.planarAnglesPicked.emit(self.picked_atoms)
            self._resetPicks()
        self._updateUI()

    def _updateUI(self):
        """
        Updates UI to account for changes
        """
        self._updateCancelString()
        self._updatePickProgressString()

    def _updatePickProgressString(self):
        pick_string = f'{len(self.picked_atoms)} / 6 atoms selected'
        self.ui.number_selected_lbl.setText(pick_string)
        color = 'yellow' if len(self.picked_atoms) >= 1 else 'white'
        formatted_color = 'QLabel {color: ' + color + ';}'
        self.ui.number_selected_lbl.setStyleSheet(formatted_color)

    def _updateCancelString(self):
        text = 'Start Over' if len(self.picked_atoms) >= 1 else 'Cancel'
        self.ui.cancel_btn.setText(text)

    def _handleCancel(self):
        # If we have more than one pick, reset all picks
        if len(self.picked_atoms) >= 1:
            self._resetPicks()
        # Otherwise, quit the banner
        else:
            self._handleClose()

    def _handleClose(self):
        self.stopPicking()
        self.closeBanner.emit()

    def _resetPicks(self):
        self.stopPicking()
        self.startPicking()