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

# cpp_dir = os.path.join('src', 'torch_dwn', 'custom_operators', 'cpp')
cuda_dir = os.path.join('src', 'torch_dwn', 'custom_operators', 'cuda')

ext_modules = []#CppExtension(filename[:-3], [os.path.join(cpp_dir, filename)]) for filename in os.listdir(cpp_dir)]
if torch.cuda.is_available():
    for filename in os.listdir(cuda_dir):
        if filename[-3:] == 'cpp':
            module_name = filename[:-4]
            kernel_filename = module_name + '_kernel.cu'
            ext_modules.append(CUDAExtension(module_name, [os.path.join(cuda_dir, filename), os.path.join(cuda_dir, kernel_filename)]))
setup(
    name='torch_dwn',
    ext_modules=ext_modules,
    cmdclass={'build_ext': BuildExtension},
    package_dir={"": "src"},
    packages=setuptools.find_packages(where="src"),
    version="1.0.3",
    author="Alan T. L. Bacellar",
    author_email="alanbacellar@gmail.com",
    description="Differentiable Weightless Neural Networks (DWN) pytorch module",
    url="https://github.com/alanbacellar/DWN",
    install_requires=[
        'torch'
    ]
)