from __future__ import print_function
from setuptools import setup, find_packages
import os
import shutil
import platform

# make the faiss python package dir
shutil.rmtree("faiss", ignore_errors=True)
os.mkdir("faiss")
shutil.copytree("contrib", "faiss/contrib")
shutil.copyfile("__init__.py", "faiss/__init__.py")
shutil.copyfile("loader.py", "faiss/loader.py")
shutil.copyfile("libfaiss_python_callbacks.so", "faiss/libfaiss_python_callbacks.so")
shutil.copyfile("../libfaiss.so", "faiss/libfaiss.so")

# 复制所有的 .so 库到 faiss 目录
intel_mkl_libs_dir = "/opt/intel/compilers_and_libraries_2020.4.304/linux/mkl/lib/intel64_lin"
for file in os.listdir(intel_mkl_libs_dir):
    if file.endswith(".so"):
        shutil.copyfile(os.path.join(intel_mkl_libs_dir, file), os.path.join("faiss", file))

ext = ".pyd" if platform.system() == 'Windows' else ".so"
prefix = "Release/" * (platform.system() == 'Windows')

swigfaiss_generic_lib = f"{prefix}_swigfaiss{ext}"
swigfaiss_avx2_lib = f"{prefix}_swigfaiss_avx2{ext}"

found_swigfaiss_generic = os.path.exists(swigfaiss_generic_lib)
found_swigfaiss_avx2 = os.path.exists(swigfaiss_avx2_lib)

assert (found_swigfaiss_generic or found_swigfaiss_avx2), \
    f"Could not find {swigfaiss_generic_lib} or " \
    f"{swigfaiss_avx2_lib}. Faiss may not be compiled yet."

if found_swigfaiss_generic:
    print(f"Copying {swigfaiss_generic_lib}")
    shutil.copyfile("swigfaiss.py", "faiss/swigfaiss.py")
    shutil.copyfile(swigfaiss_generic_lib, f"faiss/_swigfaiss{ext}")

if found_swigfaiss_avx2:
    print(f"Copying {swigfaiss_avx2_lib}")
    shutil.copyfile("swigfaiss_avx2.py", "faiss/swigfaiss_avx2.py")
    shutil.copyfile(swigfaiss_avx2_lib, f"faiss/_swigfaiss_avx2{ext}")

ROCM_PATH = os.getenv('ROCM_PATH')
dtk_path = ROCM_PATH + '/.info/rocm_version'
with open(dtk_path, 'r') as file:
    content = file.read().strip()
dtk_version = content.replace('.', '')
print(dtk_version)
cwd = os.path.dirname(os.path.abspath(__file__))
ver_path = os.path.join(cwd, "faiss", "__init__.py")
with open(ver_path, 'a') as f:
    f.write("__version__ = '1.7.2'\n")
    version = '1.7.2' + '+das.dtk' + dtk_version
    f.write(f"__dcu_version__ = '{version}'\n")

long_description="""
Faiss is a library for efficient similarity search and clustering of dense
vectors. It contains algorithms that search in sets of vectors of any size,
 up to ones that possibly do not fit in RAM. It also contains supporting
code for evaluation and parameter tuning. Faiss is written in C++ with
complete wrappers for Python/numpy. Some of the most useful algorithms
are implemented on the GPU. It is developed by Facebook AI Research.
"""

setup(
    name='faiss',
    version='1.7.2' + "+das.dtk" + dtk_version,
    description='A library for efficient similarity search and clustering of dense vectors',
    long_description=long_description,
    url='https://github.com/facebookresearch/faiss',
    author='Matthijs Douze, Jeff Johnson, Herve Jegou, Lucas Hosseini',
    author_email='matthijs@fb.com',
    license='MIT',
    keywords='search nearest neighbors',
    
    install_requires=['numpy'],
    packages=['faiss', 'faiss.contrib'],
    package_data={
        'faiss': ['*.so', '*.pyd'],  # 包含 .so 和 .pyd 文件
    },
    zip_safe=False,
)
