cmake_minimum_required(VERSION 3.10)
project(YOLO3D_TensorRT LANGUAGES CXX CUDA)

# Set C++ standard
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

# Set CUDA standard
set(CMAKE_CUDA_STANDARD 14)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)

# Find packages
find_package(CUDA REQUIRED)
find_package(OpenCV REQUIRED)

# Set TensorRT path (update this to your TensorRT installation)
set(TENSORRT_ROOT "/usr/local/TensorRT-8.6.1.6" CACHE PATH "Path to TensorRT installation")

# Find TensorRT
find_path(TENSORRT_INCLUDE_DIR NvInfer.h
    HINTS ${TENSORRT_ROOT}/include
    PATH_SUFFIXES include)

find_library(TENSORRT_LIBRARY nvinfer
    HINTS ${TENSORRT_ROOT}/lib
    PATH_SUFFIXES lib lib64)

find_library(TENSORRT_ONNX_PARSER nvonnxparser
    HINTS ${TENSORRT_ROOT}/lib
    PATH_SUFFIXES lib lib64)

# Check if TensorRT is found
if(NOT TENSORRT_INCLUDE_DIR OR NOT TENSORRT_LIBRARY)
    message(FATAL_ERROR "TensorRT not found. Please set TENSORRT_ROOT to your TensorRT installation path.")
endif()

# Optional: Find PCL for advanced point cloud handling
find_package(PCL QUIET)
if(PCL_FOUND)
    add_definitions(-DUSE_PCL)
    include_directories(${PCL_INCLUDE_DIRS})
    link_directories(${PCL_LIBRARY_DIRS})
    add_definitions(${PCL_DEFINITIONS})
endif()

# Optional: Find ROS for ROS integration
find_package(catkin QUIET COMPONENTS
    roscpp
    sensor_msgs
    visualization_msgs
    pcl_conversions
    pcl_ros
)

if(catkin_FOUND)
    add_definitions(-DUSE_ROS)
    catkin_package(
        INCLUDE_DIRS include
        LIBRARIES yolo3d_tensorrt
        CATKIN_DEPENDS roscpp sensor_msgs visualization_msgs pcl_conversions pcl_ros
    )
endif()

# Include directories
include_directories(
    ${CMAKE_CURRENT_SOURCE_DIR}/include
    ${CUDA_INCLUDE_DIRS}
    ${OpenCV_INCLUDE_DIRS}
    ${TENSORRT_INCLUDE_DIR}
)

if(catkin_FOUND)
    include_directories(${catkin_INCLUDE_DIRS})
endif()

# CUDA compilation flags
set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};
    -O3
    -gencode arch=compute_75,code=sm_75  # For RTX 2080
    -gencode arch=compute_86,code=sm_86  # For RTX 3090
    -gencode arch=compute_89,code=sm_89  # For RTX 4090
)

# Source files
set(SOURCES
    src/yolo3d_detector.cpp
    src/bev_converter.cu
    src/tensorrt_engine.cpp
    src/nms3d.cpp
    src/visualizer.cpp
    src/point_cloud_loader.cpp
)

# Create library
add_library(yolo3d_tensorrt SHARED ${SOURCES})

# Link libraries
target_link_libraries(yolo3d_tensorrt
    ${CUDA_LIBRARIES}
    ${OpenCV_LIBS}
    ${TENSORRT_LIBRARY}
    ${TENSORRT_ONNX_PARSER}
    cuda
    cudart
    cublas
    cudnn
)

if(PCL_FOUND)
    target_link_libraries(yolo3d_tensorrt ${PCL_LIBRARIES})
endif()

if(catkin_FOUND)
    target_link_libraries(yolo3d_tensorrt ${catkin_LIBRARIES})
endif()

# Create executable
add_executable(yolo3d_detect src/main.cpp)
target_link_libraries(yolo3d_detect yolo3d_tensorrt)

# Additional executables
add_executable(yolo3d_benchmark src/benchmark.cpp)
target_link_libraries(yolo3d_benchmark yolo3d_tensorrt)

add_executable(onnx_to_tensorrt src/onnx_to_tensorrt.cpp)
target_link_libraries(onnx_to_tensorrt yolo3d_tensorrt)

# Install targets
install(TARGETS yolo3d_tensorrt yolo3d_detect
    RUNTIME DESTINATION bin
    LIBRARY DESTINATION lib
    ARCHIVE DESTINATION lib
)

install(DIRECTORY include/
    DESTINATION include
)

# Create launch file for ROS
if(catkin_FOUND)
    install(DIRECTORY launch/
        DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION}/launch
    )
endif()

# Add tests
enable_testing()
add_subdirectory(test)

# Print configuration
message(STATUS "YOLO3D TensorRT Configuration:")
message(STATUS "  CUDA Version: ${CUDA_VERSION_STRING}")
message(STATUS "  TensorRT Include: ${TENSORRT_INCLUDE_DIR}")
message(STATUS "  TensorRT Library: ${TENSORRT_LIBRARY}")
message(STATUS "  OpenCV Version: ${OpenCV_VERSION}")
message(STATUS "  PCL Found: ${PCL_FOUND}")
message(STATUS "  ROS Found: ${catkin_FOUND}")
message(STATUS "  Build Type: ${CMAKE_BUILD_TYPE}")