Source code for schrodinger.application.qsite.input

"""
Module for reading and writing QSite input files.

See QSiteInput class for more documentation.

Copyright Schrodinger, LLC. All rights reserved.
"""
#Contributors: Mike Beachy

import os
import shutil
from collections.abc import MutableMapping

import schrodinger.infra.mm as mm
from schrodinger.infra.mmim import MMIMArgList
from schrodinger.infra.mmim import MMIMDict
from schrodinger.utils.fileutils import force_remove
from schrodinger.utils.fileutils import splitext

_qmregion_int_default = -9999
_qmregion_str_default = ""

# Module-level constants for the levels of theory that we support
QM = "qm"
NDDO = "nddo"


class Cut(object):
    """
    A class to handle representation and printing of general qm region cuts.

    To avoid having to update individual attributes in the back end,
    attributes are read-only. To modify attributes, make a new Cut instance
    with the desired values. The clone method is provided as a convenience
    method to ease creation of a Cut with modified values.

    """

    # Attributes and their defaults specified as a list of tuples to
    # maintain order in the output of _definedAttrs.
    _attrs = [
        ("molid", _qmregion_int_default),
        ("chain", _qmregion_str_default),
        ("resnum", _qmregion_int_default),
        ("inscode", _qmregion_str_default),
        ("qmatom", _qmregion_str_default),
        ("mmatom", _qmregion_str_default),
        ("theory", _qmregion_str_default),
    ]

    def __init__(self,
                 molid=None,
                 chain=None,
                 resnum=None,
                 inscode=None,
                 qmatom=None,
                 mmatom=None,
                 theory=QM):
        """
        Initializer for a qmregion cut specification.

        At least one of molid or qmatom to be defined.

        :param molid: A molecule id.
        :type  molid: int

        :param chain: A chain id.
        :type  chain: str

        :param resnum: A residue number.
        :type  resnum: int

        :param inscode: An insertion code.
        :type  inscode: str

        :param qmatom: The name of an atom to place on the quantum
                mechanical side of the cut.
        :type  qmatom: str

        :param mmatom: The name of an atom to place on the molecular
                mechanical side of the cut.
        :type  mmatom: str

        :param theory: The theoretical method to use for this cut
                ("nddo" or "qm" and case-insensitive)
        :type  theory: str

        """

        self._molid = molid
        self._chain = chain
        self._resnum = resnum
        self._inscode = inscode
        self._qmatom = qmatom
        self._mmatom = mmatom
        self._theory = theory

        if self.molid is None and self.qmatom is None:
            raise RuntimeError(
                "Invalid qmregion cut specification: need at least molid or qmatom defined."
            )

    @property
    def molid(self):
        """
        Read-only access to the molid attribute.

        """
        return self._molid

    @property
    def chain(self):
        """
        Read-only access to the chain attribute.

        """
        return self._chain

    @property
    def resnum(self):
        """
        Read-only access to the resnum attribute.

        """
        return self._resnum

    @property
    def inscode(self):
        """
        Read-only access to the resnum attribute.

        """
        return self._inscode

    @property
    def qmatom(self):
        """
        Read-only access to the qmatom attribute.

        """
        return self._qmatom

    @property
    def mmatom(self):
        """
        Read-only access to the mmatom attribute.

        """
        return self._mmatom

    @property
    def theory(self):
        """
        Read-only access to the theory attribute.

        """
        return self._theory.lower()

    def clone(self, **kwargs):
        """
        Create a copy of the current instance, modifying any attributes
        provided as keywords.

        """
        attrs = {}
        for attr in self._definedAttrs():
            attrs[attr] = getattr(self, "_" + attr)
        attrs.update(kwargs)
        return Cut(**attrs)

    def _definedAttrs(self):
        """
        Return the defined attributes for this cut.

        """
        defined = []
        for attr, default in self._attrs:
            value = getattr(self, "_" + attr)
            if value is not None and value != default:
                defined.append(attr)
        return defined

    def header(self):
        """
        Return the header definition for this cut that would be used in the
        QSite input file's &qmregion section.

        """
        attrs = self._definedAttrs()
        return " ".join("%6s" % attr for attr in attrs if attr != "inscode")

    def __str__(self):
        attrs = self._definedAttrs()
        strings = []
        for attr in attrs:
            if attr == "resnum":
                inscode = str(self._inscode)
                strings.append("%*s%s" %
                               (6 - len(inscode), self._resnum, inscode))
            elif attr == "inscode":
                continue
            else:
                strings.append("%6s" % str(getattr(self, attr)))

        return " ".join(strings)

    def __repr__(self):
        return "Cut(%s)" % ", ".join("%s=%s" % (attr, repr(getattr(self, attr)))
                                     for attr in self._definedAttrs())


