"""
Classes for displaying site selection information and options for organometallic
complexes of different VSEPR geometries
Copyright Schrodinger, LLC. All rights reserved.
"""
import os.path
from past.utils import old_div
from schrodinger.application.matsci import buildcomplex
from schrodinger.application.matsci import builderwidgets
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.ui.qt import swidgets
from schrodinger.ui.qt import utils as qtutils
FONT_SIZE = 30
SELECTED_TEXT = QtGui.QColor('red')
UNSELECTED_TEXT = QtGui.QColor('black')
UNCLICKABLE_TEXT = QtCore.Qt.lightGray
LETTERS = 'ABCDEFGHIJKLMNOP'
MASTER_SIZE = 500
MASTER_MOVEXBY = 50
MASTER_MOVEYBY = 80
DESIGNATOR_SHIFT = -50
FILE_DIR = os.path.split(__file__)[0]
[docs]class CoordinationSite(QtWidgets.QGraphicsTextItem):
    """
    QGraphicsTextItem for displaying the number of the coordination site
    """
    selected = QtCore.pyqtSignal(int)
    """ Signal when the site is selected """
    deselected = QtCore.pyqtSignal(int)
    """ Signal when the site is deselected (a selected site is clicked again """
[docs]    def __init__(self, value):
        """
        Create a CoordinationSite object
        :type value: int
        :param value: The number the user will see for this site.  Note that
            this will be one greater than the index into python arrays for this
            site.
        """
        self.string_value = str(value)
        QtWidgets.QGraphicsTextItem.__init__(self, self.string_value)
        font = self.font()
        font.setPixelSize(FONT_SIZE)
        self.setFont(font)
        self.value = value
        self.is_selected = False
        self.shifted_designator = False
        self.original_pos = self.pos()
        self.designated_pos = self.pos() 
[docs]    def setClickable(self, state):
        """
        Set whether this site is clickable
        :type state: bool
        :param state: Whether the site can be clicked on or not
        """
        if not state:
            color = UNCLICKABLE_TEXT
        else:
            color = UNSELECTED_TEXT
        self.setDefaultTextColor(color)
        self.setEnabled(state) 
[docs]    def mousePressEvent(self, event):
        """
        Catch when the mouse clicks on this site and select/deselect the site
        :type event: QMouseEvent
        :param event: The event that generated this call
        """
        self.select(not self.is_selected)
        return QtWidgets.QGraphicsTextItem.mousePressEvent(self, event) 
[docs]    def moveToPosition(self, xval, yval):
        """
        Move the label to its starting position - should only be called at the
        beginning.
        :type xval: int
        :type xval: The number of pixels to shift the label in the X direction
        :type yval: int
        :type yval: The number of pixels to shift the label in the Y direction
        """
        self.moveBy(xval, yval)
        self.original_pos = self.pos()
        self.designated_pos = self.pos() 
[docs]    def setShiftLeft(self):
        """
        Set this label as one that shifts to the left when a ligand designator
        is added to the label.
        """
        self.designated_pos.setX(self.designated_pos.x() + DESIGNATOR_SHIFT)
        self.shifted_designator = True 
[docs]    def select(self, state=True, released=False):
        """
        Select or deselect the label and emit the proper signal
        :type state: bool
        :param state: True if the label is to be selected, False if not
        :type released: bool
        :type released: True if the call to this function is because a ligand
            released it rather than via user-click.  If so, don't emit the
            deselected signal.
        """
        font = self.font()
        if state:
            color = SELECTED_TEXT
            weight = font.Bold
            self.selected.emit(self.value - 1)
        else:
            color = UNSELECTED_TEXT
            weight = font.Normal
            self.setPlainText(self.string_value)
            self.setPos(self.original_pos)
            if not released:
                self.deselected.emit(self.value - 1)
        self.setDefaultTextColor(color)
        font.setWeight(weight)
        self.setFont(font)
        self.is_selected = state 
