# Source code for MDAnalysis.guesser.default_guesser

    # -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding: utf-8 -*-
    # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
    #
    # MDAnalysis --- https://www.mdanalysis.org
    # Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
    # (see the file AUTHORS for the full list of names)
    #
    # Released under the Lesser GNU Public Licence, v2 or any higher version
    #
    # Please cite your use of MDAnalysis in published work:
    #
    # R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
    # D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
    # MDAnalysis: A Python package for the rapid analysis of molecular dynamics
    # simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
    # Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
    # doi: 10.25080/majora-629e541a-00e
    #
    # N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
    # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
    # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787

    r"""
    Default Guesser
    ================
    .. _DefaultGuesser:

    DefaultGuesser is a generic guesser class that has basic guessing methods.
    This class is a general purpose guesser that can be used with most topologies,
    but being generic makes it the least accurate among all guessers.

    Guessing behavior
    -----------------

    This section describes how each attribute is guessed by the DefaultGuesser.

    Masses
    ~~~~~~

    We first attempt to look up the mass of an atom based on its element if the
    element TopologyAttr is available. If not, we attempt to lookup the mass based
    on the atom type (``type``) TopologyAttr. If neither of these is available, we
    attempt to guess the atom type based on the atom name (``name``) and then
    lookup the mass based on the guessed atom type.

    Types
    ~~~~~

    We attempt to guess the atom type based on the atom name (``name``).
    The name is first stripped of any numbers and symbols, and then looked up in
    the :data:`MDAnalysis.guesser.tables.atomelements` table. If the name is not
    found, we continue checking variations of the name following the logic in
    :meth:`DefaultGuesser.guess_atom_element`. Ultimately, if no match is found,
    the first character of the stripped name is returned.

    Elements
    ~~~~~~~~

    This follows the same method as guessing atom types.

    Bonds
    ~~~~~

    Bonds are guessed based on the distance between atoms.
    See :meth:`DefaultGuesser.guess_bonds` for more details.

    Angles
    ~~~~~~

    Angles are guessed based on the bonds between atoms.
    See :meth:`DefaultGuesser.guess_angles` for more details.

    Dihedrals
    ~~~~~~~~~

    Dihedrals are guessed based on the angles between atoms.
    See :meth:`DefaultGuesser.guess_dihedrals` for more details.

    Improper Dihedrals
    ~~~~~~~~~~~~~~~~~~

    Improper dihedrals are guessed based on the angles between atoms.
    See :meth:`DefaultGuesser.guess_improper_dihedrals` for more details.

    Aromaticities
    ~~~~~~~~~~~~~

    Aromaticity is guessed using RDKit's GetIsAromatic method.
    See :meth:`DefaultGuesser.guess_aromaticities` for more details.

    Classes
    -------

    .. autoclass:: DefaultGuesser
       :members:
       :inherited-members:

    """
    from .base import GuesserBase
    import numpy as np
    import warnings
    import math

    import re

    from ..exceptions import NoDataError
    from ..lib import distances
    from . import tables

    class DefaultGuesser(GuesserBase):
        """
        This guesser holds generic methods (not directed to specific contexts) for
        guessing different topology attribute. It has the same methods which where
        originally found in Topology.guesser.py. The attributes that can be
        guessed by this class are:
         - masses
         - types
         - elements
         - angles
         - dihedrals
         - bonds
         - improper dihedrals
         - aromaticities

        You can use this guesser either directly through an instance, or through
        the :meth:`~MDAnalysis.core.universe.Universe.guess_TopologyAttrs` method.

        Parameters
        ----------
        universe : Universe
            The Universe to apply the guesser on
        box : np.ndarray, optional
            The box of the Universe. This is used for bond guessing.
        vdwradii : dict, optional
            Dict relating atom types: vdw radii. This is used for bond guessing
        fudge_factor : float, optional
            The factor by which atoms must overlap each other to be considered
            a bond.  Larger values will increase the number of bonds found. [0.55]
        lower_bound : float, optional
            The minimum bond length. All bonds found shorter than this length
            will be ignored. This is useful for parsing PDB with altloc records
            where atoms with altloc A and B may be very close together and
            there should be no chemical bond between them. [0.1]

        Examples
        --------
        to guess bonds for a universe::

            import MDAnalysis as mda
            from MDAnalysisTests.datafiles import two_water_gro

            u = mda.Universe(two_water_gro, context='default', to_guess=['bonds'])

        .. versionadded:: 2.8.0

        """

        context = "default"

        def __init__(
            self,
            universe,
            box=None,
            vdwradii=None,
            fudge_factor=0.55,
            lower_bound=0.1,
            **kwargs,
        ):
            super().__init__(
                universe,
                box=box,
                vdwradii=vdwradii,
                fudge_factor=fudge_factor,
                lower_bound=lower_bound,
                **kwargs,
            )
            self._guesser_methods = {
                "masses": self.guess_masses,
                "types": self.guess_types,
                "elements": self.guess_types,
                "bonds": self.guess_bonds,
                "angles": self.guess_angles,
                "dihedrals": self.guess_dihedrals,
                "impropers": self.guess_improper_dihedrals,
                "aromaticities": self.guess_aromaticities,
            }

        def guess_masses(self, atom_types=None, indices_to_guess=None):
            """Guess the mass of many atoms based upon their type.
            For guessing masses through :meth:`~MDAnalysis.core.universe.Universe.guess_TopologyAttrs`:

            First try to guess masses from atom elements, if not available,
            try to guess masses from types and if not available, try to guess
            types.

            Parameters
            ----------
            atom_types : Optional[np.ndarray]
              Atom types/elements to guess masses from
            indices_to_guess : Optional[np.ndarray]
              Mask array for partially guess masses for certain atoms

            Returns
            -------
            atom_masses : np.ndarray dtype float64

            Raises
            ------
            :exc:`ValueError`
                If there are no atom types or elements to guess mass from.

            """
            if atom_types is None:
                try:
                    atom_types = self._universe.atoms.elements
                except NoDataError:
                    try:
                        atom_types = self._universe.atoms.types
                    except NoDataError:
                        try:
                            atom_types = self.guess_types(
                                atom_types=self._universe.atoms.names
                            )
                        except NoDataError:
                            raise NoDataError(
                                "there is no reference attributes"
                                " (elements, types, or names)"
                                " in this universe to guess mass from"
                            ) from None

            if indices_to_guess is not None:
                atom_types = atom_types[indices_to_guess]

            masses = np.array(
                [self.get_atom_mass(atom) for atom in atom_types], dtype=np.float64
            )
            return masses

        def get_atom_mass(self, element):
            """Return the atomic mass in u for *element*.
            Masses are looked up in :data:`MDAnalysis.guesser.tables.masses`.

            .. Warning:: Until version 3.0.0 unknown masses are set to 0.0

            """
            try:
                return tables.masses[element]
            except KeyError:
                try:
                    return tables.masses[element.upper()]
                except KeyError:
                    warnings.warn(
                        "Unknown masses are set to 0.0 for current version, "
                        "this will be deprecated in version 3.0.0 and replaced by"
                        " Masse's no_value_label (np.nan)",
                        PendingDeprecationWarning,
                    )
                    return 0.0

        def guess_atom_mass(self, atomname):
            """Guess a mass based on the atom name.

            :func:`guess_atom_element` is used to determine the kind of atom.

            .. warning:: Until version 3.0.0 anything not recognized is simply
               set to 0.0; if you rely on the masses you might want to double-check.
            """
            return self.get_atom_mass(self.guess_atom_element(atomname))

        def guess_types(self, atom_types=None, indices_to_guess=None):
            """Guess the atom type of many atoms based on atom name

            Parameters
            ----------
            atom_types (optional)
              atoms names if types guessing is desired to be from names
            indices_to_guess (optional)
              Mask array for partially guess types for certain atoms

            Returns
            -------
            atom_types : np.ndarray dtype object

            Raises
            ------
            :exc:`ValueError`
               If there is no names to guess types from.

            """
            if atom_types is None:
                try:
                    atom_types = self._universe.atoms.names
                except NoDataError:
                    raise NoDataError(
                        "there is no reference attributes in this universe "
                        "to guess types from"
                    ) from None

            if indices_to_guess is not None:
                atom_types = atom_types[indices_to_guess]

            return np.array(
                [self.guess_atom_element(atom) for atom in atom_types],
                dtype=object,
            )

        def guess_atom_element(self, atomname):
            """Guess the element of the atom from the name.

            First all numbers and symbols are stripped from the name.
            Then the name is looked up in the
            :data:`MDAnalysis.guesser.tables.atomelements` table.
            If the name is not found, we remove the last character or
            first character from the name and check the table for both,
            with a preference for removing the last character. If the name is
            still not found, we iteratively continue to remove the last character
            or first character until we find a match. If ultimately no match
            is found, the first character of the stripped name is returned.

            If the input name is an empty string, an empty string is returned.

            The table comes from CHARMM and AMBER atom
            types, where the first character is not sufficient to determine the
            atom type. Some GROMOS ions have also been added.

            .. Warning: The translation table is incomplete.
               This will probably result in some mistakes,
               but it still better than nothing!

            See Also
            --------
            :func:`guess_atom_type`
            :mod:`MDAnalysis.guesser.tables`
            """
            NUMBERS = re.compile(r"[0-9]")  # match numbers
            SYMBOLS = re.compile(r"[*+-]")  # match *, +, -
            if atomname == "":
                return ""
            try:
                return tables.atomelements[atomname.upper()]
            except KeyError:
                # strip symbols and numbers
                no_symbols = re.sub(SYMBOLS, "", atomname)
                name = re.sub(NUMBERS, "", no_symbols).upper()

                # just in case
                if name in tables.atomelements:
                    return tables.atomelements[name]

                while name:
                    if name in tables.elements:
                        return name
                    if name[:-1] in tables.elements:
                        return name[:-1]
                    if name[1:] in tables.elements:
                        return name[1:]
                    if len(name) <= 2:
                        return name[0]
                    name = name[:-1]  # probably element is on left not right

                # if it's numbers
                return no_symbols

        def guess_bonds(self, atoms=None, coords=None):
            r"""Guess if bonds exist between two atoms based on their distance.

            Bond between two atoms is created, if the two atoms are within

            .. math::

                  d < f \cdot (R_1 + R_2)

            of each other, where :math:`R_1` and :math:`R_2` are the VdW radii
            of the atoms and :math:`f` is an ad-hoc *fudge_factor*. This is
            the `same algorithm that VMD uses`_.

            Parameters
            ----------
            atoms: AtomGroup
                atoms for which bonds should be guessed
            coords: np.ndarray, optional
                coordinates of the atoms. If not provided, the coordinates
                of the ``atoms`` in the universe are used.

            Returns
            -------
            list
                List of tuples suitable for use in Universe topology building.

            Warnings
            --------
            No check is done after the bonds are guessed to see if Lewis
            structure is correct. This is wrong and will burn somebody.

            Raises
            ------
            :exc:`ValueError`
               If inputs are malformed or `vdwradii` data is missing.

            .. _`same algorithm that VMD uses`:
               http://www.ks.uiuc.edu/Research/vmd/vmd-1.9.1/ug/node26.html

            """
            if atoms is None:
                atoms = self._universe.atoms

            if coords is None:
                coords = self._universe.atoms.positions

            if len(atoms) != len(coords):
                raise ValueError("'atoms' and 'coord' must be the same length")

            fudge_factor = self._kwargs.get("fudge_factor", 0.55)

            # so I don't permanently change it
            vdwradii = tables.vdwradii.copy()
            user_vdwradii = self._kwargs.get("vdwradii", None)
            # this should make algo use their values over defaults
            if user_vdwradii:
                vdwradii.update(user_vdwradii)

            # Try using types, then elements
            if hasattr(atoms, "types"):
                atomtypes = atoms.types
            else:
                atomtypes = self.guess_types(atom_types=atoms.names)

            # check that all types have a defined vdw
            if not all(val in vdwradii for val in set(atomtypes)):
                raise ValueError(
                    (
                        "vdw radii for types: "
                        + ", ".join(
                            [t for t in set(atomtypes) if t not in vdwradii]
                        )
                        + ". These can be defined manually using the"
                        + f" keyword 'vdwradii'"
                    )
                )

            lower_bound = self._kwargs.get("lower_bound", 0.1)

            box = self._kwargs.get("box", None)

            if box is not None:
                box = np.asarray(box)

            # to speed up checking, calculate what the largest possible bond
            # atom that would warrant attention.
            # then use this to quickly mask distance results later
            max_vdw = max([vdwradii[t] for t in atomtypes])

            bonds = []

            pairs, dist = distances.self_capped_distance(
                coords, max_cutoff=2.0 * max_vdw, min_cutoff=lower_bound, box=box
            )
            for idx, (i, j) in enumerate(pairs):
                d = (
                    vdwradii[atomtypes[i]] + vdwradii[atomtypes[j]]
                ) * fudge_factor
                if dist[idx] < d:
                    bonds.append((atoms[i].index, atoms[j].index))
            return tuple(bonds)

        def guess_angles(self, bonds=None):
            """Given a list of Bonds, find all angles that exist between atoms.

            Works by assuming that if atoms 1 & 2 are bonded, and 2 & 3 are bonded,
            then (1,2,3) must be an angle.

            Parameters
            ----------
            bonds : Bonds
                 from which angles should be guessed

            Returns
            -------
            list of tuples
                List of tuples defining the angles.
                Suitable for use in u._topology

            See Also
            --------
            :meth:`guess_bonds`

            """
            from ..core.universe import Universe

            angles_found = set()

            if bonds is None:
                if hasattr(self._universe.atoms, "bonds"):
                    bonds = self._universe.atoms.bonds
                else:
                    temp_u = Universe.empty(n_atoms=len(self._universe.atoms))
                    temp_u.add_bonds(
                        self.guess_bonds(
                            self._universe.atoms, self._universe.atoms.positions
                        )
                    )
                    bonds = temp_u.atoms.bonds

            for b in bonds:
                for atom in b:
                    other_a = b.partner(atom)  # who's my friend currently in Bond
                    for other_b in atom.bonds:
                        if other_b != b:  # if not the same bond I start as
                            third_a = other_b.partner(atom)
                            desc = tuple(
                                [other_a.index, atom.index, third_a.index]
                            )
                            # first index always less than last
                            if desc[0] > desc[-1]:
                                desc = desc[::-1]
                            angles_found.add(desc)

            return tuple(angles_found)

        def guess_dihedrals(self, angles=None):
            """Given a list of Angles, find all dihedrals that exist between atoms.

            Works by assuming that if (1,2,3) is an angle, and 3 & 4 are bonded,
            then (1,2,3,4) must be a dihedral.

            Parameters
            ----------
            angles : Angles
                 from which dihedrals should be guessed

            Returns
            -------
            list of tuples
                List of tuples defining the dihedrals.
                Suitable for use in u._topology

            """
            from ..core.universe import Universe

            if angles is None:
                if hasattr(self._universe.atoms, "angles"):
                    angles = self._universe.atoms.angles

                else:
                    temp_u = Universe.empty(n_atoms=len(self._universe.atoms))

                    temp_u.add_bonds(
                        self.guess_bonds(
                            self._universe.atoms, self._universe.atoms.positions
                        )
                    )

                    temp_u.add_angles(self.guess_angles(temp_u.atoms.bonds))

                    angles = temp_u.atoms.angles

            dihedrals_found = set()

            for b in angles:
                a_tup = tuple([a.index for a in b])  # angle as tuple of numbers
                # if searching with b[0], want tuple of (b[2], b[1], b[0], +new)
                # search the first and last atom of each angle
                for atom, prefix in zip(
                    [b.atoms[0], b.atoms[-1]], [a_tup[::-1], a_tup]
                ):
                    for other_b in atom.bonds:
                        if not other_b.partner(atom) in b:
                            third_a = other_b.partner(atom)
                            desc = prefix + (third_a.index,)
                            if desc[0] > desc[-1]:
                                desc = desc[::-1]
                            dihedrals_found.add(desc)

            return tuple(dihedrals_found)

        def guess_improper_dihedrals(self, angles=None):
            """Given a list of Angles, find all improper dihedrals
            that exist between atoms.

            Works by assuming that if (1,2,3) is an angle, and 2 & 4 are bonded,
            then (2, 1, 3, 4) must be an improper dihedral.
            ie the improper dihedral is the angle between the planes formed by
            (1, 2, 3) and (1, 3, 4)

            Returns
            -------
                List of tuples defining the improper dihedrals.
                Suitable for use in u._topology

            """

            from ..core.universe import Universe

            if angles is None:
                if hasattr(self._universe.atoms, "angles"):
                    angles = self._universe.atoms.angles

                else:
                    temp_u = Universe.empty(n_atoms=len(self._universe.atoms))

                    temp_u.add_bonds(
                        self.guess_bonds(
                            self._universe.atoms, self._universe.atoms.positions
                        )
                    )

                    temp_u.add_angles(self.guess_angles(temp_u.atoms.bonds))

                    angles = temp_u.atoms.angles

            dihedrals_found = set()

            for b in angles:
                atom = b[1]  # select middle atom in angle
                # start of improper tuple
                a_tup = tuple([b[a].index for a in [1, 2, 0]])
                # if searching with b[1], want tuple of (b[1], b[2], b[0], +new)
                # search the first and last atom of each angle
                for other_b in atom.bonds:
                    other_atom = other_b.partner(atom)
                    # if this atom isn't in the angle I started with
                    if other_atom not in b:
                        desc = a_tup + (other_atom.index,)
                        if desc[0] > desc[-1]:
                            desc = desc[::-1]
                        dihedrals_found.add(desc)

            return tuple(dihedrals_found)

        def guess_atom_charge(self, atoms):
            """Guess atom charge from the name.

            .. Warning:: Not implemented; simply returns 0.
            """
            # TODO: do something slightly smarter, at least use name/element
            return 0.0

        def guess_aromaticities(self, atomgroup=None):
            """Guess aromaticity of atoms using RDKit

            Returns
            -------
            aromaticities : numpy.ndarray
                Array of boolean values for the aromaticity of each atom

            """
            if atomgroup is None:
                atomgroup = self._universe.atoms

            mol = atomgroup.convert_to("RDKIT")
            return np.array([atom.GetIsAromatic() for atom in mol.GetAtoms()])

        def guess_gasteiger_charges(self, atomgroup):
            """Guess Gasteiger partial charges using RDKit

            Parameters
            ----------
            atomgroup : mda.core.groups.AtomGroup
                Atoms for which the charges will be guessed

            Returns
            -------
            charges : numpy.ndarray
                Array of float values representing the charge of each atom

            """

            mol = atomgroup.convert_to("RDKIT")
            from rdkit.Chem.rdPartialCharges import ComputeGasteigerCharges

            ComputeGasteigerCharges(mol, throwOnParamFailure=True)
            return np.array(
                [
                    atom.GetDoubleProp("_GasteigerCharge")
                    for atom in mol.GetAtoms()
                ],
                dtype=np.float32,
            )