class HydrogenCap(object):
    """
    A class to handle representation and printing of hydrogen caps for the
    qm region.

    """

    def __init__(self, qm=None, mm=None, distance=0.0, theory=QM):
        """
        Initialize the hydrogen cap. A quantum mechanical atom and molecular
        mechanical atom must always be specified.

        :param qm: The name of the quantum mechanical atom.
        :type  qm: str

        :param mm: The name of the molecular mechanical atom.
        :type  mm: str

        :param distance: The hydrogen cap distance.
        :type  distance: float

        :param theory: The theoretical model to use for the QM atom
                (Either "nddo" or "qm", case-insensitive)
        :type  theory: str

        """
        if qm is None or mm is None:
            raise RuntimeError("Invalid qmregion hcap specification: "
                               "qm and mm must both be specified.")

        self._qm = str(qm)
        self._mm = str(mm)
        self._distance = distance
        self._theory = str(theory).lower()

    @property
    def qm(self):
        """
        Read-only access to the _qm attribute.

        """
        return self._qm

    @property
    def mm(self):
        """
        Read-only access to the _mm attribute.

        """
        return self._mm

    @property
    def theory(self):
        """
        Read-only access to the _theory attribute.

        """
        return self._theory

    @property
    def distance(self):
        """
        Read-only access to the distance attribute.

        """
        return self._distance

    def clone(self, **kwargs):
        """
        Create a copy of the current instance, modifying any attributes
        provided as keywords.

        """
        attrs = {}
        attrs["qm"] = self._qm
        attrs["mm"] = self._mm
        attrs["theory"] = self._theory
        attrs["distance"] = self._distance

        attrs.update(kwargs)
        return HydrogenCap(**attrs)

    def header(self):
        """
        Return the header definition for this hydrogen cap that would be
        used in the QSite input file's &qmregion section.

        """
        if self.distance == 0.0:
            if self.theory == NDDO:
                return "hcapnddo hcapmm"
            else:
                return "hcapqm hcapmm"
        else:
            if self.theory == NDDO:
                return "hcapnddo hcapmm hcapdist"
            else:
                return "hcapqm hcapmm hcapdist"

    def __str__(self):
        if self.distance == 0.0:
            if self.theory == NDDO:
                return "%8s %6s" % (self.qm, self.mm)
            else:
                return "%6s %6s" % (self.qm, self.mm)
        else:
            if self.theory == NDDO:
                return "%8s %6s %8.6f" % (self.qm, self.mm, self.distance)
            else:
                return "%6s %6s %8.6f" % (self.qm, self.mm, self.distance)

    def __repr__(self):
        if self.distance:
            dstring = ", distance=%s" % repr(self.distance)
        else:
            dstring = ""

        if self.theory == NDDO:
            return "HydrogenCap(nddo=%s, mm=%s%s, theory=%s)" % (repr(
                self.qm), repr(self.mm), dstring, repr(self.theory))
        else:
            return "HydrogenCap(qm=%s, mm=%s%s)" % (repr(self.qm), repr(
                self.mm), dstring)


