cmake_minimum_required(VERSION 3.22)
project(pydftb_torch LANGUAGES CXX)

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

option(DFTB_TORCH_FORCE_PERFORMANCE_FLAGS "Force high-performance compiler flags for local validation builds" ON)
option(DFTB_TORCH_ENABLE_IPO "Enable interprocedural optimization when the compiler supports it" ON)
option(DFTB_TORCH_ENABLE_NATIVE_OPT "Enable CPU-native code generation for non-portable local builds" OFF)
option(DFTB_TORCH_ENABLE_FAST_MATH "Enable unsafe fast-math flags; disabled by default for numerical reproducibility" OFF)

option(DFTB_TORCH_USE_PYTHON_TORCH "Use Python torch package to locate LibTorch" ON)
if(NOT DEFINED DFTB_TORCH_PYBIND11_EXTRAS)
  if(DFTB_TORCH_FORCE_PERFORMANCE_FLAGS OR CMAKE_BUILD_TYPE MATCHES "^(Release|RelWithDebInfo|MinSizeRel)$")
    set(DFTB_TORCH_PYBIND11_EXTRAS ON CACHE BOOL "Enable pybind11 LTO/strip extras for the extension module" FORCE)
  else()
    set(DFTB_TORCH_PYBIND11_EXTRAS OFF CACHE BOOL "Enable pybind11 LTO/strip extras for the extension module" FORCE)
  endif()
endif()

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

find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)

execute_process(
  COMMAND "${Python_EXECUTABLE}" -c "import pybind11, sys; print(pybind11.get_cmake_dir())"
  OUTPUT_VARIABLE PYBIND11_PREFIX_FROM_PYTHON
  OUTPUT_STRIP_TRAILING_WHITESPACE
  ERROR_QUIET
  TIMEOUT 20
)
if(PYBIND11_PREFIX_FROM_PYTHON)
  list(PREPEND CMAKE_PREFIX_PATH "${PYBIND11_PREFIX_FROM_PYTHON}")
  # Also seed the concrete package cache variable.  Some frontends inherit a
  # broad CMAKE_PREFIX_PATH but still spend a long time in repeated package
  # discovery when the src/cpp subproject calls find_package(pybind11).
  set(pybind11_DIR "${PYBIND11_PREFIX_FROM_PYTHON}" CACHE PATH "pybind11 CMake package directory" FORCE)
endif()


if(DFTB_TORCH_USE_PYTHON_TORCH)
  execute_process(
    COMMAND "${Python_EXECUTABLE}" -c "import sys, torch, pathlib; print(sys.prefix); print(torch.utils.cmake_prefix_path); print(pathlib.Path(torch.__file__).parent / 'lib')"
    OUTPUT_VARIABLE _PY_TORCH_INFO
    OUTPUT_STRIP_TRAILING_WHITESPACE
    ERROR_QUIET
    TIMEOUT 20
  )
  string(REPLACE "\n" ";" _PY_TORCH_INFO_LIST "${_PY_TORCH_INFO}")
  list(LENGTH _PY_TORCH_INFO_LIST _PY_TORCH_INFO_LEN)
  if(_PY_TORCH_INFO_LEN GREATER_EQUAL 3)
    list(GET _PY_TORCH_INFO_LIST 0 DFTB_TORCH_PYTHON_PREFIX)
    list(GET _PY_TORCH_INFO_LIST 1 TORCH_PREFIX_FROM_PYTHON)
    list(GET _PY_TORCH_INFO_LIST 2 DFTB_TORCH_PYTHON_LIBDIR)
    list(PREPEND CMAKE_PREFIX_PATH "${DFTB_TORCH_PYTHON_PREFIX}" "${TORCH_PREFIX_FROM_PYTHON}")
    set(Torch_DIR "${TORCH_PREFIX_FROM_PYTHON}/Torch" CACHE PATH "Torch CMake package directory" FORCE)
    set(DFTB_TORCH_PYTHON_LIBDIR "${DFTB_TORCH_PYTHON_LIBDIR}" CACHE PATH "Python torch library directory" FORCE)
    if(DEFINED ENV{LIBTORCH_HOME})
      list(APPEND CMAKE_IGNORE_PREFIX_PATH "$ENV{LIBTORCH_HOME}")
    endif()
    message(STATUS "Using Python torch CMake package: ${Torch_DIR}")
  else()
    message(FATAL_ERROR "DFTB_TORCH_USE_PYTHON_TORCH=ON but Python torch could not be queried")
  endif()
elseif(DEFINED ENV{LIBTORCH_HOME})
  list(PREPEND CMAKE_PREFIX_PATH "$ENV{LIBTORCH_HOME}")
endif()

add_subdirectory(src/cpp)
