# ops/setup.py
import os
import torch
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension

def make_cuda_ext(name, module, sources, sources_cuda=[], extra_args=[]):
    define_macros = []
    extra_compile_args = {"cxx": [] + extra_args}

    if torch.cuda.is_available() or os.getenv("FORCE_CUDA", "0") == "1":
        define_macros += [("WITH_CUDA", None)]
        extension = CUDAExtension
        extra_compile_args["nvcc"] = extra_args + [
            "-D__CUDA_NO_HALF_OPERATORS__",
            "-D__CUDA_NO_HALF_CONVERSIONS__",
            "-D__CUDA_NO_HALF2_OPERATORS__",
        ]
        sources += sources_cuda
    else:
        print(f"Compiling {name} without CUDA")
        extension = CppExtension

    return extension(
        name=f"{module}.{name}",
        sources=[os.path.join(module.replace(".", "/"), "src", p) for p in sources],
        include_dirs=[os.path.join(module.replace(".", "/"), "src")],
        define_macros=define_macros,
        extra_compile_args=extra_compile_args,
    )

setup(
    name="deformable_aggregation_ext",
    packages=["deformable_aggregation_ext"],   # 指定包名
    ext_modules=[
        make_cuda_ext(
            name="deformable_aggregation_ext",
            module="deformable_aggregation_ext",  # 指向包名
            sources=[
                "deformable_aggregation.cpp",
                "deformable_aggregation_cuda.cu"
            ],
        )
    ],
    cmdclass={"build_ext": BuildExtension},
)