# Path: flare/kernels/kernels.py
# @njit
# def coordination_number(rij, cij, r_cut, cutoff_func):
#     """Pairwise contribution to many-body descriptor based on number of
#         atoms in the environment
# 
#     Args:
#         rij (float): distance between atoms i and j
#         cij (float): Component of versor of rij along given direction
#         r_cut (float): cutoff hyperparameter
#         cutoff_func (callable): cutoff function
#     Return:
#         float: the value of the pairwise many-body contribution
#         float: the value of the derivative of the pairwise many-body
#         contribution w.r.t. the central atom displacement
#     """
# 
#     fij, fdij = cutoff_func(r_cut, rij, cij)
# 
#     return fij, fdij
# 
# @njit
# def q_value_mc(
#     distances, r_cut, ref_species, species, cutoff_func, q_func=coordination_number
# ):
#     """Compute value of many-body many components descriptor based
#     on distances of atoms in the local many-body environment.
# 
#     Args:
#         distances (np.ndarray): distances between atoms i and j
#         r_cut (float): cutoff hyperparameter
#         ref_species (int): species to consider to compute the contribution
#         species (np.ndarray): atomic species of neighbours
#         cutoff_func (callable): cutoff function
#         q_func (callable): many-body pairwise descrptor function
# 
#     Return:
#         float: the value of the many-body descriptor
#     """
# 
#     q = 0.0
# 
#     for i in range(len(distances)):
#         if species[i] == ref_species:
#             q_, _ = q_func(distances[i], 0, r_cut, cutoff_func)
#             q += q_
# 
#     return q


from math import sqrt
from numba import njit
from flare.kernels.kernels import coordination_number, q_value_mc
import numpy as np
import flare.kernels.cutoffs as cf
                and (manybody_cutoff_list is not None)
            ):
                bs = spec_mask[species_list[s]]
                mbtype = manybody_mask[bs + ben]
                r_cut = manybody_cutoff_list[mbtype]

            qs_neigh[i, s] = q_value_mc(
                neigh_bond_array[:, 0],
                r_cut,
                species_list[s],
                neigh_etypes,
                cutoff_func,
            )

    # get grad from each neighbor atom
    for i in range(n_bonds):
        if (
            sepcut
            and (spec_mask is not None)
            and (manybody_mask is not None)
            and (manybody_cutoff_list is not None)
        ):
            be = spec_mask[etypes[i]]
            mbtype = manybody_mask[bcn + be]
            r_cut = manybody_cutoff_list[mbtype]

        ri = bond_array_mb[i, 0]
        for d in range(3):
            ci = bond_array_mb[i, d + 1]


# Path: flare/kernels/kernels.py
# @njit
# def coordination_number(rij, cij, r_cut, cutoff_func):
#     """Pairwise contribution to many-body descriptor based on number of
#         atoms in the environment
# 
#     Args:
#         rij (float): distance between atoms i and j
#         cij (float): Component of versor of rij along given direction
#         r_cut (float): cutoff hyperparameter
#         cutoff_func (callable): cutoff function
#     Return:
#         float: the value of the pairwise many-body contribution
#         float: the value of the derivative of the pairwise many-body
#         contribution w.r.t. the central atom displacement
#     """
# 
#     fij, fdij = cutoff_func(r_cut, rij, cij)
# 
#     return fij, fdij
# 
# @njit
# def q_value_mc(
#     distances, r_cut, ref_species, species, cutoff_func, q_func=coordination_number
# ):
#     """Compute value of many-body many components descriptor based
#     on distances of atoms in the local many-body environment.
# 
#     Args:
#         distances (np.ndarray): distances between atoms i and j
#         r_cut (float): cutoff hyperparameter
#         ref_species (int): species to consider to compute the contribution
#         species (np.ndarray): atomic species of neighbours
#         cutoff_func (callable): cutoff function
#         q_func (callable): many-body pairwise descrptor function
# 
#     Return:
#         float: the value of the many-body descriptor
#     """
# 
#     q = 0.0
# 
#     for i in range(len(distances)):
#         if species[i] == ref_species:
#             q_, _ = q_func(distances[i], 0, r_cut, cutoff_func)
#             q += q_
# 
#     return q


# Path: flare/utils/env_getarray.py
from math import sqrt
from numba import njit
from flare.kernels.kernels import coordination_number, q_value_mc
import numpy as np
import flare.kernels.cutoffs as cf

                and (manybody_cutoff_list is not None)
            ):
                bs = spec_mask[species_list[s]]
                mbtype = manybody_mask[bs + ben]
                r_cut = manybody_cutoff_list[mbtype]

            qs_neigh[i, s] = q_value_mc(
                neigh_bond_array[:, 0],
                r_cut,
                species_list[s],
                neigh_etypes,
                cutoff_func,
            )

    # get grad from each neighbor atom
    for i in range(n_bonds):
        if (
            sepcut
            and (spec_mask is not None)
            and (manybody_mask is not None)
            and (manybody_cutoff_list is not None)
        ):
            be = spec_mask[etypes[i]]
            mbtype = manybody_mask[bcn + be]
            r_cut = manybody_cutoff_list[mbtype]

        ri = bond_array_mb[i, 0]
        for d in range(3):
            ci = bond_array_mb[i, d + 1]


            ____, q_neigh_grads[i, d] = coordination_number(ri, ci, r_cut, cutoff_func)