Source code for schrodinger.application.jaguar.gui.input_tab_widgets_pka

import enum
import itertools
import re
from collections import namedtuple

import schrodinger
from schrodinger.application.jaguar import input as jaginput
from schrodinger.application.jaguar.gui.ui import define_smarts_panel_ui
from schrodinger.application.jaguar.gui.ui import pka_smarts_page_ui
from schrodinger.infra import mm
from schrodinger.models import mappers
from schrodinger.models import parameters
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import Qt
from schrodinger.structutils import analyze
from schrodinger.ui.qt import basewidgets
from schrodinger.ui.qt import delegates as qt_delegates
from schrodinger.ui.qt import pop_up_widgets
from schrodinger.ui.qt import table_helper

from . import input_tab_widgets
from .input_tab_widgets import SORT_ROLE
from .input_tab_widgets import WS_INCLUDE_ONLY

maestro = schrodinger.get_maestro()
JAG_NAME_STRUCTURE_ROLE = Qt.UserRole + 100
PKA_ATOM_ROLE = Qt.UserRole + 101

# We add the 100 to make sure that we don't conflict with the
# input_tab_widgets roles

TEXT_COLOR_PKA_ADD = QtGui.QColor('#00b0b3')  # darker cyan
TEXT_COLOR_PKA_REMOVE = QtGui.QColor('#e69500')  # darker orange

AtomType = enum.Enum('AtomType', ['Hydrogen', 'NonHydrogen', 'NonAtom'])


def _get_atom_type(atom):
    """
    Return whether the given `atom` is a hydrogen, non-hydrogen, or non-atom.

    :param atom: An atom represented as a string of the chemical symbol with its
        atom number within the respective structure. (e.g. "H1" for hydrogen)
    :type atom: str

    :return: whether the given `atom` is a hydrogen, non-hydrogen, or non-atom.
    :rtype: AtomType
    """
    atom_match = re.match(r'([A-z]{1,2})(\d+)', atom)
    if atom_match:
        atomic_symbol = atom_match.group(1)
        atomic_number = mm.mmelement_get_atomic_number_by_symbol(atomic_symbol)
        if atomic_number == 1:
            return AtomType.Hydrogen
        elif atomic_number > 1:
            return AtomType.NonHydrogen
    return AtomType.NonAtom


[docs]def filter_hydrogen_from_list(atom_list, return_h):
    """
    Filter a given atom_list to return only hydrogen atoms or only non-hydrogen atoms.

    :param atom_list: a list of workspace atoms represented as strings (e.g. ["H1", "O2", "N3"])
        that can contain both hydrogen and non-hydrogen atoms.
    :type atom_list: list or None

    :param return_h: Whether or not to return only hydrogen atoms. If True, only
        hydrogen atoms are returned. If False, only non-hydrogen atoms are returned.
    :type return_h: bool

    :return: A filtered list of atoms represented as strings with either all hydrogen atoms removed
        or only hydrogen atoms remaining. Returns None if the filtered list is empty. Also returns
        None if the supplied atom_list is empty or is None.
    :rtype: list(str) or NoneType
    """
    if not atom_list:
        return None

    if return_h:
        atom_list = [
            atom for atom in atom_list
            if _get_atom_type(atom) == AtomType.Hydrogen
        ]
    else:
        atom_list = [
            atom for atom in atom_list
            if _get_atom_type(atom) == AtomType.NonHydrogen
        ]

    if not atom_list:
        return None
    return atom_list


ProjEntryTuplePka = namedtuple(
    "ProjEntryTuplePka",
    ("entry_id", "struc", "charge", "spin_mult", "pka_atom"))


[docs]class InputEntriesColumnsPka(object):
    """
    Column constants for the pKa selected entries table
    """
    # Note the plus sign in H⁺ is a superscript plus"
    HEADERS = [
        "ID", "In", "Entry Title", "Charge", "Spin Mult.", "pKa Atom (Add H⁺)",
        "pKa Atom (Remove H⁺)"
    ]
    NUM_COLS = len(HEADERS)
    ID, INCLUSION, TITLE, CHARGE, SPIN_MULT, PKA_ATOM_ADD, PKA_ATOM_REMOVE = list(
        range(NUM_COLS))
    BASIS = -1  # Needed for compatibility with non-pKa classes
    THEORY = -1  # Needed for compatibility with non-pKa classes


[docs]class ProjEntryPka(input_tab_widgets.ProjEntry):
    """
    Builds upon `ProjEntry` by introducing support for storing data for the
    "Add H+" and "Remove H+" pKa columns.
    """
    PKA_ATOM_PROP = "s_m_pKa_atom"

    PKA_VALID, PKA_INVALID, PKA_MISSING, PKA_COLUMN_INVALID = list(range(4))

    COLUMN = InputEntriesColumnsPka

[docs]    def __init__(self, row=None):
        super(ProjEntryPka, self).__init__(row)
        self.pka_atom_add = None
        self.pka_atom_remove = None