class _GenDict(MutableMapping):
    """
    A dictionary-like class for getting and setting keywords in the &gen
    section of qsite input, which controls the Jaguar portion of the
    calculation.

    Note the 'default' method, which is non-standard for dictionaries. It
    retrieves the default value for a keyword.

    """

    def __init__(self, handle):
        self.handle = handle

    def default(self, key):
        """
        Return the default value for the provided keyword.
        """
        type_ = mm.mmjag_key_type(self.handle, key)
        if type_ == mm.MMJAG_INT:
            return mm.mmjag_key_int_def(self.handle, key)
        elif type_ == mm.MMJAG_REAL:
            return mm.mmjag_key_real_def(self.handle, key)
        else:
            return mm.mmjag_key_char_def(self.handle, key)

    def __getitem__(self, key):
        """
        Get a value for a &gen section keyword.
        """
        type_ = mm.mmjag_key_type(self.handle, key)
        if type_ == mm.MMJAG_INT:
            return mm.mmjag_key_int_get(self.handle, key)
        elif type_ == mm.MMJAG_REAL:
            return mm.mmjag_key_real_get(self.handle, key)
        else:
            return mm.mmjag_key_char_get(self.handle, key)

    def __setitem__(self, key, value):
        """
        Set a value for the &gen section.

        """
        if isinstance(value, int):
            mm.mmjag_key_int_set(self.handle, key, value)
        elif isinstance(value, float):
            mm.mmjag_key_real_set(self.handle, key, value)
        else:
            mm.mmjag_key_char_set(self.handle, key, value)

        return

    def __delitem__(self, key):
        """
        Delete a keyword from the &gen section.

        """
        mm.mmjag_key_delete(self.handle, key)
        return

    def keys(self):
        """
        Return the keys for the &gen section. Note that only keys with
        non-default values will be returned.

        """
        genkeys = []
        for i in range(1, mm.mmjag_key_nondef_count(self.handle) + 1):
            genkeys.append(mm.mmjag_key_nondef_get(self.handle, i))
        return genkeys

    def __len__(self):
        return len(self.keys())

    def __iter__(self):
        return iter(self.keys())


class _MopacDict(MutableMapping):
    """
    A dictionary-like class for getting and setting keywords in the &mopac
    section of qsite input, which controls the Jaguar/MOPAC portion of the
    calculation.

    Note the 'default' method, which is non-standard for dictionaries. It
    retrieves the default value for a keyword.

    """

    def __init__(self, handle):
        self.handle = handle

    def default(self, key):
        """
        Return the default value for the provided keyword.
        """
        type_ = mm.mmjag_mopac_key_type(self.handle, key)
        if type_ == mm.MMJAG_INT:
            return mm.mmjag_mopac_key_int_def(self.handle, key)
        elif type_ == mm.MMJAG_REAL:
            return mm.mmjag_mopac_key_real_def(self.handle, key)
        else:
            return mm.mmjag_mopac_key_char_def(self.handle, key)

    def __getitem__(self, key):
        """
        Get a value for a &mopac section keyword.
        """
        type_ = mm.mmjag_mopac_key_type(self.handle, key)
        if type_ == mm.MMJAG_INT:
            return mm.mmjag_mopac_key_int_get(self.handle, key)
        elif type_ == mm.MMJAG_REAL:
            return mm.mmjag_mopac_key_real_get(self.handle, key)
        else:
            return mm.mmjag_mopac_key_char_get(self.handle, key)

    def __setitem__(self, key, value):
        """
        Set a value for the &mopac section.

        """
        if isinstance(value, int):
            mm.mmjag_mopac_key_int_set(self.handle, key, value)
        elif isinstance(value, float):
            mm.mmjag_mopac_key_real_set(self.handle, key, value)
        else:
            mm.mmjag_mopac_key_char_set(self.handle, key, value, False)

        return

    def __delitem__(self, key):
        """
        Delete a keyword from the &mopac section.

        """
        mm.mmjag_mopac_key_delete(self.handle, key)
        return

    def keys(self):
        """
        Return the keys for the &mopac section. Note that only keys with
        non-default values will be returned.

        """
        mopkeys = []
        for i in range(1, mm.mmjag_mopac_key_nondef_count(self.handle) + 1):
            mopkeys.append(mm.mmjag_mopac_key_nondef_get(self.handle, i))
        return mopkeys

    def __len__(self):
        return len(self.keys())

    def __iter__(self):
        return iter(self.keys())


