cmake_minimum_required(VERSION 2.8.12)
project(greedrl_C_ LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 14)

find_package(PythonInterp REQUIRED)
execute_process(COMMAND "python" "-c" 
"
import os
import torch
from distutils import sysconfig as s
print(s.get_python_inc(plat_specific=True))
print(s.get_config_var('EXT_SUFFIX'))
print(os.path.dirname(torch.__file__))
"
RESULT_VARIABLE _PYTHON_SUCCESS
OUTPUT_VARIABLE _PYTHON_VALUES
ERROR_VARIABLE _PYTHON_ERROR_VALUE)

if(NOT _PYTHON_SUCCESS MATCHES 0)
    message("_PYTHON_SUCCESS: ${_PYTHON_SUCCESS}")
    message("_PYTHON_VALUES: ${_PYTHON_VALUES}")
    message("_PYTHON_ERROR_VALUE: ${_PYTHON_ERROR_VALUE}")
    message(FATAL_ERROR "get python config error!")
endif()

string(REGEX REPLACE "\n" ";" _PYTHON_VALUES ${_PYTHON_VALUES})
list(GET _PYTHON_VALUES 0 PYTHON_INCLUDE_DIR)
list(GET _PYTHON_VALUES 1 PYTHON_EXT_SUFFIX)
list(GET _PYTHON_VALUES 2 TORCH_HOME)

include_directories(
    ${PYTHON_INCLUDE_DIR}
    ${TORCH_HOME}/include
    ${TORCH_HOME}/include/TH
    ${TORCH_HOME}/include/THC
    ${TORCH_HOME}/include/torch/csrc/api/include
)

string(LENGTH "${CMAKE_SOURCE_DIR}/" SOURCE_PATH_LENGTH)
add_compile_options(-DSOURCE_PATH_LENGTH=${SOURCE_PATH_LENGTH})
add_compile_options(-D_GLIBCXX_USE_CXX11_ABI=0 -fvisibility=hidden -fopenmp)

if(${CMAKE_CXX_COMPILER_ID} STREQUAL "Clang")
    add_link_options(-undefined dynamic_lookup)
endif()

file(GLOB_RECURSE CSRC_CPP csrc/*.cpp)

add_library(greedrl_c MODULE ${CSRC_CPP})
set_target_properties(greedrl_c PROPERTIES PREFIX "")
set_target_properties(greedrl_c PROPERTIES SUFFIX "${PYTHON_EXT_SUFFIX}")
target_compile_options(greedrl_c PRIVATE -Wno-sign-conversion -O3)
target_link_libraries(greedrl_c c10 torch torch_cpu torch_python)
target_link_directories(greedrl_c PRIVATE ${TORCH_HOME}/lib)

find_package(CUDA)
if(CUDA_FOUND)
    enable_language(CUDA)
    file(GLOB_RECURSE CSRC_CU csrc/*.cu)
    add_library(greedrl_cu OBJECT ${CSRC_CU})
    target_compile_options(greedrl_cu PRIVATE -keep -Xptxas -v --expt-relaxed-constexpr --expt-extended-lambda -O3)
    set_target_properties(greedrl_cu PROPERTIES POSITION_INDEPENDENT_CODE ON CUDA_ARCHITECTURES "70;75;80")
    add_compile_definitions(CUDA_FOUND)
    include_directories(${CUDA_INCLUDE_DIRS})
    target_link_libraries(greedrl_c torch_cuda greedrl_cu)
    target_link_directories(greedrl_c PRIVATE ${TORCH_HOME}/lib)
endif()