# CMakeLists.txt for the control module

# Find required dependencies
find_package(Eigen3 REQUIRED)
find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)

find_package(pybind11 REQUIRED)
set(MAP_LOADER_DEPS)
# Keep artifacts in the active build tree so different entrypoints do not
# overwrite each other's outputs in the source tree.
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/lib)
set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/lib)
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/lib)

if(Torch_FOUND)
    list(APPEND MAP_LOADER_DEPS "${TORCH_LIBRARIES}")
endif()

if(Torch_FOUND)
  set(MPLOAD_SHARED_RPATH "/usr/local/lib")
  if(DEFINED STANDARD_BOTS_LIBTORCH_ROOT)
    if(EXISTS "${STANDARD_BOTS_LIBTORCH_ROOT}/lib")
      list(APPEND MPLOAD_SHARED_RPATH "${STANDARD_BOTS_LIBTORCH_ROOT}/lib")
    endif()
  endif()
  if(DEFINED Torch_INSTALL_PREFIX AND Torch_INSTALL_PREFIX)
    if(EXISTS "${Torch_INSTALL_PREFIX}/lib")
      list(APPEND MPLOAD_SHARED_RPATH "${Torch_INSTALL_PREFIX}/lib")
    endif()
  endif()

  if(Python3_FOUND AND Python3_EXECUTABLE)
    execute_process(
      COMMAND ${Python3_EXECUTABLE} -c "import torch, pathlib; print(pathlib.Path(torch.__file__).resolve().parent)"
      OUTPUT_VARIABLE _torch_site_dir
      OUTPUT_STRIP_TRAILING_WHITESPACE
      ERROR_QUIET)
    if(_torch_site_dir)
      get_filename_component(_torch_site_parent "${_torch_site_dir}/.." REALPATH)
      if(EXISTS "${_torch_site_parent}/torch.libs")
        list(APPEND MPLOAD_SHARED_RPATH "${_torch_site_parent}/torch.libs")
      endif()
    endif()
  endif()

  file(GLOB _torch_lib_folders
       "${PROJECT_SOURCE_DIR}/venv/lib/python*/site-packages/torch.libs"
       "${PROJECT_SOURCE_DIR}/../../venv/lib/python*/site-packages/torch.libs")
  foreach(_torch_lib_dir IN LISTS _torch_lib_folders)
    if(IS_DIRECTORY "${_torch_lib_dir}")
      list(APPEND MPLOAD_SHARED_RPATH "${_torch_lib_dir}")
    endif()
  endforeach()

  list(REMOVE_DUPLICATES MPLOAD_SHARED_RPATH)
  string(REPLACE ";" ":" MPLOAD_SHARED_RPATH_STRING "${MPLOAD_SHARED_RPATH}")
  message(STATUS "Map loader module RPATH: ${MPLOAD_SHARED_RPATH_STRING}")
endif()

# Torch-free core used by the base_shaper extension.
add_library(base_shaper_core
    include/base_shaper.cpp
)

set_target_properties(base_shaper_core PROPERTIES
    POSITION_INDEPENDENT_CODE ON
)

target_include_directories(base_shaper_core
    PUBLIC
        ${CMAKE_CURRENT_SOURCE_DIR}
        ${EIGEN3_INCLUDE_DIR}
)

target_compile_features(base_shaper_core PRIVATE cxx_std_17)

target_link_libraries(base_shaper_core
    PUBLIC
        pybind11::headers
)

if(Torch_FOUND)
  # Support both packaged SDK layout and monorepo editable layout.
  set(_ROBOT_DYN_INCLUDE_CANDIDATES
      "${CMAKE_CURRENT_LIST_DIR}/../util/native/include"
      "${CMAKE_CURRENT_LIST_DIR}/../../../../src/util/native/include")
  set(ROBOT_DYNAMICS_INCLUDE_DIR "")
  foreach(_cand IN LISTS _ROBOT_DYN_INCLUDE_CANDIDATES)
    get_filename_component(_cand_real "${_cand}" REALPATH)
    if(EXISTS "${_cand_real}/robot_dynamics.hpp")
      set(ROBOT_DYNAMICS_INCLUDE_DIR "${_cand_real}")
      break()
    endif()
  endforeach()
  if(NOT ROBOT_DYNAMICS_INCLUDE_DIR)
    message(FATAL_ERROR "Could not locate robot_dynamics.hpp in packaged or monorepo util/native/include")
  endif()
  if(NOT TARGET robot_dynamics_core)
    add_library(robot_dynamics_core STATIC
        "${ROBOT_DYNAMICS_INCLUDE_DIR}/robot_dynamics.cpp"
    )
    set_target_properties(robot_dynamics_core PROPERTIES
        POSITION_INDEPENDENT_CODE ON
    )
    target_include_directories(robot_dynamics_core
        PUBLIC
            "${ROBOT_DYNAMICS_INCLUDE_DIR}"
            ${EIGEN3_INCLUDE_DIR}
    )
    target_link_libraries(robot_dynamics_core
        PUBLIC
            Eigen3::Eigen
    )
    target_compile_features(robot_dynamics_core PUBLIC cxx_std_17)
  endif()

  # Torch-enabled core used only by the covalent extension.
  add_library(control_core
      include/covalent_wrapper.cpp
      include/map_loader.cpp
  )

  set_target_properties(control_core PROPERTIES
      POSITION_INDEPENDENT_CODE ON
  )

  target_include_directories(control_core
      PUBLIC
          ${CMAKE_CURRENT_SOURCE_DIR}
          ${EIGEN3_INCLUDE_DIR}
          ${TORCH_INCLUDE_DIRS}
          ${ROBOT_DYNAMICS_INCLUDE_DIR}
      PRIVATE
          ${Python3_INCLUDE_DIRS}
  )

  target_compile_features(control_core PRIVATE cxx_std_17)

  target_link_libraries(control_core
      PUBLIC
          pybind11::headers
          robot_dynamics_core
          ${MAP_LOADER_DEPS}
  )
endif()

# Create Python module
pybind11_add_module(base_shaper
    include/base_shaper_bindings.cpp
)

target_link_libraries(base_shaper
    PRIVATE
        base_shaper_core
)

# Set include directories for Python module
target_include_directories(base_shaper
    PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}
        ${EIGEN3_INCLUDE_DIR}
)

# Set compile options for Python module
target_compile_features(base_shaper PRIVATE cxx_std_17)

if(Torch_FOUND)
  pybind11_add_module(covalent
      include/covalent_wrapper_binding.cpp
  )

  target_link_libraries(covalent
      PRIVATE
          control_core
  )

  target_include_directories(covalent
      PRIVATE
          ${CMAKE_CURRENT_SOURCE_DIR}
          ${EIGEN3_INCLUDE_DIR}
  )

  target_compile_features(covalent PRIVATE cxx_std_17)
endif()

install(TARGETS base_shaper
    LIBRARY DESTINATION reforge_core/control
    RUNTIME DESTINATION reforge_core/control
)

if(TARGET covalent)
  install(TARGETS covalent
      LIBRARY DESTINATION reforge_core/control
      RUNTIME DESTINATION reforge_core/control
  )
endif()

install(FILES
    include/base_shaper.hpp
    include/covalent_wrapper.cpp
    include/map_loader.cpp
    DESTINATION reforge_core/control/include
)
