"""
Widget for letting the user show/hide or enable/disable features in a
hypothesis. All features are shown in a grid, and each has a checkbox next
to it.
"""
from collections import defaultdict
from past.utils import old_div
from schrodinger.application.phase import constants
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import Qt
from schrodinger.Qt.QtCore import pyqtSignal
from schrodinger.ui.qt import table_helper
NUM_COLUMNS = 4
[docs]class FeatureSelector(QtWidgets.QScrollArea):
    """
    This frame contains widgets for selecting one or more features in a
    hypothesis (e.g. A1, A2, H1, etc.).  The editor has a checkbox for each
    feature.  This frame is intended to be embedded in a layout.
    """
    # Gets emitted when the feature selection changes. Value is a set of
    # feature names:
    selectionChanged = pyqtSignal(set)
[docs]    def __init__(self, parent=None):
        super(FeatureSelector, self).__init__(parent)
        self.setWidgetResizable(True)
        self._interior_widget = QtWidgets.QWidget()
        self._interior_layout = QtWidgets.QVBoxLayout(self._interior_widget)
        self._features_layout = QtWidgets.QGridLayout()
        self._interior_layout.addLayout(self._features_layout)
        self._interior_layout.addStretch()
        self.setWidget(self._interior_widget)
        self._checkboxes = {}
        self.feature_labels = [] 
[docs]    def reset(self):
        """
        Unchecks all available checkboxes.
        """
        for checkbox in self._checkboxes.values():
            checkbox.setChecked(False) 
[docs]    def clear(self):
        """
        Remove all checkboxes from the features layout.
        """
        self.feature_labels = []
        while True:
            child = self._features_layout.takeAt(0)
            if not child:
                break
            child.widget().deleteLater()
        self._checkboxes = {} 
[docs]    def setFeatures(self, feature_names):
        """
        Initialize the editor with checkbox for each specified feature.
        """
        # Clear the dicts so that previous widgets can be deleted:
        self._checkboxes = {}
        # Clear previous contents of layout:
        self.clear()
        self.feature_labels = feature_names[:]
        # Split up the features into 4 columns:
        for i, feature_label in enumerate(self.feature_labels):
            row = old_div(i, NUM_COLUMNS)
            col = i % NUM_COLUMNS
            feature_type = feature_label[0]
            qcolor = constants.FEATURE_QCOLORS(feature_type)
            palette = QtGui.QPalette()
            palette.setColor(QtGui.QPalette.WindowText, qcolor)
            checkbox = QtWidgets.QCheckBox(feature_label)
            checkbox.setPalette(palette)
            checkbox.setStyleSheet('font-weight: bold;')
            checkbox.toggled.connect(self._selectionChanged)
            self._checkboxes[feature_label] = checkbox
            self._features_layout.addWidget(checkbox, row, col) 
[docs]    def getFeatureCheckbox(self, feature_name):
        """
        Return the checkbox for the feature.
        :param feature_name: The name of the feature (e.g. "A2")
        :type feature_name: str
        """
        return self._checkboxes[feature_name] 
    def _selectionChanged(self):
        """
        Emit the selectionChanged signal with a set of checked features.
        """
        selections = self.getSelectedFeatures()
        self.selectionChanged.emit(selections)
[docs]    def setSelectedFeatures(self, features):
        """
        :param features: The features to select/check.
        :type features: set of str
        Checks the checkboxes corresponding to the given set of feature,
        and unchecks the other checkboxes.
        """
        for feature_name in self.feature_labels:
            checkbox = self.getFeatureCheckbox(feature_name)
            checkbox.setChecked(feature_name in features)
        self._selectionChanged() 
[docs]    def getSelectedFeatures(self):
        """
        Return a set of selected features (their names).
        """
        selections = set()
        for feature_name in self.feature_labels:
            checkbox = self.getFeatureCheckbox(feature_name)
            if checkbox.isChecked():
                selections.add(feature_name)
        return selections  
[docs]class FeatureRow(object):
    """
    Data class that contains information about single feature in a hypothesis.
    This can be a regular feature or excluded volume.
    """
[docs]    def __init__(self, hypo_eid, hypo_name, feature_name, is_xvol, use_feature):
        """
        Initialize feature data.
        :param hypo_eid: hypothesis entry id
        :type hypo_eid: int
        :param hypo_name: hypothesis name. For custom features hypothesis
            name should be set to None!
        :type hypo_name: str
        :param is_xvol: True if this is excluded volume and False otherwise
        :type is_xvol: bool
        :param feature_name: feature name (empty string in case of excluded
            volume)
        :type feature_name: str
        :param use_feature: indicates whether checkbox for this feature
            is toggled on or off
        :type use_feature: bool
        """
        self.hypo_eid = hypo_eid
        self.hypo_name = hypo_name
        self.is_xvol = is_xvol
        self.feature_name = feature_name
        self.use_feature = use_feature  
class FeatureColumns(table_helper.TableColumns):
    HypoName = table_helper.Column("Hypothesis Name")
    FeatureName = table_helper.Column(
        "Feature", checkable=True, tooltip="Toggle on/off to use this feature.")