[docs]    def update(self, row):
        """
        Builds upon ProjEntryPka.update() to update the pKa atom data
        if there are no pka atoms at all.
        """
        super(ProjEntryPka, self).update(row)
        if self.pka_atom_add is self.pka_atom_remove is None:
            pka_str = row[self.PKA_ATOM_PROP]
            # This is done to ensure that self.pka_atom is a list.
            # Fixes PANEL-1806.
            if pka_str:
                pka_list = [x.strip() for x in pka_str.split(',')]
                self.pka_atom_add = filter_hydrogen_from_list(pka_list,
                                                              return_h=False)
                self.pka_atom_remove = filter_hydrogen_from_list(pka_list,
                                                                 return_h=True)

[docs]    def getStructureWithJagNames(self):
        """
        Return the entry structure with jaguar atom naming applied

        :return: The structure with jaguar atom naming applied
        :rtype: `schrodinger.structure.Structure`
        """

        struc = self.getStructure()
        jaginput.apply_jaguar_atom_naming(struc)
        return struc

[docs]    def reset(self):
        super(ProjEntryPka, self).reset()
        self.pka_atom_add = None
        self.pka_atom_remove = None

[docs]    def getPkaAtom(self, col_num):
        """
        Get either pka_atom_add or pka_atom_remove given the supplied column. Raise
        an error if supplied table column is not one of the two pKa atom columns.

        :param col_num: The table column number for which to return a pKa_atom list.
        :type col_num: int

        :return: A list of pKa atoms or None if there are no pKa atoms in the
            corresponding pka_atom attribute.
        :rtype: list(str) or None
        """
        if col_num == self.COLUMN.PKA_ATOM_ADD:
            return self.pka_atom_add
        elif col_num == self.COLUMN.PKA_ATOM_REMOVE:
            return self.pka_atom_remove
        else:
            raise ValueError

[docs]    def getPkaAtoms(self):
        """
        Get a combined list of all atoms in pka_atom_add and pka_atom_remove.

        :return: A list of pKa atoms or None if there are no pKa atoms in either
            pka_atom attribute.
        :rtype: list(str) or None
        """
        pka_list = list(
            itertools.chain.from_iterable(
                pka_list for pka_list in (self.pka_atom_add,
                                          self.pka_atom_remove)
                if pka_list is not None))
        if not pka_list:
            return None
        return pka_list

[docs]    def checkPkaAtom(self, col_num):
        """
        Make sure that a valid pKa atom(s) is specified in the given pKa atom column.

        :param col_num: The table column number to display pKa data for.
        :type col_num: int

        :return: PKA_VALID if all valid pKa atoms are specified, PKA_INVALID if
            any invalid pKa atom is specified, and PKA_MISSING if no pKa atom is
            specified.
        :rtype: int

        :raises ValueError: if the column number passed into getPkaAtom() corresponds
            to neither the "Add H+" nor the "Remove H+" column.
        """
        pka_atom_names = self.getPkaAtom(col_num)
        if pka_atom_names is None:
            return self.PKA_MISSING
        struc = self.getStructureWithJagNames()
        struc_atom_names = [atom.name for atom in struc.atom]
        for pka_atom_name in pka_atom_names:
            if pka_atom_name not in struc_atom_names:
                return self.PKA_INVALID
        return self.PKA_VALID

[docs]    def checkPkaAtoms(self):
        """
        Aggregate the results of checkPkaAtom() on both pKa atom columns.

        :return: PKA_VALID if all pKa atoms in this project entry are valid,
            PKA_INVALID if any invalid pKa atom is specified anywhere, and
            PKA_MISSING if no pKa atom is specified.
        """
        pka_atom_add_check = self.checkPkaAtom(self.COLUMN.PKA_ATOM_ADD)
        pka_atom_remove_check = self.checkPkaAtom(self.COLUMN.PKA_ATOM_REMOVE)

        if pka_atom_add_check == self.PKA_INVALID or pka_atom_remove_check == self.PKA_INVALID:
            return self.PKA_INVALID
        elif pka_atom_add_check == pka_atom_remove_check:
            if pka_atom_add_check == self.PKA_MISSING:
                return self.PKA_MISSING
        return self.PKA_VALID

[docs]    def getPkaAtomObjs(self):
        """
        Get a list of currently selected pKa atom(s) object(s) from both pKa
        columns.

        :return: If the currently selected pKa atom(s) is valid, returns the
            list of atoms itself. Otherwise, returns None.
        :rtype: list(_StructureAtom) or NoneType
        """
        pka_atom_names = self.getPkaAtoms()

        if pka_atom_names is None:
            return None
        struc = self.getStructureWithJagNames()
        atom_objs = [
            atom_obj for atom_obj in struc.atom
            if atom_obj.name in pka_atom_names
        ]
        if atom_objs:
            return atom_objs
        else:
            return None


[docs]class PickingModes(enum.Enum):
    MANUAL = 1
    AUTO = 2
    SMARTS = 3