class QMRegion(object):
    """
    A class to provide access to a &qmregion specification and the ability
    to create one on the fly.

    """

    _list_defaults = {
        "molid": _qmregion_int_default,
        "chain": _qmregion_str_default,
        "resnum": _qmregion_int_default,
        "inscode": _qmregion_str_default,
        "qmatom": _qmregion_str_default,
        "seatom": _qmregion_str_default,
        "mmatom": _qmregion_str_default,
        "hcap": 0,
        "hcap_distance": 0.0,
        "theory": QM,
    }

    def __init__(self, handle):
        """
        :param handle: An MMIM handle.
        :param type: int

        """
        self.handle = handle

        self.molid = MMIMArgList(handle, mm.MMIM_QSITE_QMMOL)
        self.chain = MMIMArgList(handle, mm.MMIM_QSITE_QMCHN)
        self.resnum = MMIMArgList(handle, mm.MMIM_QSITE_QMRES)
        self.inscode = MMIMArgList(handle, mm.MMIM_QSITE_QMINST)
        self.qmatom = MMIMArgList(handle, mm.MMIM_QSITE_QMATOM)
        self.seatom = MMIMArgList(handle, mm.MMIM_QSITE_SEATOM)
        self.mmatom = MMIMArgList(handle, mm.MMIM_QSITE_MMATOM)
        self.hcap = MMIMArgList(handle, mm.MMIM_QSITE_HCAP)
        self.hcap_distance = MMIMArgList(handle, mm.MMIM_QSITE_HCAPDIST)
        self.theory = MMIMArgList(handle, mm.MMIM_QSITE_THEORY)

    def __len__(self):
        return mm.mmim_handle_get_indexed_arg_int_count(self.handle,
                                                        mm.MMIM_QSITE_QMMOL)

    def __getitem__(self, index):
        if self.hcap[index]:
            return HydrogenCap(
                qm=self.qmatom[index],
                mm=self.mmatom[index],
                distance=self.hcap_distance[index],
                theory=self.theory[index],
            )
        else:
            return Cut(
                self.molid[index],
                self.chain[index],
                self.resnum[index],
                self.inscode[index],
                self.qmatom[index],
                self.mmatom[index],
                self.theory[index],
            )

    def __setitem__(self, index, qmregion):
        if isinstance(qmregion, Cut):
            for attr in qmregion._definedAttrs():
                arglist = getattr(self, attr)
                arglist[index] = getattr(qmregion, attr)
        elif isinstance(qmregion, HydrogenCap):
            index = self.__len__() - 1
            self.hcap[index] = 1
            self.theory[index] = qmregion.theory
            if self.theory[index] == NDDO:
                self.seatom[index] = qmregion.qm
            else:
                self.qmatom[index] = qmregion.qm
            self.mmatom[index] = qmregion.mm
            self.hcap_distance[index] = qmregion.distance
        else:
            raise Exception("QMRegion elements can only be set to instances "
                            "of Cut or HydrogenCap.")

    def _growLists(self):
        """
        Extend all MMIMArgLists by one element with their default values.

        """

        for list_attr, default in QMRegion._list_defaults.items():
            getattr(self, list_attr).append(default)

    def modify(self, index, **kwargs):
        """
        A method to modify attributes of a specific element in the qmregion
        list in-place.

        This method is provided because Cut and HydrogenCap instances have
        read-only values.

        """
        for k, v in kwargs.items():
            arglist = getattr(self, k)
            if v is None:
                v = QMRegion._list_defaults[k]
            arglist[index] = v

    def append(self, qmregion):
        """
        Add a new QM region to the list.

        :param qmregion: A specification of a QM region cut or hydrogen cap.
        :type  qmregion: Cut or HydrogenCap

        """
        if isinstance(qmregion, Cut) or isinstance(qmregion, HydrogenCap):
            self._growLists()
            index = self.__len__() - 1
            self.__setitem__(index, qmregion)
        else:
            raise RuntimeError("The append method only takes instances of "
                               "Cut or HydrogenCap.")


