# Source code for MDAnalysis.lib.NeighborSearch

    # -*- 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.1 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
    #

    """
    Neighbor Search wrapper for MDAnalysis --- :mod:`MDAnalysis.lib.NeighborSearch`
    ===============================================================================

    This module contains classes that allow neighbor searches directly with
    `AtomGroup` objects from `MDAnalysis`.
    """
    import numpy as np
    from MDAnalysis.lib.distances import capped_distance
    from MDAnalysis.lib.util import unique_int_1d
    from MDAnalysis.core.groups import AtomGroup, SegmentGroup, ResidueGroup
    import numpy.typing as npt
    from typing import Optional, Union, List

    class AtomNeighborSearch(object):
        """This class can be used to find all atoms/residues/segments within the
        radius of a given query position.

        For the neighbor search, this class is a wrapper around
        :class:`~MDAnalysis.lib.distances.capped_distance`.
        """

        def __init__(
            self, atom_group: AtomGroup, box: Optional[npt.ArrayLike] = None
        ) -> None:
            """

            Parameters
            ----------
            atom_list : AtomGroup
              list of atoms
            box : array-like or ``None``, optional, default ``None``
              Simulation cell dimensions in the form of
              :attr:`MDAnalysis.trajectory.timestep.Timestep.dimensions` when
              periodic boundary conditions should be taken into account for
              the calculation of contacts.
            """
            self.atom_group = atom_group
            self._u = atom_group.universe
            self._box = box

        def search(
            self, atoms: AtomGroup, radius: float, level: str = "A"
        ) -> Optional[Union[AtomGroup, ResidueGroup, SegmentGroup]]:
            """
            Return all atoms/residues/segments that are within *radius* of the
            atoms in *atoms*.

            Parameters
            ----------
            atoms : AtomGroup, MDAnalysis.core.groups.AtomGroup
              AtomGroup object
            radius : float
              Radius for search in Angstrom.
            level : str
              char (A, R, S). Return atoms(A), residues(R) or segments(S) within
              *radius* of *atoms*.

            Returns
            -------
            AtomGroup : :class:`~MDAnalysis.core.groups.AtomGroup`
              When ``level='A'``, AtomGroup is being returned.
            ResidueGroup : :class:`~MDAnalysis.core.groups.ResidueGroup`
              When ``level='R'``, ResidueGroup is being returned.
            SegmentGroup : :class:`~MDAnalysis.core.groups.SegmentGroup`
              When ``level='S'``, SegmentGroup is being returned.

            .. versionchanged:: 2.0.0
               Now returns :class:`AtomGroup` (when empty this is now an empty
               :class:`AtomGroup` instead of an empty list), :class:`ResidueGroup`,
               or a :class:`SegmentGroup`
            """
            unique_idx = []
            try:
                # For atom groups, take the positions attribute
                position = atoms.positions
            except AttributeError:
                # For atom, take the position attribute
                position = atoms.position
            pairs = capped_distance(
                position,
                self.atom_group.positions,
                radius,
                box=self._box,
                return_distances=False,
            )

            if pairs.size > 0:
                unique_idx = unique_int_1d(np.asarray(pairs[:, 1], dtype=np.intp))
            return self._index2level(unique_idx, level)

        def _index2level(
            self, indices: List[int], level: str
        ) -> Union[AtomGroup, ResidueGroup, SegmentGroup]:
            """Convert list of atom_indices in a AtomGroup to either the
            Atoms or segments/residues containing these atoms.

            Parameters
            ----------
            indices
               list of atom indices
            level : str
              char (A, R, S). Return atoms(A), residues(R) or segments(S) within
              *radius* of *atoms*.
            """
            atomgroup = self.atom_group[indices]
            if level == "A":
                return atomgroup
            elif level == "R":
                return atomgroup.residues
            elif level == "S":
                return atomgroup.segments
            else:
                raise NotImplementedError(
                    "{0}: level not implemented".format(level)
                )