import ctypes
import os
import numpy as np

_lib_path = os.path.abspath('./librsr.so')
_lib = None
if os.path.exists(_lib_path):
    try:
        _lib = ctypes.CDLL(_lib_path)
        # void fmm_octree_evaluate(const float* src_points, const float* trg_points, float* potentials, int N, int M)
        _lib.fmm_octree_evaluate.argtypes = [
            np.ctypeslib.ndpointer(dtype=np.float32, ndim=1, flags='C_CONTIGUOUS'),
            np.ctypeslib.ndpointer(dtype=np.float32, ndim=1, flags='C_CONTIGUOUS'),
            np.ctypeslib.ndpointer(dtype=np.float32, ndim=1, flags='C_CONTIGUOUS'),
            ctypes.c_int, ctypes.c_int
        ]
        _lib.fmm_octree_evaluate.restype = None
    except Exception as e:
        print(f"Failed to load FMM C++ symbol: {e}")

def evaluate(src_points, trg_points):
    """
    src_points: (N, 3) float32
    trg_points: (M, 3) float32
    Returns: potentials (M,)
    """
    if _lib is None:
        raise ImportError("C++ FMM Library not loaded")

    N = src_points.shape[0]
    M = trg_points.shape[0]
    potentials = np.zeros(M, dtype=np.float32)

    src_flat = src_points.flatten().astype(np.float32)
    trg_flat = trg_points.flatten().astype(np.float32)

    _lib.fmm_octree_evaluate(src_flat, trg_flat, potentials, N, M)
    return potentials