class QSiteInput(object):

    def __init__(self, name):
        """
        Create a QSite input object, either from a jobname or QSite input
        file path.

        Parameters

        name (str)
            One of:
              jobname -   Default parameters are used.
            or
              file path - Parameters are read from specified file.
                          File must be in current working directory.

        """
        mm.mmerr_initialize()
        mm.mmjag_initialize(mm.error_handler)
        mm.mmim_initialize(mm.error_handler)

        self._gen = None
        self._mmkey = None
        self._mopac = None
        self._qmregion = None
        self._structure_file = None

        if name.endswith('.in'):
            if not os.path.exists(name):
                raise IOError("File does not exist: %s" % name)

            # Make sure that the input file is in CWD:
            if not os.path.exists(os.path.basename(name)):
                raise ValueError(
                    "File must be in current working directory: %s" % name)

            self.j_handle, self.i_handle, self.ct1, self.ct2, self.ct3 = \
                    mm.mmjag_read_qsite_file(name)
        else:  # jobname specified
            self.j_handle = mm.mmjag_new()
            self.i_handle = mm.mmim_new()
            self.mmkey[mm.MMIM_JOB_NAME] = name
        self.gen["mmqm"] = 1

    def __del__(self):
        """
        Delete the mmjag handle and clean up the mmlibs.

        """
        try:
            mm.mmjag_delete(self.j_handle)
        except:
            pass
        mm.mmim_terminate()
        mm.mmjag_terminate()
        mm.mmerr_terminate()

    # &gen section

    @property
    def gen(self):
        """
        Provide dictionary-like access to &gen section

        """
        if self._gen is None:
            self._gen = _GenDict(self.j_handle)
        return self._gen

    # &mmkey section

    @property
    def mmkey(self):
        """
        Provide dictionary-like access to the &mmkey section.

        """
        if self._mmkey is None:
            self._mmkey = MMIMDict(self.i_handle)
        return self._mmkey

    # &mopac section

    @property
    def mopac(self):
        """
        Provide dictionary-like access to &mopac section

        """
        if self._mopac is None:
            self._mopac = _MopacDict(self.j_handle)
        return self._mopac

    # &qmregion section

    @property
    def qmregion(self):
        """
        Provide access to the &qmregion section.

        """
        if self._qmregion is None:
            self._qmregion = QMRegion(self.i_handle)
        return self._qmregion

    def setStructureFile(self, filename):
        """
        Set the mae structure file to be used. If the structure file name
        doesn't match the current job/file name it will be copied to the
        current directory and renamed to <jobname>.<ext>.

        """

        # Local path to <filename>:
        base = os.path.basename(filename)
        prefix, ext = splitext(filename)

        if prefix != self.jobname or not os.path.exists(base):
            new_filename = self.jobname + ext
            force_remove(new_filename)
            shutil.copy(filename, new_filename)
            filename = new_filename

        mm.mmjag_directive_set(self.j_handle, "MAEFILE", filename)

    def getStructureFile(self):
        """
        Return the mae structure that is used.

        """
        return mm.mmjag_directive_get(self.j_handle, "MAEFILE")

    def addQMMolecule(self, index, theory=QM):
        """
        Add a QM molecule to the QM region specification.

        This creates a Cut instance and adds it to the qmregion attribute.

        """
        cut = Cut(molid=index, theory=theory)
        self.qmregion.append(cut)

    def addQMCut(self, **kwargs):
        """
        Add a QM cut to the QM region specification.

        The keywords here are the same as the Cut class constructor. Calling
        this method is equivalent to creating a Cut instance and adding it
        to the QSiteInput instance by calling qmregion.append().

        """
        cut = Cut(**kwargs)
        self.qmregion.append(cut)

    def addHydrogenCap(self, qm=None, mm=None, distance=0.0, theory=QM):
        """
        Add a hydrogen cap to the QM region specification.

        The keywords here are the same as the HydrogenCap class constructor.
        Calling this method is equivalent to creating a HydrogenCap instance
        and adding it to the QSiteInput instance by calling
        qmregion.append().

        """
        cap = HydrogenCap(qm, mm, distance, theory=theory)
        self.qmregion.append(cap)

    # Job and file names

    def _setJobName(self, name):
        """
        Set the jobname and filename properties from an input filename or
        jobname.

        This method keeps the filename and jobname in sync.

        """
        self.mmkey[mm.MMIM_JOB_NAME] = name

    def _getJobName(self):
        """
        Return the name of the job.

        """
        return self.mmkey[mm.MMIM_JOB_NAME]

    jobname = property(_getJobName,
                       _setJobName,
                       doc="get/set the jobname; "
                       "setting the jobname also sets the filename")

    # Methods

    def write(self):
        """
        Write the input file to <jobname>.in.
        The structure file will also be written to <jobname>.<ext>, where
        <ext> is the extension of the specified structure file.

        Structure file must be specified via setStructureFile() first.

        Previous file with this name is overwritten.

        """
        maefile = self.getStructureFile()

        if not maefile:
            raise RuntimeError(
                "You must specify a .mae structure file with "
                "the setStructureFile() method before writing an input "
                "file.")

        if not os.path.exists(maefile):
            raise IOError("Structure file does not exist: %s" % maefile)

        # Make sure that the name and location of the maefile is in sync
        # with the current jobname.
        self.setStructureFile(maefile)

        filename = self.jobname + '.in'

        mm.mmjag_write_qsite_file(self.j_handle, self.i_handle, filename)


#EOF