[docs]    def addDesignator(self, text):
        """
        Add a ligand designator to the label to indicate that it is currently
        assigned to that ligand coordination site.
        :type text: str
        :param text: The designator to add to the site label
        """
        if self.shifted_designator:
            self.setPlainText(text + ' ' + self.string_value)
        else:
            self.setPlainText(self.string_value + ' ' + text)
        self.setPos(self.designated_pos) 
[docs]    def released(self, index):
        """
        A ligand has released this coordination site because a different site
        was picked.
        """
        if index == self.value - 1:
            self.select(False, released=True)  
[docs]class OctahedralSiteSelector(QtWidgets.QFrame):
    """
    A Panel that puts up an octahedral geometry and allows the user to select
    sites.
    """
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    MOVEXBY = MASTER_MOVEXBY
    MOVEYBY = MASTER_MOVEYBY
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [SIZE - 90, CENTER - 23]
    LOCATION[1] = [CENTER - 2 - old_div(FONT_SIZE, 4), 60]
    LOCATION[2] = [325, 375]
    LOCATION[3] = [CENTER - 2 - old_div(FONT_SIZE, 4), SIZE - 100]
    LOCATION[4] = [75, CENTER - 23]
    LOCATION[5] = [158, 80]
    IMAGE_FILE = os.path.join(FILE_DIR, 'octahedral.png')
[docs]    def __init__(self):
        """
        Create an OctahedralSiteSelector object
        """
        QtWidgets.QFrame.__init__(self)
        layout = swidgets.SHBoxLayout(self)
        # Scene/View containers
        self.scene = QtWidgets.QGraphicsScene(0, 0, self.SIZE, self.SIZE)
        self.view = QtWidgets.QGraphicsView(self.scene)
        layout.addWidget(self.view)
        # Get the base geometry image from a file and load it.
        self.pixmap = QtGui.QPixmap(self.IMAGE_FILE)
        self.pix_item = self.scene.addPixmap(self.pixmap)
        self.pix_item.moveBy(self.MOVEXBY, self.MOVEYBY)
        # Add the site labels to the geometry
        self.sites = [
            CoordinationSite(x) for x in range(1,
                                               len(self.LOCATION) + 1)
        ]
        for index, site in enumerate(self.sites):
            self.scene.addItem(site)
            site.moveToPosition(*self.LOCATION[index])
        self.markShiftedDesignators() 
[docs]    def addLigandDesignator(self, site, designator):
        """
        Add a ligand site desginator to a selected coordination site label
        :type site: int
        :type site: The site whose label should get the designation.  This is a
            Python list index, so should have a value 1 less than the value the
            user sees in the label.
        :type designator: str
        :param designator: The string to add to the site label
        """
        self.sites[site].addDesignator(designator) 
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        This should be re-implemented for each specific geometry image.
        """
        self.sites[4].setShiftLeft()
        self.sites[5].setShiftLeft() 
[docs]    def deselectSelectedSites(self, ignore=None):
        """
        Deselect selected sites other than the ignored one
        :type ignore: int or None
        :param ignore: Deselect any site other than this one (zero-based). If
            ignore=None, all sites are deselected.
        """
        for site in self.sites:
            if site.is_selected and (ignore is None or
                                     site.value != ignore + 1):
                site.select(state=False) 
[docs]    def quietlySelectSite(self, index):
        """
        Select the given site but do not emit any selection signals
        :type index: int
        :param index: The zero-based index of the site to ignore
        """
        site = self.sites[index]
        with qtutils.suppress_signals(site):
            site.select(state=True) 
[docs]    def setSiteClickable(self, index, state):
        """
        Set whether the given site is clickable
        :type index: int
        :param index: The zero-based index of the site
        :type state: bool
        :param state: Whether the site should be clickable
        """
        self.sites[index].setClickable(state)  
[docs]class TrigonalBipyramidalSiteSelector(OctahedralSiteSelector):
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    MOVEXBY = 60
    MOVEYBY = 45
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [CENTER - 20 - old_div(FONT_SIZE, 4), 20]
    LOCATION[1] = [CENTER - 20 - old_div(FONT_SIZE, 4), SIZE - 70]
    LOCATION[2] = [SIZE - 80, CENTER - 24]
    LOCATION[3] = [55, CENTER - 122]
    LOCATION[4] = [55, CENTER + 75]
    IMAGE_FILE = os.path.join(FILE_DIR, 'trigonal_bipyramidal.png')
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        """
        self.sites[3].setShiftLeft()
        self.sites[4].setShiftLeft()  
