# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
# except for the third-party components listed below.
# Hunyuan 3D does not impose any additional limitations beyond what is outlined
# in the repsective licenses of these third-party components.
# Users must comply with all terms and conditions of original licenses of these third-party
# components and must ensure that the usage of the third party components adheres to
# all relevant laws and regulations.

# For avoidance of doubts, Hunyuan 3D means the large language models and
# their software and algorithms, including trained model weights, parameters (including
# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
# fine-tuning enabling code and other elements of the foregoing made publicly available
# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.

from setuptools import setup, find_packages
import os
import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension

# build custom rasterizer

# CUDA include path: prefer conda env CUDA headers to match torch's CUDA version
_cuda_home = os.environ.get("CUDA_HOME", os.environ.get("CUDA_PATH", "/usr/local/cuda"))
_cuda_include = os.path.join(_cuda_home, "targets", "x86_64-linux", "include")
if not os.path.isdir(_cuda_include):
    _cuda_include = os.path.join(_cuda_home, "include")

custom_rasterizer_module = CUDAExtension(
    "custom_rasterizer_kernel",
    [
        "lib/custom_rasterizer_kernel/rasterizer.cpp",
        "lib/custom_rasterizer_kernel/grid_neighbor.cpp",
        "lib/custom_rasterizer_kernel/rasterizer_gpu.cu",
    ],
    include_dirs=[_cuda_include],
    # -D__GLIBC_USE_IEC_60559_FUNCS_EXT_C23=0 prevents glibc 2.38+ from declaring
    # sinpi/cospi/etc that conflict with CUDA 12.8 crt/math_functions.h on modern glibc.
    extra_compile_args={
        "nvcc": [],
        "cxx": [],
    },
)

setup(
    packages=find_packages(),
    version="0.1",
    name="custom_rasterizer",
    include_package_data=True,
    package_dir={"": "."},
    ext_modules=[
        custom_rasterizer_module,
    ],
    cmdclass={"build_ext": BuildExtension},
)