# Build the libann shared library for pyann
#
# This CMakeLists.txt is designed to be used as a subdirectory build
# that produces a shared library (.dll, .so, .dylib) from the ann source.
#
# Usage:
#   mkdir build && cd build
#   cmake ..
#   cmake --build . --config Release
#
# The resulting library will be copied to pyann/lib/

cmake_minimum_required(VERSION 3.15)
project(pyann_native VERSION 0.1.0 LANGUAGES C)

# Options
option(USE_BLAS "Use OpenBLAS for acceleration" OFF)
option(USE_CBLAS "Use CBLAS for acceleration" OFF)
option(USE_ACCELERATE "Use Apple Accelerate for BLAS (macOS only)" OFF)
option(USE_SCIPY_OPENBLAS "Use the prebuilt OpenBLAS from the scipy-openblas32 Python package" OFF)
option(USE_METAL "Use Apple Metal GPU acceleration (macOS only)" OFF)
option(USE_CUDA "Use NVIDIA CUDA GPU acceleration" OFF)

# Find the ann source (expects it as a submodule or sibling directory)
set(ANN_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/ann" CACHE PATH "Path to ann source directory")

if(NOT EXISTS "${ANN_SOURCE_DIR}/ann.c")
    message(FATAL_ERROR 
        "Could not find ann source at ${ANN_SOURCE_DIR}\n"
        "Please either:\n"
        "  1. Clone ann as a submodule: git submodule add https://github.com/mseminatore/ann.git ann\n"
        "  2. Set ANN_SOURCE_DIR to the correct path: cmake -DANN_SOURCE_DIR=/path/to/ann .."
    )
endif()

# Source files for the shared library
set(ANN_SOURCES
    ${ANN_SOURCE_DIR}/ann.c
    ${ANN_SOURCE_DIR}/tensor.c
    ${ANN_SOURCE_DIR}/json.c
    ${CMAKE_CURRENT_SOURCE_DIR}/src/native/pyann_info.c
)

# Create shared library
add_library(ann SHARED ${ANN_SOURCES})

# Set properties for shared library
set_target_properties(ann PROPERTIES
    POSITION_INDEPENDENT_CODE ON
    C_STANDARD 99
    C_STANDARD_REQUIRED ON
)
# Wheels can't hold the symlinks a versioned library needs, so only
# versioned in development builds
if(NOT SKBUILD)
    set_target_properties(ann PROPERTIES VERSION ${PROJECT_VERSION} SOVERSION 0)
endif()

# Include directories
target_include_directories(ann PUBLIC ${ANN_SOURCE_DIR})

# Platform-specific settings
if(WIN32)
    # Windows: export all symbols
    set_target_properties(ann PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON)
    set(LIB_OUTPUT_NAME "ann")
elseif(APPLE)
    set(LIB_OUTPUT_NAME "libann")
    set_target_properties(ann PROPERTIES
        INSTALL_RPATH "@loader_path"
        BUILD_WITH_INSTALL_RPATH TRUE
    )
else()
    # Linux
    set(LIB_OUTPUT_NAME "libann")
    target_link_libraries(ann PRIVATE m)  # Math library
    set_target_properties(ann PROPERTIES
        INSTALL_RPATH "$ORIGIN"
        BUILD_WITH_INSTALL_RPATH TRUE
    )
endif()

# BLAS support
if(USE_CBLAS)
    target_compile_definitions(ann PRIVATE USE_CBLAS=1 USE_BLAS=1)
    set_target_properties(ann PROPERTIES C_STANDARD 11)

    # Find CBLAS
    if(WIN32)
        set(CBLAS_ROOT "C:/opt/cblas" CACHE PATH "CBLAS installation directory")
    else()
        set(CBLAS_ROOT "/opt/cblas" CACHE PATH "CBLAS installation directory")
    endif()

    if(EXISTS "${CBLAS_ROOT}/include/cblas.h")
        target_include_directories(ann PRIVATE ${CBLAS_ROOT}/include)
        target_link_directories(ann PRIVATE ${CBLAS_ROOT}/lib)
        target_link_libraries(ann PRIVATE cblas)
    else()
        message(WARNING "CBLAS not found at ${CBLAS_ROOT}")
    endif()
elseif(USE_ACCELERATE)
    if(NOT APPLE)
        message(FATAL_ERROR "USE_ACCELERATE is only supported on macOS")
    endif()
    # libann includes <cblas.h>; point it at Accelerate's copy.
    execute_process(COMMAND xcrun --show-sdk-path OUTPUT_VARIABLE MACOS_SDK
                    OUTPUT_STRIP_TRAILING_WHITESPACE)
    set(VECLIB_HEADERS "${MACOS_SDK}/System/Library/Frameworks/Accelerate.framework/Frameworks/vecLib.framework/Headers")
    if(NOT EXISTS "${VECLIB_HEADERS}/cblas.h")
        message(FATAL_ERROR "Accelerate cblas.h not found at ${VECLIB_HEADERS}")
    endif()
    # Accelerate names cblas_saxpby catlas_saxpby (same signature)
    target_compile_definitions(ann PRIVATE USE_BLAS=1 cblas_saxpby=catlas_saxpby)
    target_include_directories(ann PRIVATE ${VECLIB_HEADERS})
    # The legacy CBLAS header is marked deprecated on macOS 13.3+
    target_compile_options(ann PRIVATE -Wno-deprecated-declarations)
    target_link_libraries(ann PRIVATE "-framework Accelerate")
elseif(USE_SCIPY_OPENBLAS)
    # Prebuilt, runtime-dispatched OpenBLAS (the one numpy and scipy ship) for
    # Linux and Windows wheels. Its symbols are prefixed with scipy_ so it can't
    # clash with another OpenBLAS in the same process.
    find_package(Python COMPONENTS Interpreter REQUIRED)
    execute_process(
        COMMAND ${Python_EXECUTABLE} -c
            "import scipy_openblas32 as s; from importlib.metadata import files; print(s.get_include_dir()); print(s.get_lib_dir()); print(next((f.locate().as_posix() for f in files('scipy-openblas32') if f.name == 'LICENSE.txt'), ''))"
        OUTPUT_VARIABLE SOB_PATHS
        RESULT_VARIABLE SOB_RESULT
        OUTPUT_STRIP_TRAILING_WHITESPACE)
    if(NOT SOB_RESULT EQUAL 0)
        message(FATAL_ERROR "USE_SCIPY_OPENBLAS requires the scipy-openblas32 package (pip install scipy-openblas32)")
    endif()
    string(REGEX REPLACE "\r?\n" ";" SOB_PATHS "${SOB_PATHS}")
    list(GET SOB_PATHS 0 SOB_INCLUDE_DIR)
    list(GET SOB_PATHS 1 SOB_LIB_DIR)
    list(LENGTH SOB_PATHS SOB_PATH_COUNT)
    if(SOB_PATH_COUNT GREATER 2)
        list(GET SOB_PATHS 2 SOB_LICENSE)
    endif()
    find_library(SOB_LIBRARY NAMES scipy_openblas libscipy_openblas
                 PATHS ${SOB_LIB_DIR} NO_DEFAULT_PATH)
    if(NOT SOB_LIBRARY)
        message(FATAL_ERROR "scipy-openblas32 library not found in ${SOB_LIB_DIR}")
    endif()
    # libann calls the unprefixed names
    foreach(fn sgemm sgemv sger saxpy saxpby sscal sasum)
        target_compile_definitions(ann PRIVATE cblas_${fn}=scipy_cblas_${fn})
    endforeach()
    target_compile_definitions(ann PRIVATE USE_BLAS=1)
    target_include_directories(ann PRIVATE ${SOB_INCLUDE_DIR})
    target_link_libraries(ann PRIVATE ${SOB_LIBRARY})
    # Development builds load OpenBLAS from the package; wheels bundle a copy
    # next to libann instead (see the install rules below)
    if(NOT SKBUILD)
        set_property(TARGET ann APPEND PROPERTY INSTALL_RPATH "${SOB_LIB_DIR}")
    endif()
    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
        target_compile_definitions(ann PRIVATE _GNU_SOURCE)
    endif()
elseif(USE_BLAS)
    target_compile_definitions(ann PRIVATE USE_BLAS=1)
    
    if(WIN32)
        message(WARNING "OpenBLAS not supported on Windows, use USE_CBLAS instead")
    else()
        set(OPENBLAS_ROOT "/opt/OpenBLAS" CACHE PATH "OpenBLAS installation directory")
        if(EXISTS "${OPENBLAS_ROOT}/include/cblas.h")
            target_include_directories(ann PRIVATE ${OPENBLAS_ROOT}/include)
            target_link_directories(ann PRIVATE ${OPENBLAS_ROOT}/lib)
            target_link_libraries(ann PRIVATE openblas)
        else()
            # Try system OpenBLAS
            find_package(BLAS)
            if(BLAS_FOUND)
                target_link_libraries(ann PRIVATE ${BLAS_LIBRARIES})
            endif()
        endif()
    endif()
endif()

# Report the BLAS in use through pyann_blas_backend()
if(USE_CBLAS)
    set(PYANN_BLAS_NAME "cblas")
elseif(USE_ACCELERATE)
    set(PYANN_BLAS_NAME "accelerate")
elseif(USE_BLAS OR USE_SCIPY_OPENBLAS)
    set(PYANN_BLAS_NAME "openblas")
endif()
if(PYANN_BLAS_NAME)
    set_source_files_properties(src/native/pyann_info.c PROPERTIES
        COMPILE_DEFINITIONS "PYANN_BLAS_NAME=\"${PYANN_BLAS_NAME}\"")
endif()

# Hypertune library (optional)
if(EXISTS "${ANN_SOURCE_DIR}/ann_hypertune.c")
    target_sources(ann PRIVATE ${ANN_SOURCE_DIR}/ann_hypertune.c)
endif()

# Metal GPU support (macOS only)
if(USE_METAL)
    if(NOT APPLE)
        message(FATAL_ERROR "USE_METAL is only supported on macOS/Apple platforms")
    endif()
    target_compile_definitions(ann PRIVATE USE_METAL=1)
    target_sources(ann PRIVATE ${ANN_SOURCE_DIR}/tensor_metal.m)
    set_source_files_properties(${ANN_SOURCE_DIR}/tensor_metal.m PROPERTIES
        COMPILE_FLAGS "-x objective-c -fobjc-arc"
        LANGUAGE C
    )
    target_link_libraries(ann PRIVATE "-framework Metal" "-framework Foundation" "-framework MetalPerformanceShaders")
endif()

# CUDA GPU support
if(USE_CUDA)
    enable_language(CUDA)
    find_package(CUDAToolkit REQUIRED)
    target_compile_definitions(ann PRIVATE USE_CUDA=1)
    target_sources(ann PRIVATE ${ANN_SOURCE_DIR}/tensor_cuda.cu ${ANN_SOURCE_DIR}/shaders.cu)
    target_link_libraries(ann PRIVATE CUDA::cudart CUDA::cublas)
endif()

if(SKBUILD)
    # Building a wheel: install the library into the pyann package
    install(TARGETS ann
        LIBRARY DESTINATION pyann/lib
        RUNTIME DESTINATION pyann/lib
    )
    if(USE_SCIPY_OPENBLAS)
        # Ship OpenBLAS (and its runtime dependencies) next to libann so the
        # wheel works without scipy-openblas32 installed, including source
        # installs that never go through auditwheel or delvewheel.
        file(GLOB SOB_RUNTIME_LIBS "${SOB_LIB_DIR}/*.so*" "${SOB_LIB_DIR}/*.dll" "${SOB_LIB_DIR}/*.dylib")
        install(FILES ${SOB_RUNTIME_LIBS} DESTINATION pyann/lib)
        # OpenBLAS, LAPACK, libgfortran and libquadmath carry their own licenses
        if(SOB_LICENSE AND EXISTS "${SOB_LICENSE}")
            install(FILES "${SOB_LICENSE}" DESTINATION pyann/lib RENAME LICENSE-scipy-openblas32.txt)
        else()
            message(FATAL_ERROR "scipy-openblas32 license file not found; it must ship with the bundled libraries")
        endif()
    endif()
    set(PYANN_LIB_DIR "wheel: pyann/lib")
else()
    # Development build: copy the library next to the sources
    set(PYANN_LIB_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src/pyann/lib")
    file(MAKE_DIRECTORY ${PYANN_LIB_DIR})

    add_custom_command(TARGET ann POST_BUILD
        COMMAND ${CMAKE_COMMAND} -E copy $<TARGET_FILE:ann> ${PYANN_LIB_DIR}/
        COMMAND ${CMAKE_COMMAND} -E create_symlink $<TARGET_FILE_NAME:ann> ${PYANN_LIB_DIR}/libann${CMAKE_SHARED_LIBRARY_SUFFIX}
        COMMENT "Copying shared library to ${PYANN_LIB_DIR}"
    )

    install(TARGETS ann
        LIBRARY DESTINATION lib
        RUNTIME DESTINATION bin
        ARCHIVE DESTINATION lib
    )
endif()

# Print summary
message(STATUS "")
message(STATUS "pyann native library configuration:")
message(STATUS "  ANN source: ${ANN_SOURCE_DIR}")
message(STATUS "  BLAS: ${USE_BLAS}")
message(STATUS "  CBLAS: ${USE_CBLAS}")
message(STATUS "  Accelerate: ${USE_ACCELERATE}")
message(STATUS "  scipy-openblas32: ${USE_SCIPY_OPENBLAS}")
message(STATUS "  Metal: ${USE_METAL}")
message(STATUS "  CUDA: ${USE_CUDA}")
message(STATUS "  Output: ${PYANN_LIB_DIR}")
message(STATUS "")
