# fmt: off

import itertools

import numpy as np
import scipy.sparse.csgraph as csgraph
from scipy import sparse as sp
from scipy.spatial import cKDTree

from ase.cell import Cell
from ase.data import atomic_numbers, covalent_radii
from ase.geometry import (
    complete_cell,
    find_mic,
    minkowski_reduce,
    wrap_positions,
)
from ase.utils import deprecated


def natural_cutoffs(atoms, mult=1, **kwargs):
    """Generate a radial cutoff for every atom based on covalent radii.

    The covalent radii are a reasonable cutoff estimation for bonds in
    many applications such as neighborlists, so function generates an
    atoms length list of radii based on this idea.

    * atoms: An atoms object
    * mult: A multiplier for all cutoffs, useful for coarse grained adjustment
    * kwargs: Symbol of the atom and its corresponding cutoff,
      used to override the covalent radii
    """
    return [kwargs.get(atom.symbol, covalent_radii[atom.number] * mult)
            for atom in atoms]


def build_neighbor_list(atoms, cutoffs=None, **kwargs):
    """Automatically build and update a NeighborList.

    Parameters:

    atoms : :class:`~ase.Atoms` object
        Atoms to build Neighborlist for.
    cutoffs: list of floats
        Radii for each atom. If not given it will be produced by calling
        :func:`ase.neighborlist.natural_cutoffs`
    kwargs: arbitrary number of options
        Will be passed to the constructor of
        :class:`~ase.neighborlist.NeighborList`

    Returns:

    return: :class:`~ase.neighborlist.NeighborList`
        A :class:`~ase.neighborlist.NeighborList` instance (updated).
    """
    if cutoffs is None:
        cutoffs = natural_cutoffs(atoms)

    nl = NeighborList(cutoffs, **kwargs)
    nl.update(atoms)

    return nl


def get_distance_matrix(graph, limit=3):
    """Get Distance Matrix from a Graph.

    Parameters:

    graph: array, matrix or sparse matrix, 2 dimensions (N, N)
        Graph representation of the connectivity.
        See `scipy doc <https://docs.scipy.org/doc/scipy/reference/generated\
/scipy.sparse.csgraph.dijkstra.html#scipy.sparse.csgraph.dijkstra>`_
        for reference.
    limit: integer
        Maximum number of steps to analyze. For most molecular information,
        three should be enough.

    Returns:

    return: scipy.sparse.csr_matrix, shape (N, N)
        A scipy.sparse.csr_matrix. All elements that are not connected within
        *limit* steps are set to zero.

    This is a potential memory bottleneck, as csgraph.dijkstra produces a
    dense output matrix. Here we replace all np.inf values with 0 and
    transform back to csr_matrix.
    Why not dok_matrix like the connectivity-matrix? Because row-picking
    is most likely and this is super fast with csr.
    """
    mat = csgraph.dijkstra(graph, directed=False, limit=limit)
    mat[mat == np.inf] = 0
    return sp.csr_matrix(mat, dtype=np.int8)


def get_distance_indices(distanceMatrix, distance):
    """Get indices for each node that are distance or less away.

    Parameters:

    distanceMatrix: any one of scipy.sparse matrices (NxN)
        Matrix containing distance information of atoms. Get it e.g. with
        :func:`~ase.neighborlist.get_distance_matrix`
    distance: integer
        Number of steps to allow.

    Returns:

    return: list of length N
        List of length N. return[i] has all indices connected to item i.

    The distance matrix only contains shortest paths, so when looking for
    distances longer than one, we need to add the lower values for cases
    where atoms are connected via a shorter path too.
    """
    shape = distanceMatrix.get_shape()
    indices = []
    # iterate over rows
    for i in range(shape[0]):
        row = distanceMatrix.getrow(i)[0]
        # find all non-zero
        found = sp.find(row)
        # screen for smaller or equal distance
        equal = np.where(found[-1] <= distance)[0]
        # found[1] contains the indexes
        indices.append([found[1][x] for x in equal])
    return indices


def mic(dr, cell, pbc=True):
    """
    Apply minimum image convention to an array of distance vectors.

    Parameters:

    dr : array_like
        Array of distance vectors.
    cell : array_like
        Simulation cell.
    pbc : array_like, optional
        Periodic boundary conditions in x-, y- and z-direction. Default is to
        assume periodic boundaries in all directions.

    Returns:

    dr : array
        Array of distance vectors, wrapped according to the minimum image
        convention.
    """
    dr, _ = find_mic(dr, cell, pbc)
    return dr


