import enum
import schrodinger
from schrodinger.application.msv.gui import gui_models
from schrodinger.application.msv.gui import popups
from schrodinger.application.msv.gui import stylesheets
from schrodinger.application.msv.gui.viewconstants import \
    ALIGN_SELECTED_ONLY_CB_TT
from schrodinger.application.msv.gui.viewconstants import \
    MULTI_ALIGN_DISABLED_TT
from schrodinger.application.msv.gui.viewconstants import \
    PROF_ALIGN_DISABLED_TT
from schrodinger.application.msv.gui.viewconstants import AlignType
from schrodinger.application.msv.gui.viewconstants import MultAlnAlgorithm
from schrodinger.application.msv.gui.viewconstants import SeqAlnMode
from schrodinger.application.msv.gui.viewconstants import StructAlnMode
from schrodinger.application.msv.gui.viewconstants import StructAlnRefASLMode
from schrodinger.application.msv.gui.viewconstants import \
    StructAlnSequenceRepresents
from schrodinger.application.msv.gui.viewconstants import StructAlnTransform
from schrodinger.models import mappers
from schrodinger.protein import annotation
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtWidgets
from schrodinger.ui.qt import atomselector
from schrodinger.ui.qt import delegates
from schrodinger.ui.qt import mapperwidgets
from schrodinger.ui.qt import pop_up_widgets
from schrodinger.ui.qt import utils as qt_utils
from schrodinger.ui.qt import widgetmixins
from schrodinger.ui.qt.mapperwidgets import plptable
from . import alignment_pane_ui
from . import binding_site_aln_dialog_ui
from . import multiple_seq_aln_dialog_ui
from . import pairwise_seq_aln_dialog_ui
from . import protein_structure_aln_dialog_ui
maestro = schrodinger.get_maestro()
INFO_BTN_TOOLTIP = """This feature can only be used with the <i>Existing
entries</i> option, and only when the Reference sequence has structure.<br>The
Reference entry must have at least 2 chains that will be used in the
alignment. The entries to be aligned can have only 1 chain in the alignment.
"""
SELECTED_RES_TT = """
<p style='white-space:pre'>
This option pertains to sequence selection only.
Residue selection will automatically restrict the
sequence alignment to those residue blocks.
</p>"""
LOCK_GAPS_CB_TT = '''Preserve existing gaps in the Reference sequence (more gaps may still be
added).  When multiple pairwise alignments are requested simultaneously,
gaps will be locked automatically after the first pair is aligned.'''
PAIRWISE_SUBST_MATRICES = [
    "PAM250", "BLOSUM45", "BLOSUM62", "BLOSUM80", "GONNET", "Residue Identity"
]
SEQ_ANNO = annotation.ProteinSequenceAnnotations.ANNOTATION_TYPES
SEQ_MODE_TEXT = {
    SeqAlnMode.Multiple: "Multiple sequence alignment",
    SeqAlnMode.Pairwise: "Pairwise sequence alignment",
    SeqAlnMode.Structure: "Current structure superposition",
    SeqAlnMode.Residue: "Residue numbers",
    SeqAlnMode.Profile: "Profile alignment",
    SeqAlnMode.PairwiseSS: "Pairwise with secondary structure prediction",
}  # yapf: disable
STRUCT_MODE_TEXT = [(StructAlnMode.Superimpose, "Current sequence alignment"),
                    (StructAlnMode.Structure, "Protein structure alignment"),
                    (StructAlnMode.BindingSite, "Binding site alignment")]
OPTIONS_PAGE = enum.IntEnum('OptionsPage', [
    'multiple', 'pairwise', 'current_sequence', 'profile', 'prot_structure',
    'binding_site', 'sel_binding_site', 'current_structure', 'residue_numbers',
    'pairwise_ss'
],
                            start=0)
SETTINGS_PAGE = enum.IntEnum('SettingsPage', [
    'multiple', 'no', 'pairwise', 'prot_structure', 'binding_site',
    'pairwise_ss'
],
                             start=0)
