# SPDX-License-Identifier: GPL-3.0-or-later
# PyGFN0Torch - LibTorch-native GFN0-xTB parameter laboratory
# Author: ss0832
#
# This top-level CMake entry point supports two Torch providers:
#
#   1. Python torch wheel (default for pip/scikit-build-core builds)
#      The extension links against the C++ libraries shipped with the active
#      Python torch package. This avoids ABI/version skew between Python torch
#      and an unrelated external LibTorch ZIP.
#
#   2. Explicit LibTorch ZIP (for non-Python C++ development)
#      Select with -DGFN0_TORCH_PROVIDER=libtorch and set LIBTORCH_HOME.
#
# For Python extension builds, do not silently prefer LIBTORCH_HOME just because
# it is present in the environment. Mixing Torch_DIR from the Python wheel with
# libtorch.so from an external ZIP can produce import-time undefined symbols.
cmake_minimum_required(VERSION 3.20)
project(pygfn0_torch
    VERSION 0.1.0
    DESCRIPTION "Differentiable GFN0-xTB Hamiltonian on LibTorch"
    LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)

if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
    set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
endif()

option(GFN0_BUILD_PYBIND "Build the pybind11 Python extension module" ON)
option(GFN0_BUILD_CLI    "Build the gfn0-cli command-line tool"       ON)
option(GFN0_BUILD_TESTS  "Build native C++ test executables"          OFF)
option(GFN0_BUILD_EXAMPLES "Build native C++ example executables"     OFF)

set(GFN0_TORCH_PROVIDER "python" CACHE STRING
    "Torch provider: 'python' uses the active Python torch wheel; 'libtorch' uses LIBTORCH_HOME.")
set_property(CACHE GFN0_TORCH_PROVIDER PROPERTY STRINGS python libtorch)

find_package(Python3 COMPONENTS Interpreter REQUIRED)

set(_gfn0_torch_lib_dir "")

if(GFN0_TORCH_PROVIDER STREQUAL "libtorch")
    if(NOT DEFINED ENV{LIBTORCH_HOME})
        message(FATAL_ERROR
            "GFN0_TORCH_PROVIDER=libtorch requires LIBTORCH_HOME to point to a valid LibTorch installation")
    endif()
    if(NOT EXISTS "$ENV{LIBTORCH_HOME}/share/cmake/Torch/TorchConfig.cmake")
        message(FATAL_ERROR
            "LIBTORCH_HOME is set to '$ENV{LIBTORCH_HOME}', but TorchConfig.cmake was not found under "
            "$ENV{LIBTORCH_HOME}/share/cmake/Torch")
    endif()
    set(Torch_DIR "$ENV{LIBTORCH_HOME}/share/cmake/Torch" CACHE PATH "Path to TorchConfig.cmake" FORCE)
    set(_gfn0_torch_lib_dir "$ENV{LIBTORCH_HOME}/lib")
    message(STATUS "Using external LibTorch provider: ${Torch_DIR}")
elseif(GFN0_TORCH_PROVIDER STREQUAL "python")
    execute_process(
        COMMAND "${Python3_EXECUTABLE}" -c
            "import os, torch; print(torch.utils.cmake_prefix_path); print(os.path.join(os.path.dirname(torch.__file__), 'lib'))"
        RESULT_VARIABLE _gfn0_torch_probe_rc
        OUTPUT_VARIABLE _gfn0_torch_probe_out
        ERROR_VARIABLE _gfn0_torch_probe_err
        OUTPUT_STRIP_TRAILING_WHITESPACE)
    if(NOT _gfn0_torch_probe_rc EQUAL 0)
        message(FATAL_ERROR
            "Could not import the active Python torch package while locating Torch.\n"
            "Python: ${Python3_EXECUTABLE}\n"
            "Error: ${_gfn0_torch_probe_err}\n"
            "Install torch in this environment, or use -DGFN0_TORCH_PROVIDER=libtorch with LIBTORCH_HOME.")
    endif()
    string(REPLACE "\n" ";" _gfn0_torch_probe_lines "${_gfn0_torch_probe_out}")
    list(GET _gfn0_torch_probe_lines 0 _gfn0_torch_cmake_prefix)
    list(GET _gfn0_torch_probe_lines 1 _gfn0_torch_lib_dir)

    list(PREPEND CMAKE_PREFIX_PATH "${_gfn0_torch_cmake_prefix}")
    if(EXISTS "${_gfn0_torch_cmake_prefix}/Torch/TorchConfig.cmake")
        set(Torch_DIR "${_gfn0_torch_cmake_prefix}/Torch" CACHE PATH "Path to TorchConfig.cmake" FORCE)
    elseif(EXISTS "${_gfn0_torch_cmake_prefix}/share/cmake/Torch/TorchConfig.cmake")
        set(Torch_DIR "${_gfn0_torch_cmake_prefix}/share/cmake/Torch" CACHE PATH "Path to TorchConfig.cmake" FORCE)
    endif()
    message(STATUS "Using Python torch provider")
    message(STATUS "  Python executable : ${Python3_EXECUTABLE}")
    message(STATUS "  torch cmake prefix: ${_gfn0_torch_cmake_prefix}")
    message(STATUS "  torch lib dir     : ${_gfn0_torch_lib_dir}")