def primitive_neighbor_list(quantities, pbc, cell, positions, cutoff,
                            numbers=None, self_interaction=False,
                            use_scaled_positions=False, max_nbins=1e6):
    """Compute a neighbor list for an atomic configuration.

    Atoms outside periodic boundaries are mapped into the box. Atoms
    outside nonperiodic boundaries are included in the neighbor list
    but complexity of neighbor list search for those can become n^2.

    The neighbor list is sorted by first atom index 'i', but not by second
    atom index 'j'.

    Parameters:

    quantities: str
        Quantities to compute by the neighbor list algorithm. Each character
        in this string defines a quantity. They are returned in a tuple of
        the same order. Possible quantities are

            * 'i' : first atom index
            * 'j' : second atom index
            * 'd' : absolute distance
            * 'D' : distance vector
            * 'S' : shift vector (number of cell boundaries crossed by the bond
              between atom i and j). With the shift vector S, the
              distances D between atoms can be computed from:
              D = positions[j]-positions[i]+S.dot(cell)
    pbc: array_like
        3-tuple indicating giving periodic boundaries in the three Cartesian
        directions.
    cell: 3x3 matrix
        Unit cell vectors.
    positions: list of xyz-positions
        Atomic positions.  Anything that can be converted to an ndarray of
        shape (n, 3) will do: [(x1,y1,z1), (x2,y2,z2), ...]. If
        use_scaled_positions is set to true, this must be scaled positions.
    cutoff: float or dict
        Cutoff for neighbor search. It can be:

            * A single float: This is a global cutoff for all elements.
            * A dictionary: This specifies cutoff values for element
              pairs. Specification accepts element numbers of symbols.
              Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
            * A list/array with a per atom value: This specifies the radius of
              an atomic sphere for each atoms. If spheres overlap, atoms are
              within each others neighborhood. See
              :func:`~ase.neighborlist.natural_cutoffs`
              for an example on how to get such a list.
    self_interaction: bool
        Return the atom itself as its own neighbor if set to true.
        Default: False
    use_scaled_positions: bool
        If set to true, positions are expected to be scaled positions.
    max_nbins: int
        Maximum number of bins used in neighbor search. This is used to limit
        the maximum amount of memory required by the neighbor list.

    Returns:

    i, j, ... : array
        Tuple with arrays for each quantity specified above. Indices in `i`
        are returned in ascending order 0..len(a)-1, but the order of (i,j)
        pairs is not guaranteed.

    """

    # Naming conventions: Suffixes indicate the dimension of an array. The
    # following convention is used here:
    #     c: Cartesian index, can have values 0, 1, 2
    #     i: Global atom index, can have values 0..len(a)-1
    #     xyz: Bin index, three values identifying x-, y- and z-component of a
    #          spatial bin that is used to make neighbor search O(n)
    #     b: Linearized version of the 'xyz' bin index
    #     a: Bin-local atom index, i.e. index identifying an atom *within* a
    #        bin
    #     p: Pair index, can have value 0 or 1
    #     n: (Linear) neighbor index

    # Return empty neighbor list if no atoms are passed here
    if len(positions) == 0:
        empty_types = dict(i=(int, (0, )),
                           j=(int, (0, )),
                           D=(float, (0, 3)),
                           d=(float, (0, )),
                           S=(int, (0, 3)))
        retvals = []
        for i in quantities:
            dtype, shape = empty_types[i]
            retvals += [np.array([], dtype=dtype).reshape(shape)]
        if len(retvals) == 1:
            return retvals[0]
        else:
            return tuple(retvals)

    # Compute reciprocal lattice vectors.
    b1_c, b2_c, b3_c = np.linalg.pinv(cell).T

    # Compute distances of cell faces.
    l1 = np.linalg.norm(b1_c)
    l2 = np.linalg.norm(b2_c)
    l3 = np.linalg.norm(b3_c)
    face_dist_c = np.array([1 / l1 if l1 > 0 else 1,
                            1 / l2 if l2 > 0 else 1,
                            1 / l3 if l3 > 0 else 1])

    if isinstance(cutoff, dict):
        max_cutoff = max(cutoff.values())
    else:
        if np.isscalar(cutoff):
            max_cutoff = cutoff
        else:
            cutoff = np.asarray(cutoff)
            max_cutoff = 2 * np.max(cutoff)

    # We use a minimum bin size of 3 A
    bin_size = max(max_cutoff, 3)
    # Compute number of bins such that a sphere of radius cutoff fits into
    # eight neighboring bins.
    nbins_c = np.maximum((face_dist_c / bin_size).astype(int), [1, 1, 1])
    nbins = np.prod(nbins_c)
    # Make sure we limit the amount of memory used by the explicit bins.
    while nbins > max_nbins:
        nbins_c = np.maximum(nbins_c // 2, [1, 1, 1])
        nbins = np.prod(nbins_c)

    # Compute over how many bins we need to loop in the neighbor list search.
    neigh_search_x, neigh_search_y, neigh_search_z = \
        np.ceil(bin_size * nbins_c / face_dist_c).astype(int)

    # If we only have a single bin and the system is not periodic, then we
    # do not need to search neighboring bins
    neigh_search_x = 0 if nbins_c[0] == 1 and not pbc[0] else neigh_search_x
    neigh_search_y = 0 if nbins_c[1] == 1 and not pbc[1] else neigh_search_y
    neigh_search_z = 0 if nbins_c[2] == 1 and not pbc[2] else neigh_search_z

    # Sort atoms into bins.
    if use_scaled_positions:
        scaled_positions_ic = positions
        positions = np.dot(scaled_positions_ic, cell)
    else:
        scaled_positions_ic = np.linalg.solve(complete_cell(cell).T,
                                              positions.T).T
    bin_index_ic = np.floor(scaled_positions_ic * nbins_c).astype(int)
    cell_shift_ic = np.zeros_like(bin_index_ic)

    for c in range(3):
        if pbc[c]:
            # (Note: np.divmod does not exist in older numpies)
            cell_shift_ic[:, c], bin_index_ic[:, c] = \
                divmod(bin_index_ic[:, c], nbins_c[c])
        else:
            bin_index_ic[:, c] = np.clip(bin_index_ic[:, c], 0, nbins_c[c] - 1)

    # Convert Cartesian bin index to unique scalar bin index.
    bin_index_i = (bin_index_ic[:, 0] +
                   nbins_c[0] * (bin_index_ic[:, 1] +
                                 nbins_c[1] * bin_index_ic[:, 2]))

    # atom_i contains atom index in new sort order.
    atom_i = np.argsort(bin_index_i)
    bin_index_i = bin_index_i[atom_i]

    # Find max number of atoms per bin
    max_natoms_per_bin = np.bincount(bin_index_i).max()

    # Sort atoms into bins: atoms_in_bin_ba contains for each bin (identified
    # by its scalar bin index) a list of atoms inside that bin. This list is
    # homogeneous, i.e. has the same size *max_natoms_per_bin* for all bins.
    # The list is padded with -1 values.
    atoms_in_bin_ba = -np.ones([nbins, max_natoms_per_bin], dtype=int)
    for i in range(max_natoms_per_bin):
        # Create a mask array that identifies the first atom of each bin.
        mask = np.append([True], bin_index_i[:-1] != bin_index_i[1:])
        # Assign all first atoms.
        atoms_in_bin_ba[bin_index_i[mask], i] = atom_i[mask]

        # Remove atoms that we just sorted into atoms_in_bin_ba. The next
        # "first" atom will be the second and so on.
        mask = np.logical_not(mask)
        atom_i = atom_i[mask]
        bin_index_i = bin_index_i[mask]

    # Make sure that all atoms have been sorted into bins.
    assert len(atom_i) == 0
    assert len(bin_index_i) == 0

    # Now we construct neighbor pairs by pairing up all atoms within a bin or
    # between bin and neighboring bin. atom_pairs_pn is a helper buffer that
    # contains all potential pairs of atoms between two bins, i.e. it is a list
    # of length max_natoms_per_bin**2.
    atom_pairs_pn = np.indices((max_natoms_per_bin, max_natoms_per_bin),
                               dtype=int)
    atom_pairs_pn = atom_pairs_pn.reshape(2, -1)

    # Initialized empty neighbor list buffers.
    first_at_neightuple_nn = []
    secnd_at_neightuple_nn = []
    cell_shift_vector_x_n = []
    cell_shift_vector_y_n = []
    cell_shift_vector_z_n = []

    # This is the main neighbor list search. We loop over neighboring bins and
    # then construct all possible pairs of atoms between two bins, assuming
    # that each bin contains exactly max_natoms_per_bin atoms. We then throw
    # out pairs involving pad atoms with atom index -1 below.
    binz_xyz, biny_xyz, binx_xyz = np.meshgrid(np.arange(nbins_c[2]),
                                               np.arange(nbins_c[1]),
                                               np.arange(nbins_c[0]),
                                               indexing='ij')
    # The memory layout of binx_xyz, biny_xyz, binz_xyz is such that computing
    # the respective bin index leads to a linearly increasing consecutive list.
    # The following assert statement succeeds:
    #     b_b = (binx_xyz + nbins_c[0] * (biny_xyz + nbins_c[1] *
    #                                     binz_xyz)).ravel()
    #     assert (b_b == np.arange(np.prod(nbins_c))).all()

    # First atoms in pair.
    _first_at_neightuple_n = atoms_in_bin_ba[:, atom_pairs_pn[0]]
    for dz in range(-neigh_search_z, neigh_search_z + 1):
        for dy in range(-neigh_search_y, neigh_search_y + 1):
            for dx in range(-neigh_search_x, neigh_search_x + 1):
                # Bin index of neighboring bin and shift vector.
                shiftx_xyz, neighbinx_xyz = divmod(binx_xyz + dx, nbins_c[0])
                shifty_xyz, neighbiny_xyz = divmod(biny_xyz + dy, nbins_c[1])
                shiftz_xyz, neighbinz_xyz = divmod(binz_xyz + dz, nbins_c[2])
                neighbin_b = (neighbinx_xyz + nbins_c[0] *
                              (neighbiny_xyz + nbins_c[1] * neighbinz_xyz)
                              ).ravel()

                # Second atom in pair.
                _secnd_at_neightuple_n = \
                    atoms_in_bin_ba[neighbin_b][:, atom_pairs_pn[1]]

                # Shift vectors.
                _cell_shift_vector_x_n = \
                    np.resize(shiftx_xyz.reshape(-1, 1),
                              (max_natoms_per_bin**2, shiftx_xyz.size)).T
                _cell_shift_vector_y_n = \
                    np.resize(shifty_xyz.reshape(-1, 1),
                              (max_natoms_per_bin**2, shifty_xyz.size)).T
                _cell_shift_vector_z_n = \
                    np.resize(shiftz_xyz.reshape(-1, 1),
                              (max_natoms_per_bin**2, shiftz_xyz.size)).T

                # We have created too many pairs because we assumed each bin
                # has exactly max_natoms_per_bin atoms. Remove all surperfluous
                # pairs. Those are pairs that involve an atom with index -1.
                mask = np.logical_and(_first_at_neightuple_n != -1,
                                      _secnd_at_neightuple_n != -1)
                if mask.sum() > 0:
                    first_at_neightuple_nn += [_first_at_neightuple_n[mask]]
                    secnd_at_neightuple_nn += [_secnd_at_neightuple_n[mask]]
                    cell_shift_vector_x_n += [_cell_shift_vector_x_n[mask]]
                    cell_shift_vector_y_n += [_cell_shift_vector_y_n[mask]]
                    cell_shift_vector_z_n += [_cell_shift_vector_z_n[mask]]

    # Flatten overall neighbor list.
    first_at_neightuple_n = np.concatenate(first_at_neightuple_nn)
    secnd_at_neightuple_n = np.concatenate(secnd_at_neightuple_nn)
    cell_shift_vector_n = np.transpose([np.concatenate(cell_shift_vector_x_n),
                                        np.concatenate(cell_shift_vector_y_n),
                                        np.concatenate(cell_shift_vector_z_n)])

    # Add global cell shift to shift vectors
    cell_shift_vector_n += cell_shift_ic[first_at_neightuple_n] - \
        cell_shift_ic[secnd_at_neightuple_n]

    # Remove all self-pairs that do not cross the cell boundary.
    if not self_interaction:
        m = np.logical_not(np.logical_and(
            first_at_neightuple_n == secnd_at_neightuple_n,
            (cell_shift_vector_n == 0).all(axis=1)))
        first_at_neightuple_n = first_at_neightuple_n[m]
        secnd_at_neightuple_n = secnd_at_neightuple_n[m]
        cell_shift_vector_n = cell_shift_vector_n[m]

    # For nonperiodic directions, remove any bonds that cross the domain
    # boundary.
    for c in range(3):
        if not pbc[c]:
            m = cell_shift_vector_n[:, c] == 0
            first_at_neightuple_n = first_at_neightuple_n[m]
            secnd_at_neightuple_n = secnd_at_neightuple_n[m]
            cell_shift_vector_n = cell_shift_vector_n[m]

    # Sort neighbor list.
    i = np.argsort(first_at_neightuple_n)
    first_at_neightuple_n = first_at_neightuple_n[i]
    secnd_at_neightuple_n = secnd_at_neightuple_n[i]
    cell_shift_vector_n = cell_shift_vector_n[i]

    # Compute distance vectors.
    distance_vector_nc = positions[secnd_at_neightuple_n] - \
        positions[first_at_neightuple_n] + \
        cell_shift_vector_n.dot(cell)
    abs_distance_vector_n = \
        np.sqrt(np.sum(distance_vector_nc * distance_vector_nc, axis=1))

    # We have still created too many pairs. Only keep those with distance
    # smaller than max_cutoff.
    mask = abs_distance_vector_n < max_cutoff
    first_at_neightuple_n = first_at_neightuple_n[mask]
    secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
    cell_shift_vector_n = cell_shift_vector_n[mask]
    distance_vector_nc = distance_vector_nc[mask]
    abs_distance_vector_n = abs_distance_vector_n[mask]

    if isinstance(cutoff, dict) and numbers is not None:
        # If cutoff is a dictionary, then the cutoff radii are specified per
        # element pair. We now have a list up to maximum cutoff.
        per_pair_cutoff_n = np.zeros_like(abs_distance_vector_n)
        for (atomic_number1, atomic_number2), c in cutoff.items():
            try:
                atomic_number1 = atomic_numbers[atomic_number1]
            except KeyError:
                pass
            try:
                atomic_number2 = atomic_numbers[atomic_number2]
            except KeyError:
                pass
            if atomic_number1 == atomic_number2:
                mask = np.logical_and(
                    numbers[first_at_neightuple_n] == atomic_number1,
                    numbers[secnd_at_neightuple_n] == atomic_number2)
            else:
                mask = np.logical_or(
                    np.logical_and(
                        numbers[first_at_neightuple_n] == atomic_number1,
                        numbers[secnd_at_neightuple_n] == atomic_number2),
                    np.logical_and(
                        numbers[first_at_neightuple_n] == atomic_number2,
                        numbers[secnd_at_neightuple_n] == atomic_number1))
            per_pair_cutoff_n[mask] = c
        mask = abs_distance_vector_n < per_pair_cutoff_n
        first_at_neightuple_n = first_at_neightuple_n[mask]
        secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
        cell_shift_vector_n = cell_shift_vector_n[mask]
        distance_vector_nc = distance_vector_nc[mask]
        abs_distance_vector_n = abs_distance_vector_n[mask]
    elif not np.isscalar(cutoff):
        # If cutoff is neither a dictionary nor a scalar, then we assume it is
        # a list or numpy array that contains atomic radii. Atoms are neighbors
        # if their radii overlap.
        mask = abs_distance_vector_n < \
            cutoff[first_at_neightuple_n] + cutoff[secnd_at_neightuple_n]
        first_at_neightuple_n = first_at_neightuple_n[mask]
        secnd_at_neightuple_n = secnd_at_neightuple_n[mask]
        cell_shift_vector_n = cell_shift_vector_n[mask]
        distance_vector_nc = distance_vector_nc[mask]
        abs_distance_vector_n = abs_distance_vector_n[mask]

    # Assemble return tuple.
    retvals = []
    for q in quantities:
        if q == 'i':
            retvals += [first_at_neightuple_n]
        elif q == 'j':
            retvals += [secnd_at_neightuple_n]
        elif q == 'D':
            retvals += [distance_vector_nc]
        elif q == 'd':
            retvals += [abs_distance_vector_n]
        elif q == 'S':
            retvals += [cell_shift_vector_n]
        else:
            raise ValueError('Unsupported quantity specified.')
    if len(retvals) == 1:
        return retvals[0]
    else:
        return tuple(retvals)


def neighbor_list(quantities, a, cutoff, self_interaction=False,
                  max_nbins=1e6):
    """Compute a neighbor list for an atomic configuration.

    Atoms outside periodic boundaries are mapped into the box. Atoms
    outside nonperiodic boundaries are included in the neighbor list
    but complexity of neighbor list search for those can become n^2.

    The neighbor list is sorted by first atom index 'i', but not by second
    atom index 'j'.

    Parameters
    ----------
    quantities: str
        Quantities to compute by the neighbor list algorithm. Each character
        in this string defines a quantity. They are returned in a tuple of
        the same order. Possible quantities are:

           * 'i' : first atom index
           * 'j' : second atom index
           * 'd' : absolute distance
           * 'D' : distance vector
           * 'S' : shift vector (number of cell boundaries crossed by the bond
             between atom i and j). With the shift vector S, the
             distances D between atoms can be computed from:
             D = a.positions[j]-a.positions[i]+S.dot(a.cell)
    a: :class:`ase.Atoms`
        Atomic configuration.
    cutoff: float or dict
        Cutoff for neighbor search. It can be:

            * A single float: This is a global cutoff for all elements.
            * A dictionary: This specifies cutoff values for element
              pairs. Specification accepts element numbers of symbols.
              Example: {(1, 6): 1.1, (1, 1): 1.0, ('C', 'C'): 1.85}
            * A list/array with a per atom value: This specifies the radius of
              an atomic sphere for each atoms. If spheres overlap, atoms are
              within each others neighborhood. See
              :func:`~ase.neighborlist.natural_cutoffs`
              for an example on how to get such a list.

    self_interaction: bool
        Return the atom itself as its own neighbor if set to true.
        Default: False
    max_nbins: int
        Maximum number of bins used in neighbor search. This is used to limit
        the maximum amount of memory required by the neighbor list.

    Returns
    -------
    i, j, ...: array
        Tuple with arrays for each quantity specified above. Indices in `i`
        are returned in ascending order 0..len(a), but the order of (i,j)
        pairs is not guaranteed.

    Examples
    --------

    >>> import numpy as np
    >>> from ase.build import bulk, molecule

    1. Coordination counting

    >>> atoms = molecule('isobutane')
    >>> i = neighbor_list('i', atoms, 1.85)
    >>> coord = np.bincount(i, minlength=len(atoms))

    2. Coordination counting with different cutoffs for each pair of species

    >>> cutoff = {('H', 'H'): 1.1, ('C', 'H'): 1.3, ('C', 'C'): 1.85}
    >>> i = neighbor_list('i', atoms, cutoff)
    >>> coord = np.bincount(i, minlength=len(atoms))

    3. Pair distribution function

    >>> atoms = bulk('Cu', cubic=True) * 3
    >>> atoms.rattle(0.5, rng=np.random.default_rng(42))
    >>> cutoff = 5.0
    >>> d = neighbor_list('d', atoms, cutoff)
    >>> hist, bin_edges = np.histogram(d, bins=100, range=(0.0, cutoff))
    >>> hist = hist / len(atoms)  # per atom
    >>> rho_mean = len(atoms) / atoms.cell.volume
    >>> dv = 4.0 * np.pi * (bin_edges[1:] ** 3 - bin_edges[:-1] ** 3) / 3.0
    >>> rho = hist / dv
    >>> pdf = rho / rho_mean

    4. Forces of a pair potential

    >>> natoms = len(atoms)
    >>> i, j, d, D = neighbor_list('ijdD', atoms, 5.0)
    >>> # Lennard-Jones potential
    >>> eps = 1.0
    >>> sgm = 1.0
    >>> epairs = 4.0 * eps * ((sgm / d) ** 12 - (sgm / d) ** 6)
    >>> energy = 0.5 * epairs.sum()  # correct double-counting
    >>> dd = -4.0 * eps * (12 * (sgm / d) ** 13 - 6 * (sgm / d) ** 7) / sgm
    >>> dd = (dd * (D.T / d)).T
    >>> fx = -1.0 * np.bincount(i, weights=dd[:, 0], minlength=natoms)
    >>> fy = -1.0 * np.bincount(i, weights=dd[:, 1], minlength=natoms)
    >>> fz = -1.0 * np.bincount(i, weights=dd[:, 2], minlength=natoms)

    5. Force-constant matrix of a pair potential

    >>> i, j, d, D = neighbor_list('ijdD', atoms, 5.0)
    >>> epairs = 1.0 / d  # Coulomb potential
    >>> forces = (D.T / d**3).T  # (npairs, 3)
    >>> # second derivative
    >>> d2 = 3.0 * D[:, :, None] * D[:, None, :] / d[:, None, None] ** 5
    >>> for k in range(3):
    ...     d2[:, k, k] -= 1.0 / d**3
    >>> # force-constant matrix
    >>> fc = np.zeros((natoms, 3, natoms, 3))
    >>> for ia in range(natoms):
    ...     for ja in range(natoms):
    ...         fc[ia, :, ja, :] -= d2[(i == ia) & (j == ja), :, :].sum(axis=0)
    >>> for ia in range(natoms):
    ...     fc[ia, :, ia, :] -= fc[ia].sum(axis=1)

    """
    return primitive_neighbor_list(quantities, a.pbc,
                                   a.get_cell(complete=True),
                                   a.positions, cutoff, numbers=a.numbers,
                                   self_interaction=self_interaction,
                                   max_nbins=max_nbins)


def first_neighbors(natoms, first_atom):
    """
    Compute an index array pointing to the ranges within the neighbor list that
    contain the neighbors for a certain atom.

    Parameters:

    natoms : int
        Total number of atom.
    first_atom : array_like
        Array containing the first atom 'i' of the neighbor tuple returned
        by the neighbor list.

    Returns:

    seed : array
        Array containing pointers to the start and end location of the
        neighbors of a certain atom. Neighbors of atom k have indices from s[k]
        to s[k+1]-1.
    """
    if len(first_atom) == 0:
        return np.zeros(natoms + 1, dtype=int)
    # Create a seed array (which is returned by this function) populated with
    # -1.
    seed = -np.ones(natoms + 1, dtype=int)

    first_atom = np.asarray(first_atom)

    # Mask array contains all position where the number in the (sorted) array
    # with first atoms (in the neighbor pair) changes.
    mask = first_atom[:-1] != first_atom[1:]

    # Seed array needs to start at 0
    seed[first_atom[0]] = 0
    # Seed array needs to stop at the length of the neighbor list
    seed[-1] = len(first_atom)
    # Populate all intermediate seed with the index of where the mask array is
    # true, i.e. the index where the first_atom array changes.
    seed[first_atom[1:][mask]] = (np.arange(len(mask)) + 1)[mask]

    # Now fill all remaining -1 value with the value in the seed array right
    # behind them. (There are no neighbor so seed[i] and seed[i+1] must point)
    # to the same index.
    mask = seed == -1
    while mask.any():
        seed[mask] = seed[np.arange(natoms + 1)[mask] + 1]
        mask = seed == -1
    return seed


def get_connectivity_matrix(nl, sparse=True):
    """Return connectivity matrix for a given NeighborList (dtype=numpy.int8).

    A matrix of shape (nAtoms, nAtoms) will be returned.
    Connected atoms i and j will have matrix[i,j] == 1, unconnected
    matrix[i,j] == 0. If bothways=True the matrix will be symmetric,
    otherwise not!

    If *sparse* is True, a scipy csr matrix is returned.
    If *sparse* is False, a numpy matrix is returned.

    Note that the old and new neighborlists might give different results
    for periodic systems if bothways=False.

    Example:

    Determine which molecule in a system atom 1 belongs to.

    >>> from scipy import sparse

    >>> from ase.build import molecule
    >>> from ase.neighborlist import get_connectivity_matrix
    >>> from ase.neighborlist import natural_cutoffs
    >>> from ase.neighborlist import NeighborList

    >>> mol = molecule('CH3CH2OH')
    >>> cutoffs = natural_cutoffs(mol)
    >>> neighbor_list = NeighborList(
    ...     cutoffs, self_interaction=False, bothways=True)
    >>> neighbor_list.update(mol)
    True

    >>> matrix = neighbor_list.get_connectivity_matrix()
    >>> # or: matrix = get_connectivity_matrix(neighbor_list.nl)
    >>> n_components, component_list = sparse.csgraph.connected_components(
    ...    matrix)
    >>> idx = 1
    >>> molIdx = component_list[idx]
    >>> print("There are {} molecules in the system".format(n_components))
    There are 1 molecules in the system
    >>> print("Atom {} is part of molecule {}".format(idx, molIdx))
    Atom 1 is part of molecule 0
    >>> molIdxs = [i for i in range(len(component_list))
    ...            if component_list[i] == molIdx]
    >>> print("Atoms are part of molecule {}: {}".format(molIdx, molIdxs))
    Atoms are part of molecule 0: [0, 1, 2, 3, 4, 5, 6, 7, 8]
    """

    nAtoms = len(nl.cutoffs)

    if nl.nupdates <= 0:
        raise RuntimeError(
            'Must call update(atoms) on your neighborlist first!')

    if sparse:
        matrix = sp.dok_matrix((nAtoms, nAtoms), dtype=np.int8)
    else:
        matrix = np.zeros((nAtoms, nAtoms), dtype=np.int8)

    for i in range(nAtoms):
        for idx in nl.get_neighbors(i)[0]:
            matrix[i, idx] = 1

    return matrix


class NewPrimitiveNeighborList:
    """Neighbor list object. Wrapper around neighbor_list and first_neighbors.

    cutoffs: list of float
        List of cutoff radii - one for each atom. If the spheres (defined by
        their cutoff radii) of two atoms overlap, they will be counted as
        neighbors.
    skin: float
        If no atom has moved more than the skin-distance since the
        last call to the
        :meth:`~ase.neighborlist.NewPrimitiveNeighborList.update()`
        method, then the neighbor list can be reused. This will save
        some expensive rebuilds of the list, but extra neighbors outside
        the cutoff will be returned.
    sorted: bool
        Sort neighbor list.
    self_interaction: bool
        Should an atom return itself as a neighbor?
    bothways: bool
        Return all neighbors.  Default is to return only "half" of
        the neighbors.

    Example:

    >>> from ase.build import bulk
    >>> from ase.neighborlist import NewPrimitiveNeighborList

    >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
    >>> atoms = bulk('Cu', 'fcc', a=3.6)
    >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
    True
    >>> indices, offsets = nl.get_neighbors(0)
    """

    def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
                 bothways=False, use_scaled_positions=False):
        self.cutoffs = np.asarray(cutoffs) + skin
        self.skin = skin
        self.sorted = sorted
        self.self_interaction = self_interaction
        self.bothways = bothways
        self.nupdates = 0
        self.use_scaled_positions = use_scaled_positions

    def update(self, pbc, cell, positions, numbers=None):
        """Make sure the list is up to date."""

        if self.nupdates == 0:
            self.build(pbc, cell, positions, numbers=numbers)
            return True

        if ((self.pbc != pbc).any() or (self.cell != cell).any() or
                ((self.positions - positions)**2).sum(1).max() > self.skin**2):
            self.build(pbc, cell, positions, numbers=numbers)
            return True

        return False

    def build(self, pbc, cell, positions, numbers=None):
        """Build the list.
        """
        self.pbc = np.array(pbc, copy=True)
        self.cell = np.array(cell, copy=True)
        self.positions = np.array(positions, copy=True)

        pair_first, pair_second, offset_vec = \
            primitive_neighbor_list(
                'ijS', pbc, cell, positions, self.cutoffs, numbers=numbers,
                self_interaction=self.self_interaction,
                use_scaled_positions=self.use_scaled_positions)

        if len(positions) > 0 and not self.bothways:
            offset_x, offset_y, offset_z = offset_vec.T

            mask = offset_z > 0
            mask &= offset_y == 0
            mask |= offset_y > 0
            mask &= offset_x == 0
            mask |= offset_x > 0
            mask |= (pair_first <= pair_second) & (offset_vec == 0).all(axis=1)

            pair_first = pair_first[mask]
            pair_second = pair_second[mask]
            offset_vec = offset_vec[mask]

        if len(positions) > 0 and self.sorted:
            mask = np.argsort(pair_first * len(pair_first) +
                              pair_second)
            pair_first = pair_first[mask]
            pair_second = pair_second[mask]
            offset_vec = offset_vec[mask]

        self.pair_first = pair_first
        self.pair_second = pair_second
        self.offset_vec = offset_vec

        # Compute the index array point to the first neighbor
        self.first_neigh = first_neighbors(len(positions), pair_first)

        self.nupdates += 1

    def get_neighbors(self, a):
        """Return neighbors of atom number a.

        A list of indices and offsets to neighboring atoms is
        returned.  The positions of the neighbor atoms can be
        calculated like this:

        >>> from ase.build import bulk
        >>> from ase.neighborlist import NewPrimitiveNeighborList

        >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
        >>> atoms = bulk('Cu', 'fcc', a=3.6)
        >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
        True
        >>> indices, offsets = nl.get_neighbors(0)
        >>> for i, offset in zip(indices, offsets):
        ...     print(atoms.positions[i] + offset @ atoms.get_cell())
        ... # doctest: +SKIP

        Notice that if get_neighbors(a) gives atom b as a neighbor,
        then get_neighbors(b) will not return a as a neighbor - unless
        bothways=True was used."""

        return (self.pair_second[self.first_neigh[a]:self.first_neigh[a + 1]],
                self.offset_vec[self.first_neigh[a]:self.first_neigh[a + 1]])