COMBO_PAGE_MAPPING = {
    SeqAlnMode.Multiple: (OPTIONS_PAGE.multiple, SETTINGS_PAGE.multiple),
    SeqAlnMode.Pairwise: (OPTIONS_PAGE.pairwise, SETTINGS_PAGE.pairwise),
    SeqAlnMode.Structure: (OPTIONS_PAGE.current_structure, SETTINGS_PAGE.no),
    SeqAlnMode.Residue: (OPTIONS_PAGE.residue_numbers, SETTINGS_PAGE.no),
    SeqAlnMode.Profile: (OPTIONS_PAGE.profile, SETTINGS_PAGE.no),
    SeqAlnMode.PairwiseSS: (OPTIONS_PAGE.pairwise_ss, SETTINGS_PAGE.pairwise_ss),
    StructAlnMode.Superimpose: (OPTIONS_PAGE.current_sequence, SETTINGS_PAGE.no),
    StructAlnMode.Structure: (OPTIONS_PAGE.prot_structure, SETTINGS_PAGE.prot_structure),
    StructAlnMode.BindingSite: (OPTIONS_PAGE.binding_site, SETTINGS_PAGE.binding_site)
}  # yapf: disable
def _set_not_implemented(widget):
    """
    Mark a widget as not implemented
    """
    widget.setEnabled(False)
    font = widget.font()
    font.setStrikeOut(True)
    widget.setFont(font)
    widget.setToolTip("Not implemented")
[docs]class MsvComboBoxMixin:
[docs]    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.setStyleSheet(stylesheets.LIGHT_COMBOBOX)
        # MSV-1962: Workaround for a bug where the combobox background color
        # was being applied to the selected item background on Linux
        self.setItemDelegate(QtWidgets.QStyledItemDelegate(self)) 
    def _getWidth(self):
        """
        Return the correct width for the text.
        This method is needed to prevent Macs from truncating text that is too
        long
        """
        text_items = (self.itemText(i) for i in range(self.count()))
        font_metrics = self.fontMetrics()
        text_width = max((font_metrics.width(item) for item in text_items),
                         default=0)
        padding = 20
        return text_width + padding 
[docs]class MsvMappableComboBox(MsvComboBoxMixin, mapperwidgets.MappableComboBox):
    pass 
[docs]class MsvEnumComboBox(MsvComboBoxMixin, mapperwidgets.EnumComboBox):
    pass 
[docs]class AlignmentPane(mappers.MapperMixin, widgetmixins.InitMixin, BasePopUp):
    """
    Alignment pane for align options
    :ivar alignmentRequested: Signal emitted when "Align" is clicked.
    :vartype sequenceAlignmentRequested: QtCore.pyqtSignal
    """
    model_class = gui_models.PageModel
    ui_module = alignment_pane_ui
    alignmentRequested = QtCore.pyqtSignal()
[docs]    def initSetUp(self):
        super().initSetUp()
        self._previous_align_selected = True
        self.setObjectName('alignment_pane')
        self.setStyleSheet(stylesheets.ALIGNMENT_PANE)
        self._setUpComboBoxes()
        self._setUpProteinStructureAlnDialog()
        self._setUpBindingSiteAlnDialog()
        self._setUpMultipleSeqAlnDialog()
        self._setUpPairwiseSeqAlnDialog()
        ui = self.ui
        ui.dropdown_options_stacked_widget.setEnum(AlignType)
        self._mode_btn_grp = mapperwidgets.MappableButtonGroup({
            ui.sequences_rb: AlignType.Sequence,
            ui.structures_rb: AlignType.Structure,
        })
        ui.selected_res_info_btn.setToolTip(SELECTED_RES_TT)
        ui.clear_constraints_btn.pressed.connect(
            self._onClearConstraintsClicked)
        ui.ss_clear_constraints_btn.pressed.connect(
            self._onClearConstraintsClicked)
        ui.align_btn.clicked.connect(self._onAlignClicked)
        self.structure_dialog.alignBtnUpdateRequested.connect(
            self.updateAlignBtnVisibility)
        self.ui.lock_gaps_cb.setToolTip(LOCK_GAPS_CB_TT)
        self.ui.selected_only_cb.setToolTip(ALIGN_SELECTED_ONLY_CB_TT)
        self.ui.help_btn.clicked.connect(self._showHelp) 
    def _showHelp(self):
        qt_utils.help_dialog('MSV_ALIGN_PANE', parent=self)
[docs]    def updateAlignBtnVisibility(self, valid):
        """
        Enable/disable the Align button. Also, set the tooltip when disabled.
        :param valid: Whether to enable or disable the Align button.
        :type valid: bool
        """
        self.ui.align_btn.setEnabled(valid)
        tooltip = "" if valid else (
            "Chain mapping is invalid - correct values or toggle off option")
        self.ui.align_btn.setToolTip(tooltip) 