[docs]class InputEntriesModelPka(input_tab_widgets.InputEntriesModel):
    """
    The data model for the pKa selected entries table. There are two columns that
    hold pKa atom values, "pKa Atom (Add H+)" and "pKa Atom (Remove H+)". The idea being
    that only non H-atoms can be in the "Add H+" column and only H-atoms can be in the
    "Remove H+ column.

    :ivar set_pka_marker: A signal emitted when a new pKa atom should be marked
        in the workspace.  Emitted with two arguments:

          - The entry id of the structure to be marked (str)
          - The atom to be marked (`schrodinger.structure._StructureAtom`)
    :vartype set_pka_marker: `PyQt5.QtCore.pyqtSignal`

    :ivar projUpdated: A signal emitted when the project changes which
        should trigger a SMARTS search again if picking mode is SMARTS
    """

    COLUMN = InputEntriesColumnsPka
    PKA_COLUMNS = (COLUMN.PKA_ATOM_ADD, COLUMN.PKA_ATOM_REMOVE)
    set_pka_marker = QtCore.pyqtSignal(str, object)
    projUpdated = QtCore.pyqtSignal()
    ROW_CLASS = ProjEntryPka

[docs]    def __init__(self, parent):
        super(InputEntriesModelPka, self).__init__(parent)
        self.rowsAboutToBeRemoved.connect(self.removePkaMarkers)
        self.rowsInserted.connect(self.addPkaMarkers)
        self.mode = PickingModes.MANUAL
        self.smarts_is_defined = False

[docs]    def setPickingMode(self, mode):
        """
        Set the picking mode. This impacts what is shown in the pKa atom columns.

        :param mode: the new picking mode
        :type mode: PickingModes
        """
        self.mode = mode
        index_tl = self.index(0, self.COLUMN.PKA_ATOM_ADD)
        index_br = self.index(self.rowCount() - 1, self.COLUMN.PKA_ATOM_REMOVE)
        self.dataChanged.emit(index_tl, index_br)

[docs]    def data(self, index, role=Qt.DisplayRole):
        if index.column() in self.PKA_COLUMNS:
            return self._pkaData(index.row(), index.column(), role)
        else:
            return super(InputEntriesModelPka, self).data(index, role)

[docs]    def headerData(self, section, orientation, role=Qt.DisplayRole):
        """
        Override InputEntriesModel to allow special coloring of the pKa column headers.

        These colors are darker than the WS atom marker colors because they appear
        against a light background.
        """
        if role == Qt.ForegroundRole:
            if section == self.COLUMN.PKA_ATOM_ADD:
                return QtGui.QBrush(TEXT_COLOR_PKA_ADD)
            elif section == self.COLUMN.PKA_ATOM_REMOVE:
                return QtGui.QBrush(TEXT_COLOR_PKA_REMOVE)
        return super().headerData(section, orientation, role)

[docs]    def flags(self, index):
        """
        Enable or disable the pka column depending on whether manual editing
        is enabled or not

        :param index: The model index to get flags for.
        :type index: `QtCore.QModelIndex`
        """
        flags = super().flags(index)
        col = index.column()
        if col in self.PKA_COLUMNS:
            if self.mode == PickingModes.MANUAL:
                flags |= Qt.ItemIsEnabled
            else:
                flags &= ~Qt.ItemIsEnabled
        return flags

    def _pkaData(self, row_num, col_num, role):
        """
        Return data for one of the two pKa atom columns.

        :param row_num: The table row number to display data for
        :type row_num: int

        :param col_num: The table column number to display data for
        :type col_num: int

        :param role: The role to retrieve data for
        :type role: int
        """

        proj_row = self.cur_table_rows[row_num]
        pka_atoms = proj_row.getPkaAtom(col_num)

        if role == SORT_ROLE:
            return pka_atoms
        elif role == Qt.DisplayRole:
            if self.mode == PickingModes.AUTO:
                return self._getPkaColumnTextAuto(row_num)
            if pka_atoms:
                pka_text = ", ".join([str(x) for x in pka_atoms])
                return pka_text
            else:
                if self.mode == PickingModes.SMARTS:
                    return self._getPkaColumnTextSmarts(row_num)
            return None
        elif role in (Qt.ToolTipRole, Qt.BackgroundRole):
            if proj_row.checkPkaAtom(col_num) != proj_row.PKA_INVALID:
                return None
            elif role == Qt.ToolTipRole:
                return ("The specified atom name does not appear in this "
                        "structure.")
            elif role == Qt.BackgroundRole:
                return self.ERROR_BACKGROUND_BRUSH
        elif role == Qt.TextAlignmentRole:
            return Qt.AlignCenter
        elif role == JAG_NAME_STRUCTURE_ROLE:
            return proj_row.getStructureWithJagNames()
        elif role == PKA_ATOM_ROLE:
            return proj_row.getPkaAtomObjs()

    def _getPkaColumnTextAuto(self, row_num):
        """
        Return text for the pka column when the picking mode is AUTO

        :param row_num: The table row number for which the returned text
            will correspond to
        :type row_num: int

        :return: The text to use in the display role for the pka column
        :rtype: str
        """
        if row_num == 0:
            return "(auto-search)"
        else:
            return '"'

    def _getPkaColumnTextSmarts(self, row_num):
        """
        Return text to display in the pka column when the picking mode is
        SMARTS and a pka atom was not set.

        :param row_num: The table row number for which the returned text
            will correspond to
        :type row_num: int

        :return: The text to use in the display role for the pka column
        :rtype: str
        """
        if self.smarts_is_defined:
            return "(not found)"
        else:
            if row_num == 0:
                return "(undefined)"
            else:
                return '"'