class PrimitiveNeighborList:
    """Neighbor list that works without Atoms objects.

    This is less fancy, but can be used to avoid conversions between
    scaled and non-scaled coordinates which may affect cell offsets
    through rounding errors.

    Attributes
    ----------
    nupdates : int
        Number of updated times.
    """

    def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
                 bothways=False, use_scaled_positions=False):
        self.cutoffs = np.asarray(cutoffs) + skin
        self.skin = skin
        self.sorted = sorted
        self.self_interaction = self_interaction
        self.bothways = bothways
        self.nupdates = 0
        self.use_scaled_positions = use_scaled_positions

    def update(self, pbc, cell, coordinates):
        """Make sure the list is up to date.

        Returns
        -------
        bool
            True if the neighbor list is updated.
        """

        if self.nupdates == 0:
            self.build(pbc, cell, coordinates)
            return True

        if ((self.pbc != pbc).any() or (self.cell != cell).any() or (
                (self.coordinates
                 - coordinates)**2).sum(1).max() > self.skin**2):
            self.build(pbc, cell, coordinates)
            return True

        return False

    def build(self, pbc, cell, coordinates):
        """Build the list.

        Coordinates are taken to be scaled or not according
        to self.use_scaled_positions.
        """
        self.pbc = pbc = np.array(pbc, copy=True)
        self.cell = cell = Cell(cell)
        self.coordinates = coordinates = np.array(coordinates, copy=True)

        if len(self.cutoffs) != len(coordinates):
            raise ValueError('Wrong number of cutoff radii: {} != {}'
                             .format(len(self.cutoffs), len(coordinates)))

        rcmax = self.cutoffs.max() if len(self.cutoffs) > 0 else 0.0

        if self.use_scaled_positions:
            positions0 = cell.cartesian_positions(coordinates)
        else:
            positions0 = coordinates

        rcell, op = minkowski_reduce(cell, pbc)
        positions = wrap_positions(positions0, rcell, pbc=pbc, eps=0)

        natoms = len(positions)
        self.nupdates += 1
        if natoms == 0:
            self.neighbors = []
            self.displacements = []
            return

        tree = cKDTree(positions, copy_data=True)
        offsets = cell.scaled_positions(positions - positions0)
        offsets = offsets.round().astype(int)

        N = _calc_expansion(rcell, pbc, rcmax)

        neighbor_indices_a = [[] for _ in range(natoms)]
        displacements_a = [[] for _ in range(natoms)]

        for n1, n2, n3 in itertools.product(range(N[0] + 1),
                                            range(-N[1], N[1] + 1),
                                            range(-N[2], N[2] + 1)):
            if n1 == 0 and (n2 < 0 or n2 == 0 and n3 < 0):
                continue

            displacement = (n1, n2, n3) @ rcell
            shift0 = (n1, n2, n3) @ op
            indices_all = tree.query_ball_point(
                positions - displacement,
                r=self.cutoffs + rcmax,
            )

            for a in range(natoms):
                indices = indices_all[a]

                if not indices:
                    continue

                indices = np.array(indices)
                delta = positions[indices] + displacement - positions[a]
                distances = np.sqrt(np.add.reduce(delta**2, axis=1))
                cutoffs = self.cutoffs[indices] + self.cutoffs[a]
                i = indices[distances < cutoffs]
                if n1 == 0 and n2 == 0 and n3 == 0:
                    if self.self_interaction:
                        i = i[i >= a]
                    else:
                        i = i[i > a]

                neighbor_indices_a[a].append(i)

                disp = shift0 + offsets[i] - offsets[a]
                displacements_a[a].append(disp)

        self.neighbors = [np.concatenate(i) for i in neighbor_indices_a]
        self.displacements = [np.concatenate(d) for d in displacements_a]

        if self.bothways:
            neighbors2 = [[] for a in range(natoms)]
            displacements2 = [[] for a in range(natoms)]
            for a in range(natoms):
                for b, disp in zip(self.neighbors[a], self.displacements[a]):
                    # avoid double counting of self interaction
                    if a == b and (disp == 0).all():
                        continue
                    neighbors2[b].append(a)
                    displacements2[b].append(-disp)
            for a in range(natoms):
                nbs = np.concatenate((self.neighbors[a], neighbors2[a]))
                disp = np.array(list(self.displacements[a]) + displacements2[a])
                # Force correct type and shape for case of no neighbors:
                self.neighbors[a] = nbs.astype(int)
                self.displacements[a] = disp.astype(int).reshape((-1, 3))

        if self.sorted:
            _sort_neighbors(self.neighbors, self.displacements)

    def get_neighbors(self, a):
        """Return neighbors of atom number a.

        A list of indices and offsets to neighboring atoms is
        returned.  The positions of the neighbor atoms can be
        calculated like this:

        >>> from ase.build import bulk
        >>> from ase.neighborlist import NewPrimitiveNeighborList

        >>> nl = NewPrimitiveNeighborList([2.3, 1.7])
        >>> atoms = bulk('Cu', 'fcc', a=3.6)
        >>> nl.update(atoms.pbc, atoms.get_cell(), atoms.positions)
        True
        >>> indices, offsets = nl.get_neighbors(0)
        >>> for i, offset in zip(indices, offsets):
        ...     print(atoms.positions[i] + offset @ atoms.get_cell())
        ... # doctest: +SKIP

        Notice that if get_neighbors(a) gives atom b as a neighbor,
        then get_neighbors(b) will not return a as a neighbor - unless
        bothways=True was used."""

        return self.neighbors[a], self.displacements[a]


