find_package(Threads REQUIRED)

add_library(elfes_tonari_geometry STATIC core/geometry.cpp)
target_compile_features(elfes_tonari_geometry PUBLIC cxx_std_20)
target_include_directories(
    elfes_tonari_geometry PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/core"
)
set_target_properties(
    elfes_tonari_geometry PROPERTIES POSITION_INDEPENDENT_CODE ON
)

add_library(
    elfes_tonari_cpu STATIC
    core/neighbors_cpu.cpp
    core/thread_pool.cpp
)
target_compile_features(elfes_tonari_cpu PUBLIC cxx_std_20)
target_include_directories(
    elfes_tonari_cpu PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/core"
)
target_link_libraries(
    elfes_tonari_cpu PUBLIC elfes_tonari_geometry Threads::Threads
)
set_target_properties(elfes_tonari_cpu PROPERTIES POSITION_INDEPENDENT_CODE ON)

add_library(elfes_tonari_torch INTERFACE)
target_include_directories(
    elfes_tonari_torch SYSTEM INTERFACE ${ELFES_TORCH_INCLUDE_DIRS}
)
target_compile_definitions(
    elfes_tonari_torch
    INTERFACE _GLIBCXX_USE_CXX11_ABI=${ELFES_TORCH_CXX11_ABI}
)
target_link_libraries(
    elfes_tonari_torch
    INTERFACE
        ${ELFES_C10_LIBRARY}
        ${ELFES_TORCH_LIBRARY}
        ${ELFES_TORCH_CPU_LIBRARY}
        ${ELFES_TORCH_PYTHON_LIBRARY}
)

function(elfes_configure_tonari_torch_module target module_name)
    target_compile_features(${target} PRIVATE cxx_std_20)
    target_compile_definitions(
        ${target}
        PRIVATE
            TORCH_API_INCLUDE_EXTENSION_H
            TORCH_EXTENSION_NAME=${module_name}
    )
    target_link_libraries(${target} PRIVATE elfes_tonari_torch)
    set_target_properties(
        ${target}
        PROPERTIES
            CXX_VISIBILITY_PRESET hidden
            VISIBILITY_INLINES_HIDDEN ON
            OUTPUT_NAME ${module_name}
            INSTALL_RPATH "$ORIGIN"
    )
endfunction()

pybind11_add_module(elfes_tonari_numpy_cpu MODULE numpy/bindings.cpp)
target_link_libraries(elfes_tonari_numpy_cpu PRIVATE elfes_tonari_cpu)
set_target_properties(
    elfes_tonari_numpy_cpu PROPERTIES OUTPUT_NAME _neighbor_numpy_cpu
)

Python_add_library(
    elfes_tonari_torch_cpu MODULE WITH_SOABI torch/bindings_cpu.cpp
)
elfes_configure_tonari_torch_module(
    elfes_tonari_torch_cpu _neighbor_torch_cpu
)
target_link_libraries(elfes_tonari_torch_cpu PRIVATE elfes_tonari_cpu)

set(elfes_tonari_python_targets
    elfes_tonari_numpy_cpu
    elfes_tonari_torch_cpu
)
if(BUILD_CUDA)
    Python_add_library(
        elfes_tonari_torch_cuda MODULE WITH_SOABI
        torch/bindings_cuda.cpp
        torch/neighbors_cuda.cu
        torch/neighbors_cell_cuda.cu
    )
    elfes_configure_tonari_torch_module(
        elfes_tonari_torch_cuda _neighbor_torch_cuda
    )
    target_compile_options(
        elfes_tonari_torch_cuda
        PRIVATE
            $<$<COMPILE_LANGUAGE:CUDA>:--expt-extended-lambda>
            $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
            $<$<COMPILE_LANGUAGE:CUDA>:-lineinfo>
    )
    target_link_libraries(
        elfes_tonari_torch_cuda
        PRIVATE
            elfes_tonari_geometry
            CUDA::cudart
            ${ELFES_C10_CUDA_LIBRARY}
            ${ELFES_TORCH_CUDA_LIBRARY}
    )
    list(APPEND elfes_tonari_python_targets elfes_tonari_torch_cuda)
endif()

install(
    TARGETS ${elfes_tonari_python_targets}
    LIBRARY DESTINATION elfes/native
    RUNTIME DESTINATION elfes/native
)