[docs]    def setData(self, index, value, role=Qt.EditRole):
        if index.column() not in self.PKA_COLUMNS:
            return super(InputEntriesModelPka, self).setData(index, value, role)
        elif role != Qt.EditRole:
            return False

        self._setPkaAtomsToIndex(index, value)
        return True

    def _setPkaAtomsToIndex(self, index, value):
        """
        Assign one or more pKa atoms to a particular table index.
        Filters out any atoms that don't belong in that column
        (e.g. a hydrogen atom in the "Add H+" column.)

        :param index: The index of the active selection delegate.
        :type index: QtCore.QModelIndex

        :param value: The pka atoms to set for the index. Given as a list of
            Jaguar atom names (e.g. ["H1", "C9"]) or None.
        :type value: list(str) or None
        """
        proj_row = self.cur_table_rows[index.row()]
        if value:
            # make value the same ordered list of unique items in the list
            value = list(dict.fromkeys(value))
        if index.column() == self.COLUMN.PKA_ATOM_ADD:
            pka_list = filter_hydrogen_from_list(value, return_h=False)
            proj_row.pka_atom_add = pka_list
        elif index.column() == self.COLUMN.PKA_ATOM_REMOVE:
            pka_list = filter_hydrogen_from_list(value, return_h=True)
            proj_row.pka_atom_remove = pka_list
        self.dataChanged.emit(index, index)
        self.set_pka_marker.emit(proj_row.entry_id, proj_row.getPkaAtomObjs())

    def _setPkaAtomsFromSmarts(self, row_num, value):
        """
        Assign one or more pka atoms to a row in the table. Clears any
        pKa atom selections not specified by SMARTS.

        :param row_num: The row to which the pka atoms will get added
        :type row_num: int

        :param value: The pka atoms to set for the row. Given as a list of
            Jaguar atom names (e.g. ["H1", "C9"]) or None.
        :type value: list(str) or None
        """
        index_add = self.index(row_num, self.COLUMN.PKA_ATOM_ADD)
        index_remove = self.index(row_num, self.COLUMN.PKA_ATOM_REMOVE)
        self._setPkaAtomsToIndex(index_add, value)
        self._setPkaAtomsToIndex(index_remove, value)

[docs]    def updatePkaAtomsFromSmarts(self, smarts_models):
        """
        Update pKa atoms of entry rows based on patterns given by a list of
        SMARTS patterns and atom positions.

        :param smarts_models: a list of smarts patterns to search for in each
            structure. The list may be empty.
        :type smarts_models: list(SmartsPageModel)
        """
        self.smarts_is_defined = False
        for row_num, row in enumerate(self.cur_table_rows):
            st = row.getStructureWithJagNames()
            atoms = []
            for m in smarts_models:
                smarts = m.smarts
                if not smarts:
                    continue
                if analyze.validate_smarts_canvas(smarts)[0] is False:
                    continue
                self.smarts_is_defined = True
                matches = analyze.evaluate_smarts_canvas(st, m.smarts)
                if matches:
                    for match in matches:
                        a_num = match[m.atom_pos - 1]
                        atom_name = st.atom[a_num].name
                        atoms.append(atom_name)
            self._setPkaAtomsFromSmarts(row_num, atoms)

[docs]    def projectUpdated(self):
        super().projectUpdated()
        self.projUpdated.emit()

[docs]    def getStructures(self):
        all_data = []
        for proj_row in self.cur_table_rows:
            row_data = proj_row.getStrucChargeAndSpinMult()
            (eid, struc, charge, spin_mult) = row_data
            pka_atom_names = proj_row.getPkaAtoms()
            cur_data = ProjEntryTuplePka(eid, struc, charge, spin_mult,
                                         pka_atom_names)
            all_data.append(cur_data)
        return all_data

[docs]    def checkPkaAtoms(self):
        """
        Make sure that all structure have a valid pKa atom selected

        :return: A tuple of
              - A list of structures with invalid pKa atoms
              - A list of structures with no pKa atom
        :rtype: tuple
        """

        bad_atoms = []
        missing_atoms = []
        for proj_row in self.cur_table_rows:
            row_status = proj_row.checkPkaAtoms()
            if row_status == proj_row.PKA_MISSING:
                missing_atoms.append(proj_row.title)
            elif row_status == proj_row.PKA_INVALID:
                bad_atoms.append(proj_row.title)
        return bad_atoms, missing_atoms

