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

"""
Module containing classes for selection of theory in Jaguar GUIs.

Copyright Schrodinger, LLC. All rights reserved.
"""

import csv

from schrodinger.application.jaguar.gui.utils import THEORY_HF
from schrodinger.application.jaguar.gui.utils import THEORY_LMP2
from schrodinger.application.jaguar.gui.utils import THEORY_RIMP2
from schrodinger.application.jaguar.jaguar_keyword_utils import LEVELS_OF_THEORY
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import Qt
from schrodinger.ui.qt import filter_list
from schrodinger.ui.qt import pop_up_widgets
from schrodinger.utils import csv_unicode

NON_DFT_POSTFIX = " (non-DFT)"
DOCS_FIELDS = ['description', 'references']


class TheoryListWidgetItem(QtWidgets.QListWidgetItem):
    """
    Custom list widget items for theory selection.
    """

    def __init__(self,
                 text,
                 method=None,
                 is_recommended=False,
                 is_dispersion_corrected_dft=False,
                 is_long_range_corrected_dft=False,
                 is_hybrid_dft=False,
                 is_meta_gga_dft=False,
                 is_gga_dft=False,
                 is_lda_dft=False,
                 is_non_dft=False):
        """
        :param text: Display text for this list item
        :type text: str

        :param method: Actual method name for this list item.
        :type method: str

        :param is_recommended: Whether this list item is recommended
        :type is_recommended: bool

        :param is_dispersion_corrected_dft: Whether this list item is dispersion-corrected DFT
        :type is_dispersion_corrected_dft: bool

        :param is_long_range_corrected_dft: Whether this list item is long-range corrected DFT
        :type is_long_range_corrected_dft: bool

        :param is_hybrid_dft: Whether this list item is hybrid DFT
        :type is_hybrid_dft: bool

        :param is_meta_gga_dft: Whether this list item is meta-GGA DFT
        :type is_meta_gga_dft: bool

        :param is_gga_dft: Whether this list item is GGA DFT
        :type is_gga_dft: bool

        :param is_lda_dft: Whether this list item is LDA DFT
        :type is_lda_dft: bool

        :param is_non_dft: Whether this list item is non-DFT
        :type is_non_dft: bool
        """
        super().__init__(text)
        self.method = method
        self.is_recommended = is_recommended
        self.is_dispersion_corrected_dft = is_dispersion_corrected_dft
        self.is_long_range_corrected_dft = is_long_range_corrected_dft
        self.is_hybrid_dft = is_hybrid_dft
        self.is_meta_gga_dft = is_meta_gga_dft
        self.is_gga_dft = is_gga_dft
        self.is_lda_dft = is_lda_dft
        self.is_non_dft = is_non_dft


class TheorySelectorFilterListPopUp(filter_list.FilterListPopUp):
    """
    Class allowing for dynamic filtering and selection of methods.

    :cvar DFT_ONLY: Whether only DFT functionals should be included
    :type DFT_ONLY: bool
    """

    DFT_ONLY = False

    def __init__(self, parent):
        list_items = self._getListItems()
        cbs = self._getFilterCheckBoxes()
        super().__init__(parent, list_items, cbs,
                         'Limit list to matching methods')

    def _getListItems(self):
        """
        :return: A tuple of the list items for the list widget.
        :rtype: tuple(TheoryListWidgetItem)
        """
        list_items = []
        with csv_unicode.reader_open(LEVELS_OF_THEORY) as fh:
            reader = csv.DictReader(fh)
            for row in reader:
                for col in DOCS_FIELDS:
                    del row[col]
                for col, val in row.items():
                    if col != 'method':
                        if val not in ("0", "1"):
                            raise RuntimeError(
                                f"Unexpected {col} value for {row['method']}: {val}"
                            )
                        row[col] = val == '1'
                is_non_dft = row['is_non_dft']
                method = row['method']
                if is_non_dft and self.DFT_ONLY:
                    continue
                disp_name = method
                if is_non_dft:
                    disp_name = disp_name + NON_DFT_POSTFIX
                list_items.append(TheoryListWidgetItem(disp_name, **row))
        return tuple(list_items)

    def _getFilterCheckBoxes(self):
        """
        :return: a tuple of the filter checkboxes for this list
        :rtype: tuple(filter_list.FilterCheckBox)
        """
        recommended_func = lambda li: li.is_recommended
        recommended_cb = filter_list.FilterCheckBox("Recommended",
                                                    recommended_func)
        disp_corrected_func = lambda li: li.is_dispersion_corrected_dft
        disp_corrected_cb = filter_list.FilterCheckBox(
            'Dispersion corrected DFT', disp_corrected_func)
        long_range_func = lambda li: li.is_long_range_corrected_dft
        long_range_cb = filter_list.FilterCheckBox('Long range corrected DFT',
                                                   long_range_func)
        hybrid_func = lambda li: li.is_hybrid_dft
        hybrid_cb = filter_list.FilterCheckBox("Hybrid DFT", hybrid_func)
        meta_gga_func = lambda li: li.is_meta_gga_dft
        meta_gga_cb = filter_list.FilterCheckBox('Meta GGA DFT', meta_gga_func)
        gga_func = lambda li: li.is_gga_dft
        gga_cb = filter_list.FilterCheckBox('GGA DFT', gga_func)
        lda_func = lambda li: li.is_lda_dft
        lda_cb = filter_list.FilterCheckBox("LDA DFT", lda_func)
        non_dft_func = lambda li: li.is_non_dft
        non_dft_cb = filter_list.FilterCheckBox("Non-DFT", non_dft_func)
        cbs = [
            recommended_cb, disp_corrected_cb, long_range_cb, hybrid_cb,
            meta_gga_cb, gga_cb, lda_cb
        ]
        if not self.DFT_ONLY:
            cbs += [non_dft_cb]
        return tuple(cbs)

    def _getListItemTextForTheory(self, theory):
        """
        :return: The list item text that is shown for a specified theory method.
        :rtype: str
        """
        if theory in [THEORY_HF, THEORY_LMP2, THEORY_RIMP2]:
            theory_text = theory + NON_DFT_POSTFIX
        else:
            theory_text = theory
        return theory_text

    def setMethod(self, theory=None):
        """
        Set the theory to the specified value.

        :param theory: Theory value to set
        :type theory: str or None

        :return: True if theory was set, False otherwise.
        :rtype: bool
        """
        if theory is None:
            self._list_widget.setCurrentItem(None)
            return
        theory = self._getListItemTextForTheory(theory)
        items = self._list_widget.findItems(theory,
                                            Qt.MatchFlag.MatchFixedString)
        if len(items) != 1:
            return False
        item = items[0]
        if item.isHidden():
            return False
        self._list_widget.setCurrentItem(item)
        return True

    def getMethod(self):
        """
        :return: The currently selected theory value.
        :rtype: str or None
        """
        current_item = self._list_widget.currentItem()
        if not current_item:
            return None
        return current_item.method

    def addMethod(self, method, category, display_name=None):
        """
        Add the specified method to the popup's available methods.

        :param method: Method to be added
        :type method: str

        :param category: Category for this method
        :type category: str

        :param display_name: Display name for this method. If not
                             specified, the method value will be used.
        :type display_name: str or None
        """
        if display_name is None:
            display_name = method
        list_item = TheoryListWidgetItem(display_name, method, category)
        self._list_widget.addItem(list_item)

    def isItemHidden(self, theory):
        # See parent class for argument specification
        theory = self._getListItemTextForTheory(theory)
        return super().isItemHidden(theory)


