# Path: flare/kernels/cutoffs.py
# @njit
# def quadratic_cutoff(r_cut: float, ri: float, ci: float):
#     """A quadratic cutoff that goes to zero smoothly at the cutoff boundary.
# 
#     Args:
#         r_cut (float): Cutoff value (in angstrom).
#         ri (float): Interatomic distance.
#         ci (float): Cartesian coordinate divided by the distance.
# 
#     Returns:
#         (float, float): Cutoff value and its derivative.
#     """
# 
#     rdiff = r_cut - ri
#     fi = rdiff * rdiff
#     fdi = 2 * rdiff * ci
# 
#     return fi, fdi


import numpy as np
from numba import njit
from math import exp, floor
from typing import Callable
from flare.kernels.cutoffs import quadratic_cutoff
from time import time




def grid_kernel(struc, *args):

    if not isinstance(struc, list):
        struc = [struc]

    kern = 0
    for env in struc:
        kern += grid_kernel_env(env, *args)

    return kern


def grid_kernel_env(
    env1,
    bodies,
    kern_type,
    get_bonds_func,
    bonds_cutoff_func,
    c2,
    etypes2,
    hyps,
    r_cut,
    grids,
    fj,
    fdj,

# Path: flare/kernels/cutoffs.py
# @njit
# def quadratic_cutoff(r_cut: float, ri: float, ci: float):
#     """A quadratic cutoff that goes to zero smoothly at the cutoff boundary.
# 
#     Args:
#         r_cut (float): Cutoff value (in angstrom).
#         ri (float): Interatomic distance.
#         ci (float): Cartesian coordinate divided by the distance.
# 
#     Returns:
#         (float, float): Cutoff value and its derivative.
#     """
# 
#     rdiff = r_cut - ri
#     fi = rdiff * rdiff
#     fdi = 2 * rdiff * ci
# 
#     return fi, fdi


# Path: flare/mgp/grid_kernels.py
import numpy as np
from numba import njit
from math import exp, floor
from typing import Callable
from flare.kernels.cutoffs import quadratic_cutoff
from time import time





def grid_kernel(struc, *args):

    if not isinstance(struc, list):
        struc = [struc]

    kern = 0
    for env in struc:
        kern += grid_kernel_env(env, *args)

    return kern


def grid_kernel_env(
    env1,
    bodies,
    kern_type,
    get_bonds_func,
    bonds_cutoff_func,
    c2,
    etypes2,
    hyps,
    r_cut,
    grids,
    fj,
    fdj,

    cutoff_func: Callable = quadratic_cutoff,