[docs]    def addPkaMarkers(self, index, start_row, end_row):
        """
        Add pKa workspace atom markers for the specified rows

        :param index: Not used, but present for Qt compatability

        :param start_row: The first row to add a pKa marker to
        :type start_row: int

        :param end_row: The last row to add a pKa marker to
        :type end_row: int
        """

        for cur_row in range(start_row, end_row + 1):
            proj_row = self.cur_table_rows[cur_row]
            self.set_pka_marker.emit(proj_row.entry_id,
                                     proj_row.getPkaAtomObjs())

[docs]    def removePkaMarkers(self, index, start_row, end_row):
        """
        Remove the pKa workspace atom markers for the specified rows

        :param index: Not used, but present for Qt compatibility

        :param start_row: The first row to add a pKa marker to
        :type start_row: int

        :param end_row: The last row to add a pKa marker to
        :type end_row: int
        """

        for cur_row in range(start_row, end_row + 1):
            proj_row = self.cur_table_rows[cur_row]
            self.set_pka_marker.emit(proj_row.entry_id, None)


[docs]class InputEntriesViewPka(table_helper.SampleDataTableViewMixin,
                          input_tab_widgets.InputEntriesView):
    """
    The view for the pKa selected entries table. SampleDataTableViewMixin makes
    sure the default table width is large enough to make all columns visible.

    :ivar set_pka_marker: A signal emitted when a new pKa atom should be marked
        in the workspace.  Emitted with two arguments:

          - The entry id of the structure to be marked (str)
          - List of atoms (`schrodinger.structure._StructureAtom`) to be marked
    :vartype set_pka_marker: `PyQt5.QtCore.pyqtSignal`
    """
    COLUMN = InputEntriesColumnsPka
    set_pka_marker = QtCore.pyqtSignal(str, object)

    def _setDelegates(self, parent):
        self._setInclusionDelegate(parent)
        self._setChargeAndSpinMultDelegates(parent)

        self._default_pka_atom_delegate_add = self.itemDelegateForColumn(
            self.COLUMN.PKA_ATOM_ADD)
        self._default_pka_atom_delegate_remove = self.itemDelegateForColumn(
            self.COLUMN.PKA_ATOM_REMOVE)

        self._manual_pka_atom_delegate = AtomSelectionDelegate(parent)
        self._manual_pka_atom_delegate.commitDataToSelected.connect(
            self.commitDataToSelected)
        self._manual_pka_atom_delegate.set_pka_marker.connect(
            self.set_pka_marker.emit)

        self.setEditablePkaAtomDelegate(True)

[docs]    def setEditablePkaAtomDelegate(self, editable):
        """
        Add the delegates to the pKa column

        :param editable: Whether to set the editable default message delegate
            or to set the Qt default delegate
        """
        if editable:
            self.setItemDelegateForColumn(self.COLUMN.PKA_ATOM_ADD,
                                          self._manual_pka_atom_delegate)
            self.setItemDelegateForColumn(self.COLUMN.PKA_ATOM_REMOVE,
                                          self._manual_pka_atom_delegate)
        else:
            self.setItemDelegateForColumn(self.COLUMN.PKA_ATOM_ADD,
                                          self._default_pka_atom_delegate_add)
            self.setItemDelegateForColumn(
                self.COLUMN.PKA_ATOM_REMOVE,
                self._default_pka_atom_delegate_remove)


[docs]class SmartsPageModel(parameters.CompoundParam):
    """
    Model for a single page which contains the a single smarts string and
    single atom position index.
    """
    atom_pos = parameters.IntParam(1)
    smarts = parameters.StringParam("")


[docs]class PkaPage(mappers.MapperMixin, basewidgets.BaseWidget):

    model_class = SmartsPageModel
    ui_module = pka_smarts_page_ui

[docs]    def initSetUp(self):
        super().initSetUp()
        self.ui.smarts_le.textChanged.connect(self.updateSpinBoxLimits)
        self.ui.from_selection_btn.clicked.connect(self.getFromSelection)

[docs]    def updateSpinBoxLimits(self):
        """
        Update limits of atom position spin box to reflect the number of
        atoms in the SMARTS string
        """
        smarts = self.ui.smarts_le.text()
        if smarts and analyze.validate_smarts_canvas(smarts)[0]:
            max_index = int(analyze.count_atoms_in_smarts(smarts))
            if max_index < self.ui.atom_position_sb.value():
                self.ui.atom_position_sb.setValue(max_index)
            self.ui.atom_position_sb.setMaximum(max_index)

[docs]    def defineMappings(self):
        return {
            self.ui.smarts_le: SmartsPageModel.smarts,
            self.ui.atom_position_sb: SmartsPageModel.atom_pos
        }