def _calc_expansion(rcell, pbc, rcmax):
    r"""Calculate expansion to contain a sphere of radius `2.0 * rcmax`.

    This function determines the minimum supercell (parallelepiped) that
    contains a sphere of radius `2.0 * rcmax`. For this, `a_1` is projected
    onto the unit vector perpendicular to `a_2 \times a_3` (i.e. the unit
    vector along the direction `b_1`) to know how many `a_1`'s the supercell
    takes to contain the sphere.
    """
    ircell = np.linalg.pinv(rcell)
    vs = np.sqrt(np.add.reduce(ircell**2, axis=0))
    ns = np.where(pbc, np.ceil(2.0 * rcmax * vs), 0.0)
    return ns.astype(int)


def _sort_neighbors(neighbors, offsets):
    """Sort neighbors first by indices and then offsets."""
    natoms = len(neighbors)
    for a in range(natoms):
        keys = (
            offsets[a][:, 2],
            offsets[a][:, 1],
            offsets[a][:, 0],
            neighbors[a]
        )
        mask = np.lexsort(keys)
        neighbors[a] = neighbors[a][mask]
        offsets[a] = offsets[a][mask]


class NeighborList:
    """Neighbor list object.

    cutoffs: list of float
        List of cutoff radii - one for each atom. If the spheres
        (defined by their cutoff radii) of two atoms overlap, they
        will be counted as neighbors. See
        :func:`~ase.neighborlist.natural_cutoffs` for an example on
        how to get such a list.

    skin: float
        If no atom has moved more than the skin-distance since the
        last call to the
        :meth:`~ase.neighborlist.NeighborList.update()` method, then
        the neighbor list can be reused.  This will save some
        expensive rebuilds of the list, but extra neighbors outside
        the cutoff will be returned.
    self_interaction: bool
        Should an atom return itself as a neighbor?
    bothways: bool
        Return all neighbors.  Default is to return only "half" of
        the neighbors.
    primitive: class
        Define which implementation to use. Older and quadratically-scaling
        :class:`~ase.neighborlist.PrimitiveNeighborList` or newer and
        linearly-scaling :class:`~ase.neighborlist.NewPrimitiveNeighborList`.

    Example:

    >>> from ase.build import molecule
    >>> from ase.neighborlist import NeighborList

    >>> atoms = molecule("CO")
    >>> nl = NeighborList([0.76, 0.66])
    >>> nl.update(atoms)
    True
    >>> indices, offsets = nl.get_neighbors(0)

    """

    def __init__(self, cutoffs, skin=0.3, sorted=False, self_interaction=True,
                 bothways=False, primitive=PrimitiveNeighborList):
        self.nl = primitive(cutoffs, skin, sorted,
                            self_interaction=self_interaction,
                            bothways=bothways)

    def update(self, atoms):
        """
        See :meth:`ase.neighborlist.PrimitiveNeighborList.update` or
        :meth:`ase.neighborlist.PrimitiveNeighborList.update`.
        """
        return self.nl.update(atoms.pbc, atoms.get_cell(complete=True),
                              atoms.positions)

    def get_neighbors(self, a):
        """
        See :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors` or
        :meth:`ase.neighborlist.PrimitiveNeighborList.get_neighbors`.
        """
        if self.nl.nupdates <= 0:
            raise RuntimeError('Must call update(atoms) on your neighborlist '
                               'first!')

        return self.nl.get_neighbors(a)

    def get_connectivity_matrix(self, sparse=True):
        """
        See :func:`~ase.neighborlist.get_connectivity_matrix`.
        """
        return get_connectivity_matrix(self.nl, sparse)

    @property
    def nupdates(self):
        """Get number of updates."""
        return self.nl.nupdates

    @property
    @deprecated(
        'Use, e.g., `sum(_.size for _ in nl.neighbors)` '
        'for `bothways=False` and `self_interaction=False`.'
    )
    def nneighbors(self):
        """Get number of neighbors.

        .. deprecated:: 3.24.0
        """
        nneighbors = sum(indices.size for indices in self.nl.neighbors)
        if self.nl.self_interaction:
            nneighbors -= len(self.nl.neighbors)
        return nneighbors // 2 if self.nl.bothways else nneighbors

    @property
    @deprecated(
        'Use, e.g., `sum(_.any(1).sum() for _ in nl.displacements)` '
        'for `bothways=False` and `self_interaction=False`.'
    )
    def npbcneighbors(self):
        """Get number of pbc neighbors.

        .. deprecated:: 3.24.0
        """
        nneighbors = sum(
            offsets.any(axis=1).sum() for offsets in self.nl.displacements
        )  # sum up all neighbors that have non-zero supercell offsets
        return nneighbors // 2 if self.nl.bothways else nneighbors