class DftTheorySelectorFilterListPopUp(TheorySelectorFilterListPopUp):
    """
    Class allowing for dynamic filtering and selection of DFT functionals
    """

    DFT_ONLY = True


class TheorySelectorFilterListToolButton(
        filter_list.ToolButtonWithFilterListPopUp):
    """
    Custom tool button with a theory selector filter list pop up.
    """

    POP_UP_CLASS = TheorySelectorFilterListPopUp

    def getMethod(self):
        """
        :return: the currently selected theory
        :rtype: str or None
        """
        return self._pop_up.getMethod()

    def setMethod(self, theory=None):
        """
        Set the current theory value.

        :param theory: Theory value to be set
        :type theory: str or None

        :return: True if the theory level was set, False otherwise.
        :rtype: bool
        """
        return self._pop_up.setMethod(theory)

    def addMethod(self, method, category, display_name=None):
        """
        Add the specified method to the popup's available methods.

        :param method: Method to be added
        :type method: str

        :param category: Category for this method
        :type category: str

        :param display_name: Display name for this method. If not
                             specified, the method value will be used.
        :type display_name: str or None
        """
        self._pop_up.addMethod(method, category, display_name)

    def applySettings(self, settings):
        """
        Apply the specified filter settings to the pop up

        :param settings: Settings to be applied
        :type settings: dict
        """
        self._pop_up.applySettings(settings)


class DftTheorySelectorFilterListToolButton(TheorySelectorFilterListToolButton):
    """
    Custom tool button with a theory selector filter list pop up for
    DFT functionals.
    """

    POP_UP_CLASS = DftTheorySelectorFilterListPopUp


class FilterTheorySelectorReadOnlyLineEdit(pop_up_widgets.LineEditWithPopUp):
    """
    A read-only line edit used as an editor for table models with a TheorySelectorFilterPopUp.

    :ivar filtersChanged: Signal emitted when filters are toggled.
                          emits a dict of current filter settings.
    :type filtersChanged: QtCore.pyQtSignal(dict)
    """

    filtersChanged = QtCore.pyqtSignal(dict)

    def __init__(self, parent):
        super().__init__(parent, TheorySelectorFilterListPopUp)
        self.setReadOnly(True)
        self._pop_up.filtersChanged.connect(self.filtersChanged)
        # FIXME - The below is intended to assist in
        # closing the popup when embedded in a table, but
        # doesn't seem to be working as expected...
        self.setFocusPolicy(Qt.StrongFocus)

    def setMethod(self, theory=None):
        """
        Set the current theory value.

        :param theory: The current theory value to be set.
        :type theory: str or None
        """
        self._pop_up.setMethod(theory)

    def popUpUpdated(self):
        """
        Update the line edit text based on the current pop up selection.
        """
        theory = self._pop_up.getMethod()
        if theory is not None:
            self.setText(theory)

    def applySettings(self, settings):
        """
        Apply the specified filter settings to the pop up.

        :param settings: Filter settings to apply
        :type settings: dict
        """
        self._pop_up.applySettings(settings)