[docs]    def getFromSelection(self):
        """
        Get smarts from selected atoms in workspace and add it to the current
        page
        """
        try:
            self.model.smarts = maestro.selected_atoms_get_smarts()
        except ValueError:
            # selected_atoms_get_smarts raises ValueError when too many atoms
            # are selected, atoms are disconnected, no atoms are selected, etc.
            self.model.smarts = ""


[docs]class SmartsSelector(pop_up_widgets.PopUp, basewidgets.BaseWidget):
    """
    A popup widget that allows users to specify any number of SMARTS strings
    to specify pka atoms as well as an atom position which specifies the atom
    in the SMARTS
    """

    closed = QtCore.pyqtSignal()
    ui_module = define_smarts_panel_ui

[docs]    def setup(self):
        """Needs to be implemented for PopUp widgets"""

[docs]    def initSetUp(self):
        super().initSetUp()
        self.addPage()
        self.ui.define_another_btn.clicked.connect(self.addPage)
        self.ui.remove_btn.clicked.connect(self.removeCurrentPage)
        self.ui.next_page_btn.clicked.connect(self.next)
        self.ui.back_page_btn.clicked.connect(self.back)

[docs]    def addPage(self):
        """
        Add an additional page for specifying SMARTS
        """
        new_widget = PkaPage(self)

        self.ui.stacked_widget.addWidget(new_widget)
        self.ui.stacked_widget.setCurrentIndex(self.ui.stacked_widget.count() -
                                               1)
        self.updatePagination()

[docs]    def removeCurrentPage(self):
        """
        Remove the current page. If there was only one page before removing,
        add another after so there is always one page.
        """
        current_widget = self.ui.stacked_widget.currentWidget()

        self.ui.stacked_widget.removeWidget(current_widget)

        if self.ui.stacked_widget.count() == 0:
            self.addPage()
        self.updatePagination()

[docs]    def next(self):
        """
        Show next page of stacked widget
        """
        current_idx = self.ui.stacked_widget.currentIndex()
        self.ui.stacked_widget.setCurrentIndex(current_idx + 1)
        self.updatePagination()

[docs]    def back(self):
        """
        Show previous page of stacked widget
        """
        current_idx = self.ui.stacked_widget.currentIndex()
        self.ui.stacked_widget.setCurrentIndex(current_idx - 1)
        self.updatePagination()

[docs]    def updatePagination(self):
        """
        Enable or disable forward and back button based on current page index
        and number of pages and update label indicating current page
        """
        current_idx = self.ui.stacked_widget.currentIndex()
        count = self.ui.stacked_widget.count()

        self.ui.back_page_btn.setEnabled(current_idx != 0)
        self.ui.next_page_btn.setEnabled(current_idx != count - 1)

        label = f"{current_idx+1}/{count}"
        self.ui.pagination_label.setText(label)

[docs]    def getModels(self):
        """
        Get the models for the Smarts Pages

        :return: A list of the models
        :rtype: list(SmartsPageModel)
        """
        count = self.ui.stacked_widget.count()
        return [self.ui.stacked_widget.widget(i).model for i in range(count)]

[docs]    def closeEvent(self, ev):
        self.closed.emit()
        super().closeEvent(ev)


[docs]class SmartsSelectorButton(pop_up_widgets.ToolButtonWithPopUp):
    """
    A tool button that opens the Smarts Selector when pressed
    """

[docs]    def __init__(self, parent):
        """
        :param parent: The Qt parent widget
        :type parent: `PyQt5.QtWidgets.QWidget`
        """
        super().__init__(parent, SmartsSelector, Qt.DownArrow, "Define")
        self.setPopupHalign(self.ALIGN_RIGHT)
        self.setPopupValign(self.ALIGN_BOTTOM)
        self.setFocusPolicy(QtCore.Qt.StrongFocus)

[docs]    def getModels(self):
        """
        Get the list of models for each SMARTS page

        :return: The list of moels for each SMARTS page
        :rtype: list(SmartsPageModel)
        """
        return self._pop_up.getModels()

[docs]    def showSmartsSelector(self):
        """
        Show the popup.
        This gets called by the input tab to manually show the popup
        """
        self._pop_up.show()


