# FindTorch
# -------
#
# Finds the Torch library
#
# This will define the following variables:
#
#   TORCH_FOUND        -- True if the system has the Torch library
#   TORCH_INCLUDE_DIRS -- The include directories for torch
#   TORCH_LIBRARIES    -- Libraries to link against
#   TORCH_CXX_FLAGS    -- Additional (required) compiler flags
#
# and the following imported targets:
#
#   torch

if ($ENV{TORCH_INSTALL_PREFIX})
  set(TORCH_INSTALL_PREFIX $ENV{TORCH_INSTALL_PREFIX})
else()
  # Assume we are in <install-prefix>/share/cmake/Torch/TorchConfig.cmake
  get_filename_component(CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH)
  get_filename_component(TORCH_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE)
endif()

# Include directories.
if (EXISTS "${TORCH_INSTALL_PREFIX}/lib/include")
  set(TORCH_INCLUDE_DIRS
    ${TORCH_INSTALL_PREFIX}/lib/include
    ${TORCH_INSTALL_PREFIX}/lib/include/torch/csrc/api/include)
else()
  set(TORCH_INCLUDE_DIRS
    ${TORCH_INSTALL_PREFIX}/include
    ${TORCH_INSTALL_PREFIX}/include/torch/csrc/api/include)
endif()

# Library dependencies.
find_package(Caffe2 REQUIRED)

find_library(TORCH_LIBRARY torch PATHS "${TORCH_INSTALL_PREFIX}/lib")
add_library(torch SHARED IMPORTED)
set(TORCH_LIBRARIES torch ${Caffe2_MAIN_LIBS})

if (0)
  if(MSVC)
    set(NVTOOLEXT_HOME "C:/Program Files/NVIDIA Corporation/NvToolsExt")
    if ($ENV{NVTOOLEXT_HOME})
      set(NVTOOLEXT_HOME $ENV{NVTOOLEXT_HOME})
    endif()
    set(TORCH_CUDA_LIBRARIES
      ${NVTOOLEXT_HOME}/lib/x64/nvToolsExt64_1.lib
      ${CUDA_LIBRARIES})
    list(APPEND TORCH_INCLUDE_DIRS ${NVTOOLEXT_HOME}/include)
  elseif(APPLE)
    set(TORCH_CUDA_LIBRARIES
      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libcudart.dylib
      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libnvrtc.dylib
      ${CUDA_TOOLKIT_ROOT_DIR}/lib/libnvToolsExt.dylib
      ${CUDA_LIBRARIES})
  else()
    set(TORCH_CUDA_LIBRARIES
      ${CUDA_CUDA_LIB}
      ${CUDA_NVRTC_LIB}
      ${CUDA_TOOLKIT_ROOT_DIR}/lib64/libnvToolsExt.so
      ${CUDA_LIBRARIES})
  endif()
  list(APPEND TORCH_LIBRARIES ${TORCH_CUDA_LIBRARIES})
endif()

# When we build libtorch with the old GCC ABI, dependent libraries must too.
set(TORCH_CXX_FLAGS "-D_GLIBCXX_USE_CXX11_ABI=")

set_target_properties(torch PROPERTIES
    IMPORTED_LOCATION "${TORCH_LIBRARY}"
    INTERFACE_INCLUDE_DIRECTORIES "${TORCH_INCLUDE_DIRS}"
    INTERFACE_COMPILE_OPTIONS "${TORCH_CXX_FLAGS}"
    CXX_STANDARD 11
)