else()
    message(FATAL_ERROR "GFN0_TORCH_PROVIDER must be 'python' or 'libtorch', got '${GFN0_TORCH_PROVIDER}'")
endif()

option(GFN0_FORCE_CPU_TORCH
    "Patch TorchConfig at configure-time to skip CUDA library lookup."
    OFF)

if(NOT GFN0_FORCE_CPU_TORCH)
    execute_process(
        COMMAND "${Python3_EXECUTABLE}" -c
            "import glob, importlib.util, os; sp=importlib.util.find_spec('torch'); base=os.path.dirname(os.path.dirname(sp.origin)) if sp else ''; nv=os.path.join(base, 'nvidia'); print(os.linesep.join([os.path.join(d, 'lib') for d in glob.glob(os.path.join(nv, '*')) if os.path.isdir(os.path.join(d, 'lib'))]))"
        OUTPUT_VARIABLE _gfn0_nvidia_libs
        OUTPUT_STRIP_TRAILING_WHITESPACE)
    if(_gfn0_nvidia_libs)
        string(REPLACE "\n" ";" _gfn0_nvidia_libs_list "${_gfn0_nvidia_libs}")
        list(APPEND CMAKE_LIBRARY_PATH ${_gfn0_nvidia_libs_list})
        message(STATUS "Adding NVIDIA wheel libs to search path: ${_gfn0_nvidia_libs_list}")
    endif()
endif()

find_package(Torch REQUIRED)

if(UNIX AND NOT APPLE)
    set(_gfn0_origin_rpath "$ORIGIN/../torch/lib")
elseif(APPLE)
    set(_gfn0_origin_rpath "@loader_path/../torch/lib")
else()
    set(_gfn0_origin_rpath "")
endif()

add_subdirectory(src/cpp)

if(GFN0_BUILD_TESTS)
    enable_testing()
    add_subdirectory(tests/cpp)
endif()

message(STATUS "")
message(STATUS "PyGFN0Torch build configuration:")
message(STATUS "  Build type      : ${CMAKE_BUILD_TYPE}")
message(STATUS "  Compiler        : ${CMAKE_CXX_COMPILER_ID} ${CMAKE_CXX_COMPILER_VERSION}")
message(STATUS "  C++ standard    : C++${CMAKE_CXX_STANDARD}")
message(STATUS "  Torch provider  : ${GFN0_TORCH_PROVIDER}")
message(STATUS "  Torch_DIR       : ${Torch_DIR}")
message(STATUS "  Torch lib dir   : ${_gfn0_torch_lib_dir}")
message(STATUS "  Extension RPATH : ${_gfn0_origin_rpath};${_gfn0_torch_lib_dir}")
message(STATUS "  Pybind module   : ${GFN0_BUILD_PYBIND}")
message(STATUS "  CLI             : ${GFN0_BUILD_CLI}")
message(STATUS "  Native tests    : ${GFN0_BUILD_TESTS}")
message(STATUS "  Native examples : ${GFN0_BUILD_EXAMPLES}")
message(STATUS "")