[docs]class AtomSelectionDelegate(input_tab_widgets.CommitMultipleDelegate,
                            qt_delegates.DefaultMessageDelegate):
    """
    A delegate for selecting a pKa atom.  The atom name can either be typed into
    the line edit or selected from the workspace.  A tool tip containing
    instructions will appear when the editor is first open and any time the user
    hovers their mouse over the editor.  If the user clicks on an atom from the
    wrong structure, the atom will be ignored and a tool tip warning will
    appear.  Clicking on an atom does not close the editor so that the user can
    immediately pick a different atom if desired.  Valid atoms will be immediately
    added to the model upon clicking.

    :ivar set_pka_marker: A signal emitted when a new pKa atom should be marked
        in the workspace.  Emitted with two arguments:

          - The entry id of the structure to be marked (str)
          - List of atoms (`schrodinger.structure._StructureAtom`) to be marked
    :vartype set_pka_marker: `PyQt5.QtCore.pyqtSignal`
    """

    MAESTRO_BANNER_REMOVE_H = "Pick protons to be removed"
    MAESTRO_BANNER_ADD_H = "Pick atoms to add a proton to"
    TOOL_TIP_INSTRUCTIONS = ("Click an atom in the workspace to\n"
                             "set it as the pKa atom or type the\n"
                             "atom name here.")
    TOOL_TIP_WRONG_EID = ("The atom you selected is not\n"
                          "part of this structure.")
    set_pka_marker = QtCore.pyqtSignal(str, object)

[docs]    def __init__(self, parent):
        super(AtomSelectionDelegate, self).__init__(parent)
        self.closeEditor.connect(maestro.picking_stop)
        self.closeEditor.connect(self._resetMarker)

[docs]    def createEditor(self, parent, option, index):
        editor = QtWidgets.QLineEdit(parent)
        editor.setGeometry(option.rect)
        editor.index = index
        editor.setToolTip(self.TOOL_TIP_INSTRUCTIONS)
        self._showToolTipTimer(editor, self.TOOL_TIP_INSTRUCTIONS, 300)
        # The 300 ms delay is to make sure that painting the editor has
        # finished.  Without it, the tooltip will appear and immediately
        # disappear (at least on Windows).  Moving the _showToolTip() call
        # to the end of the line edit's show() function doesn't avoid the
        # need for this delay.

        self._ensureEntryIncluded(index)
        if index.column() == InputEntriesColumnsPka.PKA_ATOM_ADD:
            self._startPicking(editor, index, self.MAESTRO_BANNER_ADD_H)
        elif index.column() == InputEntriesColumnsPka.PKA_ATOM_REMOVE:
            self._startPicking(editor, index, self.MAESTRO_BANNER_REMOVE_H)
        editor.new_atom_marked = False
        return editor

    def _startPicking(self, editor, index, message):
        pick_func = lambda atom_num: self._setEditorAtom(
            atom_num, editor, index)
        maestro.picking_atom_start(message, pick_func)

    def _ensureEntryIncluded(self, index):
        """
        Make sure that an entry is included in the workspace.  If it isn't,
        include only that structure.

        :param index: The model index corresponding to the structure to include
        :type index: `QtCore.QModelIndex`
        """

        model = index.model()
        row = index.row()
        in_col = model.COLUMN.INCLUSION
        in_index = model.index(row, in_col)
        included = in_index.data()
        if not included:
            model.setData(in_index, WS_INCLUDE_ONLY)

[docs]    def setEditorData(self, editor, index):
        text = index.data()
        if text is None:
            text = ""
        editor.setText(text)