[docs]class SquarePlanarSiteSelector(OctahedralSiteSelector):
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    MOVEXBY = MASTER_MOVEXBY
    MOVEYBY = MASTER_MOVEYBY
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [SIZE - 90, CENTER - 23]
    LOCATION[1] = [CENTER - 2 - old_div(FONT_SIZE, 4), 60]
    LOCATION[2] = [CENTER - 2 - old_div(FONT_SIZE, 4), SIZE - 100]
    LOCATION[3] = [75, CENTER - 23]
    IMAGE_FILE = os.path.join(FILE_DIR, 'square_planar.png')
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        """
        self.sites[3].setShiftLeft()  
[docs]class TetrahedralSiteSelector(OctahedralSiteSelector):
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    MOVEXBY = MASTER_MOVEXBY
    MOVEYBY = MASTER_MOVEYBY
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [CENTER - 2 - old_div(FONT_SIZE, 4), 60]
    LOCATION[1] = [SIZE - 120, 310]
    LOCATION[2] = [95, 310]
    LOCATION[3] = [CENTER - 4 - old_div(FONT_SIZE, 4), SIZE - 100]
    IMAGE_FILE = os.path.join(FILE_DIR, 'tetrahedral.png')
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        """
        self.sites[2].setShiftLeft()  
[docs]class TrigonalPlanarSiteSelector(OctahedralSiteSelector):
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    MOVEXBY = 60
    MOVEYBY = 40
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [SIZE - 85, CENTER - 12 - old_div(FONT_SIZE, 4)]
    LOCATION[1] = [67, CENTER - 218]
    LOCATION[2] = [70, CENTER + 182]
    IMAGE_FILE = os.path.join(FILE_DIR, 'trigonal_planar.png')
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        """
        self.sites[1].setShiftLeft()
        self.sites[2].setShiftLeft()  
[docs]class LinearSiteSelector(OctahedralSiteSelector):
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    MOVEXBY = 60
    MOVEYBY = 40
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [SIZE - 85, CENTER - 30]
    LOCATION[1] = [50, CENTER - 30]
    IMAGE_FILE = os.path.join(FILE_DIR, 'linear.png')
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        """
        self.sites[1].setShiftLeft()  