[docs]    def initLayOut(self):
        super().initLayOut()
        g = self.geometry()
        g.setWidth(0)  # Force the pane to resize to the minimum width
        self.setGeometry(g) 
[docs]    def defineMappings(self):
        M = self.model_class
        settings = M.options.align_settings
        page_target = mappers.TargetSpec(slot=self._updatePanelState)
        align_type_lhs = (self._mode_btn_grp,
                          self.ui.dropdown_options_stacked_widget, page_target)
        can_aln_tgt = mappers.TargetSpec(slot=self._updateCanAlign)
        can_aln_params = (M.menu_statuses.can_only_multiple_align,
                          M.menu_statuses.can_only_profile_align,
                          M.menu_statuses.can_aln_set_align)
        can_clear_constraints_tgt = mappers.TargetSpec(
            setter=self._updateClearConstraintsBtn)
        can_set_constraints_tgt = mappers.TargetSpec(
            setter=self._updateSetConstraintsBtn)
        return [
            (align_type_lhs, settings.align_type),
            ((self.sequence_combo, page_target), settings.seq_align_mode),
            ((self.structure_combo, page_target), settings.struct_align_mode),
            (self.ui.selected_only_cb, settings.align_only_selected_seqs),
            (self.pairwise_dialog, settings.pairwise),
            (self.multiple_dialog, settings.multiple),
            (self.structure_dialog, settings.protein_structure),
            (self.binding_site_aln_dialog, settings.binding_site),
            (self.ui.globally_conserved_residues_cb, settings.multiple.find_globally_conserved),
            (self.ui.multiple_superimpose_structures_cb, settings.multiple.superimpose_after),
            (self.ui.set_constraints_cb, settings.pairwise.set_constraints),
            (self.ui.ss_set_constraints_cb, settings.pairwise_ss.set_constraints),
            (self.ui.gpcr_cb,settings.pairwise_ss.use_gpcr_aln),
            (can_set_constraints_tgt, M.menu_statuses.can_set_constraints),
            (can_clear_constraints_tgt, M.menu_statuses.can_clear_constraints),
            (self.ui.lock_gaps_cb, settings.pairwise.lock_gaps),
            (self.ui.pairwise_superimpose_structures_cb, settings.pairwise.superimpose_after),
            (self.ui.resnum_superimpose_structures_cb, settings.residue_number.superimpose_after),
            (self.ui.align_sequences_cb, settings.protein_structure.align_seqs),
            (self.ui.force_alignment_cb, settings.protein_structure.force),
            (self.ui.superimpose_selected_cb, settings.superimpose.align_sel_res_only),
            (self.ui.align_following_superposition_cb, settings.binding_site.align_seqs),
            (self.ui.align_sel_res_sequences_cb, settings.binding_site.align_seqs),
            ((self.binding_site_cutoff_combo, page_target), settings.binding_site.binding_site_cutoff),
            (can_aln_tgt, can_aln_params),
        ]  # yapf: disable 
[docs]    def getSignalsAndSlots(self, model):
        # MSV-2583: Figure out why `res_selection_model` is getting
        # garbage collected
        self._res_sel_model_reference = model.aln.res_selection_model
        resSelectionChanged = model.aln.res_selection_model.selectionChanged
        ss = [
            (resSelectionChanged, self._updateSelectedResInfoIconVisibility),
            (model.menu_statuses.can_align_binding_siteChanged, self._updateCanAlign),
            (model.menu_statuses.can_align_selected_seqsChanged, self._updateSelectedSeqCb),
            (model.menu_statuses.can_only_multiple_alignChanged, self._updateSelectedSeqCb),
            (model.menu_statuses.can_only_profile_alignChanged, self._updateSelectedSeqCb),
            (model.options.align_settings.binding_site.align_sel_res_onlyChanged, self._updatePanelState),
        ]  # yapf: disable
        return ss 