[docs]    def setModelData(self, editor, model, index):
        """
        Set the model data at the appropriate index according to the atoms currently
        present in the editor. Gets called every time a valid atom is clicked as well
        as when the delegate is closed (i.e. can only delete atoms upon closing the
        delegate).

        An editor may only add/remove atoms that belong in its own column.
        I.e. the "Add H+" editor can only add/remove non-hydrogen atoms
        to the model. Conversely, the "Remove H+" editor may only add/remove
        hydrogen atoms to the model.

        :param editor: The line edit to enter the atom name into
        :type editor: PyQt5.QtWidgets.QLineEdit

        :param model: The table model to edit at the given index.
        :type model: InputEntriesModelPka

        :param index: The index of the table cell being edited
        :type index: PyQt5.QtCore.QModelIndex
        """
        text = editor.text()
        if text == '':
            pka_list = None
        else:
            pka_list = [s.strip() for s in text.split(',')]
        model.setData(index, pka_list)

    def _setEditorAtom(self, atom_num, editor, index):
        """
        Respond to the user clicking an atom in the workspace.  Do nothing
        if the user clicks on a hydrogen atom while the the "Remove H+" delegate
        is active or if the user clicks on a non-hydrogen atom while the "Add H+"
        delegate is active. If the atom does not belong to this structure, display
        a warning.

        Otherwise, enter the Jaguar-ified atom name into the line edit, display
        a marker on the atom, and update the model to reflect the selection.
        If atom is already included in the list of pka atoms we remove it from
        the list.

        :param atom_num: The atom number (relative to the workspace structure)
            that was selected.
        :type atom_num: int

        :param editor: The line edit to enter the atom name into
        :type editor: PyQt5.QtWidgets.QLineEdit

        :param index: The index of the cell being edited
        :type index: PyQt5.QtCore.QModelIndex
        """

        model = index.model()
        eid_index = model.index(index.row(), model.COLUMN.ID)
        editor_eid = eid_index.data()
        struc = maestro.workspace_get()
        atom = struc.atom[atom_num]

        is_hydrogen_atom = atom.element == 'H'
        if is_hydrogen_atom and index.column() == model.COLUMN.PKA_ATOM_ADD:
            return
        elif (not is_hydrogen_atom and
              index.column() == model.COLUMN.PKA_ATOM_REMOVE):
            return
        elif atom.entry_id != editor_eid:
            self._showToolTip(editor, self.TOOL_TIP_WRONG_EID)
        else:
            jag_name_struc = index.data(JAG_NAME_STRUCTURE_ROLE)
            jag_name_atom = jag_name_struc.atom[atom.number_by_entry]
            jag_name = jag_name_atom.name
            self._updateEditorText(editor, jag_name)
            pka_text = editor.text()
            pka_atoms = self._getAtomSts(jag_name_struc, pka_text)
            editor.new_atom_marked = True
            self.set_pka_marker.emit(editor_eid, pka_atoms)
            self.setModelData(editor, model, index)

    def _getAtomSts(self, struc, text):
        """
        This function returns a list of pka atoms by parsing a given text
        string that contains atom names and matching them to the atom names
        in the input structure. This is needed when user edits pka atom
        names in the input tab table.

        :param struc: current structure
        :type struc: structure.Structure

        :param text: text string that contains comma separated list
            of pka atoms.
        :type text: str

        :return: list of structure atoms or None
        :rtype: list or None
        """
        if text == "":
            return None
        else:
            atoms = [s.strip() for s in text.split(',')]
            pka_atoms = []
            for atom in struc.atom:
                if atom.name in atoms:
                    pka_atoms.append(atom)
            return pka_atoms

    def _updateEditorText(self, editor, atom_name):
        """
        This function checks whether a given atom_name is present in the
        editor's string. If not found, add it to editor's text. If found
        delete it from the editor text and reformat the text.

        :param editor: The line edit to enter the atom name into
        :type editor: PyQt5.QtWidgets.QLineEdit

        :param atom_name: Name of picked atom
        :type atom_name: string
        """

        text = editor.text()
        if text == "":
            text = atom_name
        else:
            pka_atom_names = [s.strip() for s in text.split(',')]
            if atom_name in pka_atom_names:
                pka_atom_names.remove(atom_name)
            else:
                pka_atom_names.append(atom_name)
            text = ', '.join(pka_atom_names)
        editor.setText(text)

    def _resetMarker(self, editor, hint):
        """
        If a pKa atom was selected via the user clicking in the workspace, reset
        the currently marked atom to the table value.

        :param editor: The delegate editor
        :type editor: PyQt5.QtWidgets.QLineEdit

        :param hint: Not used, but present for Qt compatibility
        """

        if editor.new_atom_marked:
            index = editor.index
            model = index.model()
            eid_index = model.index(index.row(), model.COLUMN.ID)
            editor_eid = eid_index.data()
            pka_atom_objs = index.data(PKA_ATOM_ROLE)
            self.set_pka_marker.emit(editor_eid, pka_atom_objs)

[docs]    def eventFilter(self, editor, event):
        """
        Make sure that the editor doesn't close when the user clicks on another
        window since that will prevent the user from being able to click on an
        atom.

        :param editor: The pKa atom line edit
        :type editor: PyQt5.QtWidgets.QWidget

        :param event: A Qt event
        :param event: PyQt5.QtCore.QEvent

        :note: We don't need to worry about the case where the user clicks on a
            different widget in the pKa panel after selecting an atom.  Since the
            editor was the last widget with focus in the pKa panel, it will receive
            another FocusOut event when the other widget receives focus, and that
            FocusOut event will cause the editor to close.
        """

        if event.type() == event.FocusOut and not editor.isActiveWindow():
            return True
        else:
            return super(AtomSelectionDelegate, self).eventFilter(editor, event)

    def _showToolTip(self, widget, text):
        """
        Show a tool tip centered on the specified widget

        :param widget: The widget to center the tool tip on
        :type widget: PyQt5.QtWidgets.QWidget

        :param text: The text to display in the tool tip
        :type text: str
        """

        try:
            tool_tip_point = widget.rect().center()
        except RuntimeError:
            # If this function was called via _showToolTipTimer(), then it's
            # possible that the delegate was closed before the timer fired.  If
            # that happened, then the underlying C++ object was deleted and
            # we'll get a RuntimeError here.  In that case, we can give up on
            # displaying a tooltip and ignore the runtime error.
            return
        tool_tip_point = widget.mapToGlobal(tool_tip_point)
        QtWidgets.QToolTip.showText(tool_tip_point, text, widget)

    def _showToolTipTimer(self, widget, text, delay=0):
        """
        Show a tool tip centered on the specified widget after a delay

        :param widget: The widget to center the tool tip on
        :type widget: PyQt5.QtWidgets.QWidget

        :param text: The text to display in the tool tip
        :type text: str

        :param delay: How long of a delay
        :type delay: int
        """

        func = lambda: self._showToolTip(widget, text)
        QtCore.QTimer.singleShot(delay, func)