[docs]class OrderedOctahedralSiteSelector(OctahedralSiteSelector):
    """ Overrides the parent class to put coordination sites in a more rational
    order for tri and tetradentate ligands """
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [SIZE - 90, CENTER - 23]
    LOCATION[1] = [CENTER - 2 - old_div(FONT_SIZE, 4), 60]
    LOCATION[4] = [325, 375]
    LOCATION[3] = [CENTER - 2 - old_div(FONT_SIZE, 4), SIZE - 100]
    LOCATION[2] = [75, CENTER - 23]
    LOCATION[5] = [158, 80]
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        """
        self.sites[2].setShiftLeft()
        self.sites[5].setShiftLeft()  
[docs]class OrderedSquarePlanarSiteSelector(SquarePlanarSiteSelector):
    """ Overrides the parent class to put coordination sites in a more rational
    order for tri and tetradentate ligands """
    SIZE = MASTER_SIZE
    CENTER = old_div(SIZE, 2)
    LOCATION = {}
    """ The locations of the site labels in pixels. """
    LOCATION[0] = [SIZE - 90, CENTER - 23]
    LOCATION[1] = [CENTER - 2 - old_div(FONT_SIZE, 4), 60]
    LOCATION[3] = [CENTER - 2 - old_div(FONT_SIZE, 4), SIZE - 100]
    LOCATION[2] = [75, CENTER - 23]
[docs]    def markShiftedDesignators(self):
        """
        Mark the labels that should shift left when a ligand designator is added
        to it.  We shift some labels left so they don't clash with the geometry
        image.
        """
        self.sites[2].setShiftLeft()  
[docs]class Ligand(QtWidgets.QFrame):
    """
    A row of widgets for a single ligand
    """
[docs]    def __init__(self, row, index, layout=None):
        """
        Create a Ligand object
        :type row: LigandRow
        :param row: The main panel LigandRow object that provides the data for
            this ligand
        :type index: int
        :param index: The ligand index for this ligand (top ligand is 0, next
            ligand is 1, etc.)
        :type layout: QLayout
        :param layout: The layout to place this widget row in
        """
        QtWidgets.QFrame.__init__(self)
        self.row = row
        self.index = index
        mylayout = swidgets.SHBoxLayout(self)
        # Label for this ligand
        prefix = LETTERS[index] + ') '
        self.label = builderwidgets.StructureLabel(row,
                                                   mylayout,
                                                   "",
                                                   unset_text="",
                                                   prefix=prefix)
        self.label.set(row.structure)
        self.label.setMinimumWidth(120)
        # The radio buttons for the coordination sites for this ligand
        self.buttons = []
        self.buttons.append(SiteRadioButton(index, 'R1', layout=mylayout))
        if self.row.dentation_type == buildcomplex.BIDENTATE:
            self.buttons.append(SiteRadioButton(index, 'R2', layout=mylayout))
        mylayout.addStretch()
        if layout is not None:
            layout.addWidget(self) 
[docs]    def getSlots(self):
        """
        Return which coordination sites this ligand is connected to
        :rtype: list
        :return: List of int, each integer is a coordination site index. The
            first item is for R1, the second, if it exists is for R2.  A list item
            of None is returned if a coordination site has not been set for an R
            value
        """
        return [x.site_index for x in self.buttons]  
[docs]class LigandSelector(QtWidgets.QFrame):
    """
    The master frame that stores a row for each ligand
    """
[docs]    def __init__(self, ligand_rows, layout=None):
        """
        Create a LigandSelector object
        :type ligand_rows: list
        :param ligand_rows: LigandRow objects from the master panel for each
            ligand
        :type layout: QLayout
        :param layout: The layout to place this object into
        """
        QtWidgets.QFrame.__init__(self)
        mylayout = swidgets.SVBoxLayout(self)
        self.ligands = []
        count = 0
        # Create a row for each ligand
        for row in ligand_rows:
            for mate in range(row.numberOfCopies()):
                self.ligands.append(Ligand(row, count, mylayout))
                count = count + 1
        # Put all the radio buttons into a single group so that they are
        # exclusive.
        self.button_group = swidgets.SRadioButtonGroup()
        for ligand in self.ligands:
            for button in ligand.buttons:
                self.button_group.addExistingButton(button)
        try:
            self.button_group.button(0).setChecked(True)
        except AttributeError:
            pass
        # Add a label that is slightly longer than the row labels to the frame
        # so changing radiobutton labels doesn't cause expansion/contraction.
        blank_label = QtWidgets.QLabel("")
        blank_label.setMinimumWidth(300)
        mylayout.addWidget(blank_label)
        mylayout.addStretch()
        if layout is not None:
            layout.addWidget(self) 
[docs]    def getSlots(self):
        """
        Return the coordination sites picked for all the ligands. The sites will
        be returned in the same order the ligands appear in the LigandRow
        objects in the master.
        :rtype: list
        :return: list of int, each int is a coordination site index
        """
        slots = []
        for ligand in self.ligands:
            slots.extend(ligand.getSlots())
        return slots  
[docs]class SiteSelectionDialog(QtWidgets.QDialog):
    """
    The dialog window that allows the user to select coordination sites for each
    ligand
    """
    SELECTORS = {
        buildcomplex.OCTAHEDRAL: OctahedralSiteSelector,
        buildcomplex.TRIGONAL_BIPYRAMIDAL: TrigonalBipyramidalSiteSelector,
        buildcomplex.TETRAHEDRAL: TetrahedralSiteSelector,
        buildcomplex.SQUARE_PLANAR: SquarePlanarSiteSelector,
        buildcomplex.TRIGONAL_PLANAR: TrigonalPlanarSiteSelector,
        buildcomplex.LINEAR: LinearSiteSelector
    }
[docs]    def __init__(self, geometry, rows, parent=None, defined_slots=None):
        """
        Create a SiteSelectionDialog object
        :type geometry: str
        :param geometry: The VESPR geometry of the complex
        :type rows: list
        :param row: list of LigandRow objects to extract ligands from
        :type parent: QWidget
        :param parent: The window to display this dialog over
        :type defined_slots: list
        :param defined_slots: Each item in the list is the slot that
            coordinate site will use. The first item of the list is the slot for the
            R1 site of the first ligand in the first ligand row. The second item is
            the slot for the R2 site of the first ligand in the first ligand row (if
            that ligand is bidentate) or the R1 site of the next ligand, etc. A
            slot, in this case, is the metal atom coordination site.
        """
        self.master = parent
        QtWidgets.QDialog.__init__(self, parent)
        self.setWindowTitle('Specify Ligand Coordination Sites')
        layout = swidgets.SVBoxLayout(self)
        layout.setContentsMargins(6, 6, 6, 6)
        alayout = swidgets.SHBoxLayout()
        layout.addLayout(alayout)
        # The geometry/coordination site picture
        self.diagram = self.SELECTORS[geometry]()
        alayout.addWidget(self.diagram)
        # The rows of ligands
        self.ligand_frame = LigandSelector(rows, alayout)
        for site in self.diagram.sites:
            site.selected.connect(self.siteSelected)
            site.deselected.connect(self.siteDeselected)
            for button in self.ligand_frame.button_group.buttons():
                button.released.connect(site.released)
        if defined_slots:
            buttons = self.ligand_frame.getButtons()
            if len(defined_slots) != len(buttons):
                self.master.warning('Previously defined ligand sites cannot '
                                    'be mapped to current coordination sites, '
                                    'they will not be used')
                defined_slots = []
            for slot in defined_slots:
                if slot is None:
                    continue
                self.diagram.sites[slot].select(state=True)
        # Buttons
        dbb = QtWidgets.QDialogButtonBox
        dialog_buttons = dbb(dbb.Save | dbb.Cancel)
        dialog_buttons.accepted.connect(self.accept)
        dialog_buttons.rejected.connect(self.reject)
        layout.addWidget(dialog_buttons) 
[docs]    def siteDeselected(self, site):
        """
        A coordination site was deselected by clicking on it.  Let the radio
        buttons know so the one that was holding it will know it is no longer
        holding it.
        :type site: int
        :param site: The coordination site that was deselected.  This value will
            be one less than what the user sees.  (Site 1 emits site=0)
        """
        for button in self.ligand_frame.getButtons():
            button.siteUnpicked(site) 
[docs]    def siteSelected(self, site):
        """
        A coordination site was selected.  Let the current radiobutton know that
        it has a site, update the site label, and move to the next ligand site.
        :type site: int
        :param site: The coordination site that was selected.  This value will
            be one less than what the user sees.  (Site 1 emits site=0)
        """
        button, id = self.ligand_frame.getCurrentButton()
        if button is None:
            return
        button.sitePicked(site)
        designator = LETTERS[button.ligand_index] + str(button.text())[1]
        self.diagram.addLigandDesignator(site, designator)
        self.ligand_frame.nextButton(id=id) 
[docs]    def accept(self):
        """
        Save button callback.  Store the selected ligand sites
        """
        slots = self.ligand_frame.getSlots()
        if any([x is None for x in slots]):
            self.master.warning('Not all ligand sites have been defined')
            return False
        self.master.user_defined_slots = self.ligand_frame.getSlots()
        return QtWidgets.QDialog.accept(self)