[docs]    def setModel(self, model):
        super().setModel(model)
        if model is None:
            return
        self._updateSelectedResInfoIconVisibility() 
    def _setUpComboBoxes(self):
        self.sequence_combo = MsvMappableComboBox()
        for enum_member in SeqAlnMode:
            text = SEQ_MODE_TEXT[enum_member]
            self.sequence_combo.addItem(text, enum_member)
        self.ui.sequence_combo_layout.addWidget(self.sequence_combo)
        self.structure_combo = MsvEnumComboBox(enum=StructAlnMode)
        self.structure_combo.updateItemTexts(STRUCT_MODE_TEXT)
        self.ui.structure_combo_layout.addWidget(self.structure_combo)
        self.binding_site_cutoff_combo = popups.BindingSiteDistanceComboBox(
            parent=self)
        self.ui.binding_site_cutoff_layout.addWidget(
            self.binding_site_cutoff_combo)
        post_lbl = QtWidgets.QLabel(' of ligand')
        self.ui.binding_site_cutoff_layout.addWidget(post_lbl)
    def _setUpProteinStructureAlnDialog(self):
        """
        Add settings button to create protein structure aln dialog
        """
        protein_structure_aln_btn = popups.make_pop_up_tool_button(
            parent=self,
            pop_up_class=ProteinStructureAlnPopUp,
            obj_name='protein_structure_aln_btn',
        )
        protein_structure_aln_btn.setPopupValign(
            protein_structure_aln_btn.ALIGN_BOTTOM)
        self.ui.protein_structure_alignment_settings_widget.layout().addWidget(
            protein_structure_aln_btn)
        self.structure_dialog = protein_structure_aln_btn.popup_dialog
    def _setUpBindingSiteAlnDialog(self):
        """
        Add settings button that opens the binding site aln dialog
        """
        binding_site_aln_btn = popups.make_pop_up_tool_button(
            parent=self,
            pop_up_class=BindingSiteAlnPopup,
            obj_name='binding_site_aln_btn',
        )
        binding_site_aln_btn.setPopupValign(binding_site_aln_btn.ALIGN_BOTTOM)
        self.ui.binding_site_alignment_settings_widget.layout().addWidget(
            binding_site_aln_btn)
        self.binding_site_aln_dialog = binding_site_aln_btn.popup_dialog
    def _setUpMultipleSeqAlnDialog(self):
        """
        Add settings button to create multiple seq aln dialog
        """
        multiple_seq_aln_btn = popups.make_pop_up_tool_button(
            parent=self,
            pop_up_class=MultipleAlignSettingsPopUp,
            obj_name='multiple_seq_aln_btn')
        multiple_seq_aln_btn.setPopupValign(multiple_seq_aln_btn.ALIGN_BOTTOM)
        self.ui.Multiple_sequence_alignment_settings_widget.layout().addWidget(
            multiple_seq_aln_btn)
        self.multiple_dialog = multiple_seq_aln_btn.popup_dialog
    def _setUpPairwiseSeqAlnDialog(self):
        """
        Add settings button to create pairwise seq aln dialog
        """
        pairwise_seq_aln_btn = popups.make_pop_up_tool_button(
            parent=self,
            pop_up_class=PairwiseAlignSettingsPopUp,
            obj_name='pairwise_seq_aln_btn')
        pairwise_seq_aln_btn.setPopupValign(pairwise_seq_aln_btn.ALIGN_BOTTOM)
        self.ui.pairwise_sequence_alignment_settings_widget.layout().addWidget(
            pairwise_seq_aln_btn)
        self.pairwise_dialog = pairwise_seq_aln_btn.popup_dialog
    ##############################################
    # Mapper callbacks
    ##############################################
    @QtCore.pyqtSlot()
    def _updateSelectedSeqCb(self):
        """
        Manage state of "Selected seqs" checkbox
        """
        menu_statuses = self.model.menu_statuses
        align_restricted = (menu_statuses.can_only_multiple_align or
                            menu_statuses.can_only_profile_align)
        enable = menu_statuses.can_align_selected_seqs and not align_restricted
        self.ui.selected_only_cb.setEnabled(enable)
        settings = self.model.options.align_settings
        if enable:
            # If we're enabling the checkbox, restore the value to what it was
            # when we disabled it.
            settings.align_only_selected_seqs = self._previous_align_selected
        else:
            # If we're disabling the checkbox, save the state first so
            # we can restore it later.
            self._previous_align_selected = settings.align_only_selected_seqs
            settings.align_only_selected_seqs = align_restricted
    @QtCore.pyqtSlot()
    def _updateSelectedResInfoIconVisibility(self):
        """
        Manage the visibility of the info icon next to the "Selected Only"
        checkbox.
        The info icon will be shown for sequence alignment modes if any
        residues are selected.
        """
        should_show = (
            self.model.aln.res_selection_model.hasSelection() and
            self.model.options.align_settings.align_type is AlignType.Sequence)
        self.ui.selected_res_info_btn.setVisible(should_show)
    @QtCore.pyqtSlot()
    def _updatePanelState(self):
        """
        Change options and settings pages based on current mode
        """
        ui = self.ui
        settings = self.model.options.align_settings
        align_mode = settings.getAlignMode()
        options_page, settings_page = COMBO_PAGE_MAPPING[align_mode]
        if options_page == OPTIONS_PAGE.binding_site \
                