[docs]class FeatureModel(table_helper.RowBasedTableModel):
    """
    Features model.
    """
    Column = FeatureColumns
    ROW_CLASS = FeatureRow
    CUSTOM_HYPO_TEXT = "Custom"
    XVOL_TEXT = "XVols"
    # This signal is emitted when feature check box is toggled
    featureToggled = pyqtSignal()
    @table_helper.data_method(Qt.DisplayRole, Qt.CheckStateRole)
    def _getData(self, col, feature_row, role):
        # See base class for documentation
        if col == self.Column.HypoName and role == Qt.DisplayRole:
            if feature_row.hypo_name is None:
                return self.CUSTOM_HYPO_TEXT
            else:
                return feature_row.hypo_name
        if col == self.Column.FeatureName:
            if role == Qt.DisplayRole:
                if feature_row.is_xvol:
                    return self.XVOL_TEXT
                else:
                    return feature_row.feature_name
            if role == Qt.CheckStateRole:
                if feature_row.use_feature:
                    return Qt.Checked
                else:
                    return Qt.Unchecked
    @table_helper.data_method(Qt.ForegroundRole)
    def _getForegroundColor(self, col, feature_row, role):
        # See table_helper.data_method for method documentation
        if col == self.Column.FeatureName and not feature_row.is_xvol:
            feature_type = feature_row.feature_name[0]
            return constants.FEATURE_QCOLORS(feature_type)
        # Use default foreground color
        return None
    def _setData(self, col, feature_row, value, role, row_num):
        # See table_helper._setData for method documentation
        if role == Qt.CheckStateRole and col == self.Column.FeatureName:
            feature_row.use_feature = bool(value)
            self.featureToggled.emit()
            return True
        return False
[docs]    def getSelectedFeatures(self, include_pt_feats, include_custom_feats):
        """
        Returns dictionary of checked feature names. It is keyed on hypothesis
        entry ids and contains feature names for each hypothesis.
        :param include_pt_feats: indicates that selected features that came
            from PT hypotheses should be included
        :type include_pt_feats: bool
        :param include_custom_feats: indicates that selected features that were
            manually added by the user should be included
        :type include_custom_feats: bool
        :return: dictionary of checked feature names
        :rtype: dict
        """
        feature_selection = defaultdict(list)
        for row in self.rows:
            add_pt_row = (include_pt_feats and row.hypo_name)
            add_custom_row = (include_custom_feats and row.hypo_name is None)
            add_row = add_pt_row or add_custom_row
            if row.use_feature and add_row and not row.is_xvol:
                feature_selection[row.hypo_eid].append(row.feature_name)
        return feature_selection 
[docs]    def getSelectedExcludedVolumes(self):
        """
        Returns list of hypothesis entry ids which have excluded volumes checked.
        :return: list of hypothesis entry ids, which have have excluded volumes
            checked
        :rtype: list
        """
        xvol_selection = [
            row.hypo_eid for row in self.rows if row.use_feature and row.is_xvol
        ]
        return xvol_selection 
[docs]    def toggleSelection(self, hypo_eid, feature_name):
        """
        Flips use_feature flag for a given feature.
        :param hypo_eid: feature's hypothesis entry id
        :type hypo_eid: int
        :param feature_name: feature name
        :type feature_name: str
        """
        for row_num, row in enumerate(self.rows):
            if row.hypo_eid == hypo_eid and row.feature_name == feature_name:
                row.use_feature = not row.use_feature
                self.rowChanged(row_num)
                self.featureToggled.emit()
                break 
[docs]    def clearSelection(self):
        """
        Toggles of 'use' check boxes for all features.
        """
        for row in self.rows:
            row.use_feature = False
        self.columnChanged(self.Column.FeatureName) 
[docs]    def selectedFeaturesCount(self):
        """
        Returns total number of selected features (excluding volumes).
        """
        selection = [
            row.feature_name
            for row in self.rows
            if row.use_feature and not row.is_xvol
        ]
        return len(selection) 
[docs]    def getLastCustomFeatureNum(self):
        """
        Finds all 'custom' features and returns last feature number. For
        example, if custom features are ['A1', 'D11', 'R5'] last custom
        feature number will be 11.
        """
        custom_feature_nums = [
            int(row.feature_name[1:])
            for row in self.rows
            if row.hypo_name is None
        ]
        return max(custom_feature_nums) if custom_feature_nums else 0 
[docs]    def updateFeatureNames(self, marker_features):
        """
        This function updates feature names in the model so that they are
        consistent with markers feature names. This is needed when user
        changes feature type using edit feature dialog. In this case only
        a single feature row needs to be modified.
        :param marker_features: dictionary of feature marker names keyed
            on hypothesis entry ids.
        :type marer_features: dict
        """
        # find feature that changed its name
        original_names = self._getFeatureNames()
        for row_num, row in enumerate(self.rows):
            if row.is_xvol:
                continue
            hypo_eid = row.hypo_eid
            if row.feature_name not in marker_features[hypo_eid]:
                # Only one name in marker_features would be different from the
                # ones in original_names. Here we find this name and use it as
                # a 'new' feature name.
                new_name = list(
                    set(marker_features[hypo_eid]) -
                    set(original_names[hypo_eid]))[0]
                row.feature_name = new_name
                row.use_feature = False
                self.rowChanged(row_num)
                return 
    def _getFeatureNames(self):
        """
        Returns dictionary of all feature names in the model, which is keyed on
        hypothesis entry ids.
        :return: dictionary of feature names
        :rtype: dict
        """
        features = defaultdict(list)
        for row in self.rows:
            features[row.hypo_eid].append(row.feature_name)
        return features 
# For testing purposes only:
if __name__ == "__main__":
    app = QtWidgets.QApplication([__file__])
    selector = FeatureSelector()
    selector.setFeatures([
        'A1', 'A2', 'H1', 'H2', 'P1', 'D1', 'D2', 'D3', 'D4', 'R10', 'R11',
        'R12', 'R13'
    ])
    def _print(selection):
        print(selection)
    selector.selectionChanged.connect(_print)
    selector.show()
    selector.raise_()
    app.exec()