and settings.binding_site.align_sel_res_only:
            options_page = OPTIONS_PAGE.sel_binding_site
        ui.checkable_options_stacked_widget.setCurrentIndex(options_page)
        ui.settings_stacked_widget.setCurrentIndex(settings_page)
        self._updateCanAlign()
        self._updateSelectedResInfoIconVisibility()
    @QtCore.pyqtSlot()
    def _updateCanAlign(self):
        """
        Update whether aligning is allowed
        """
        enable_align = True
        settings = self.model.options.align_settings
        statuses = self.model.menu_statuses
        align_type = settings.align_type
        align_mode = settings.getAlignMode()
        align_btn_tooltip = ''
        if not statuses.can_aln_set_align:
            enable_align = False
            align_btn_tooltip = PROF_ALIGN_DISABLED_TT
        elif align_type is AlignType.Sequence:
            is_profile = align_mode is SeqAlnMode.Profile
            must_profile = statuses.can_only_profile_align
            if is_profile or must_profile:
                enable_align = is_profile and must_profile
                if not enable_align:
                    align_btn_tooltip = (PROF_ALIGN_DISABLED_TT if is_profile
                                         else MULTI_ALIGN_DISABLED_TT)
                elif not self.model.split_chain_view:
                    enable_align = False
                    align_btn_tooltip = "Profile alignment not enabled in combined chain mode"
            elif statuses.can_only_multiple_align:
                enable_align = align_mode is SeqAlnMode.Multiple
        elif align_type is AlignType.Structure:
            if (statuses.can_only_multiple_align or
                    statuses.can_only_profile_align):
                enable_align = False
            elif align_mode == StructAlnMode.Superimpose:
                enable_align = bool(maestro)
            elif align_mode == StructAlnMode.BindingSite and settings.binding_site.align_sel_res_only:
                enable_align = statuses.can_align_binding_site
        self.ui.align_btn.setEnabled(enable_align)
        self.ui.align_btn.setToolTip(align_btn_tooltip)
    @QtCore.pyqtSlot(bool)
    def _updateClearConstraintsBtn(self, can_clear_constraints):
        """
        Update the enabled state of the clear constraints btn
        """
        self.ui.clear_constraints_btn.setEnabled(can_clear_constraints)
        self.ui.ss_clear_constraints_btn.setEnabled(can_clear_constraints)
    @QtCore.pyqtSlot(bool)
    def _updateSetConstraintsBtn(self, can_set_constraints):
        """
        Update the enabled state of the set constraints btn
        """
        self.ui.set_constraints_cb.setEnabled(can_set_constraints)
        self.ui.ss_set_constraints_cb.setEnabled(can_set_constraints)
    ##############################################
    # Widget slots
    ##############################################
    def _onClearConstraintsClicked(self):
        self.model.aln.resetPairwiseConstraints()
    def _onAlignClicked(self):
        self.close()
        self.alignmentRequested.emit() 
[docs]class MappableAtomSelector(mappers.TargetMixin, atomselector.AtomSelector):
[docs]    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.aslTextModified.connect(self.targetValueChanged) 
[docs]    def targetGetValue(self):
        return self.getAsl() 
[docs]    def targetSetValue(self, value):
        self.setAsl(value)  
[docs]class ChainMapTableSpec(plptable.TableSpec):
    @plptable.FieldColumn(gui_models.ChainName.seq_chain,
                          title="Sequence (1 per entry)")
    def sequence_chain(self, field):
        return f"{field[0]}_{field[1]}"
    reference = plptable.FieldColumn(gui_models.ChainName.reference,
                                     title="Reference Chain",
                                     editable=True)
[docs]    @reference.data_method(role=delegates.ComboBoxDelegate.COMBOBOX_ROLE)
    def populateComboBox(self, field):
        return self._reference_chains 
[docs]    def setReferenceChains(self, ref_chains):
        """
        Set all the available chains from the reference structure ready to
        be added to the combobox delegate.
        :param ref_chains: list of chains.
        :type ref_chains: list(str)
        """
        self._reference_chains = ref_chains