# The setup used here is derived with bits from:
# - https://github.com/rmcgibbo/npcuda-example

import os
import sys
import platform
import subprocess
from os.path import join as pjoin
import warnings
from setuptools import setup, find_packages
from setuptools.command.sdist import sdist
from distutils.extension import Extension
from distutils import unixccompiler
from setuptools.command.build_ext import build_ext as build_ext_orig
from setuptools.command.bdist_egg import bdist_egg
from setuptools.command.install_lib import install_lib as su_install_lib
from setuptools.command.sdist import sdist as su_sdist
from pyvkfft.version import __version__, git_version, vkfft_git_version

# Maximum number of dimensions VkFFT can handle. VkFFT sets this to 4,
# pyvkfft uses a default of 8. Set an environment variable
# VKFFT_MAX_FFT_DIMENSIONS to increase this
if 'VKFFT_MAX_FFT_DIMENSIONS' in os.environ:
    VKFFT_MAX_FFT_DIMENSIONS = os.environ['VKFFT_MAX_FFT_DIMENSIONS']
else:
    VKFFT_MAX_FFT_DIMENSIONS = 8


def find_in_path(name, path):
    """Find a file in a search path"""

    # Adapted fom http://code.activestate.com/recipes/52224
    for dir in path.split(os.pathsep):
        binpath = pjoin(dir, name)
        if os.path.exists(binpath):
            return os.path.abspath(binpath)
    return None


def locate_cuda():
    """Locate the CUDA environment on the system

    Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64'
    and values giving the absolute path to each directory.

    Starts by looking for the CUDAHOME or CUDA_PATH env variable.
    If not found, find 'nvcc' in the PATH.
    """
    # search for nvcc
    app_name = 'nvcc.exe' if platform.system() == "Windows" else 'nvcc'
    nvcc, home = None, None
    for p in ['CUDA_PATH', 'CUDA_HOME', 'CUDAHOME']:
        if p in os.environ:
            if os.path.exists(pjoin(os.environ[p], 'bin', app_name)):
                nvcc = pjoin(os.environ[p], 'bin', app_name)
                home = os.environ[p]
                break
    if nvcc is None:
        # Search the PATH for NVCC
        nvcc = find_in_path(app_name, os.environ['PATH'])
        if nvcc is not None:
            home = os.path.dirname(os.path.dirname(nvcc))
        else:
            raise EnvironmentError('The nvcc binary could not be '
                                   'located in your $PATH. Either add it to your path, '
                                   'or set $CUDA_PATH')

    if platform.system() == "Windows":
        libdir = pjoin(home, 'lib', 'x64')
        extra_compile_args = ['-O3', '--ptxas-options=-v', '-Xcompiler', '-MD',
                              f'-DVKFFT_MAX_FFT_DIMENSIONS={VKFFT_MAX_FFT_DIMENSIONS}',
                              "--use-local-env", "-ccbin",
                              os.path.dirname(find_in_path('cl.exe', os.environ['PATH']))]
        tmp = pjoin(os.path.dirname(find_in_path('cl.exe', os.environ['PATH'])).split('bin')[0],
                    'lib', 'x64')
        extra_link_args = ['-L%s' % libdir, "--use-local-env", "-ccbin",
                           os.path.dirname(find_in_path('cl.exe', os.environ['PATH'])),
                           '-L%s' % tmp]
    else:
        if os.path.exists(pjoin(home, 'lib64')):
            libdir = pjoin(home, 'lib64')
        else:
            libdir = pjoin(home, 'lib')
        extra_compile_args = ['-O3', '-g', '--ptxas-options=-v', '-std=c++11',
                              '--compiler-options=-fPIC',
                              f'-DVKFFT_MAX_FFT_DIMENSIONS={VKFFT_MAX_FFT_DIMENSIONS}']
        extra_link_args = ['--shared', '-L%s' % libdir]
    cudaconfig = {'home': home, 'nvcc': nvcc,
                  'include_dirs': [pjoin(home, 'include'), 'src/VkFFT/vkFFT'],
                  'extra_compile_args': extra_compile_args,
                  'extra_link_args': extra_link_args}
    if 'CONDA_PREFIX' in os.environ:
        cudaconfig['include_dirs'].append(f"{os.environ['CONDA_PREFIX']}/include")
        cudaconfig['include_dirs'].append(f"{os.environ['CONDA_PREFIX']}/Library/include")
    return cudaconfig


def locate_opencl():
    """
    Get the opencl configuration
    :return:
    """
    include_dirs = ['src/VkFFT/vkFFT']
    library_dirs = []
    extra_compile_args = ['-std=c++11', '-Wno-format-security',
                          f'-DVKFFT_MAX_FFT_DIMENSIONS={VKFFT_MAX_FFT_DIMENSIONS}']
    extra_link_args = None
    if platform.system() == 'Darwin':
        libraries = None
        extra_link_args = ['-Wl,-framework,OpenCL']
    elif platform.system() == "Windows":
        # Add include & lib dirs if possible from usual nvidia and AMD paths
        for path in ["CUDA_HOME", "CUDAHOME", "CUDA_PATH"]:
            if path in os.environ:
                include_dirs.append(pjoin(os.environ[path], 'include'))
                library_dirs.append(pjoin(os.environ[path], 'lib', 'x64'))
        library_dirs.append(pjoin(os.path.dirname(find_in_path('cl.exe', os.environ['PATH'])).split('bin')[0],
                                  'lib', 'x64'))
        libraries = ['OpenCL']
        extra_compile_args = [f'/DVKFFT_MAX_FFT_DIMENSIONS={VKFFT_MAX_FFT_DIMENSIONS}']
    else:
        # Linux
        libraries = ['OpenCL']
        extra_link_args = ['--shared']

    if 'CONDA_PREFIX' in os.environ:
        include_dirs.append(f"{os.environ['CONDA_PREFIX']}/include")
        include_dirs.append(f"{os.environ['CONDA_PREFIX']}/Library/include")
        library_dirs.append(f"{os.environ['CONDA_PREFIX']}/Library/Lib")

    opencl_config = {'libraries': libraries, 'extra_link_args': extra_link_args,
                     'include_dirs': include_dirs, 'library_dirs': library_dirs,
                     'extra_compile_args': extra_compile_args}
    return opencl_config


class build_ext_custom(build_ext_orig):
    """Custom `build_ext` command which will correctly compile and link
    the OpenCL and CUDA modules. The hooks are based on the name of the extension"""

    def build_extension(self, ext):
        if "cuda" in ext.name:
            # Use nvcc for compilation. This assumes all sources are .cu files
            # for this extension.
            default_compiler = self.compiler
            # Create unix compiler patched for cu
            self.compiler = unixccompiler.UnixCCompiler()
            self.compiler.src_extensions.append('.cu')
            tmp = CUDA['nvcc']  # .replace('\\\\','toto')
            self.compiler.set_executable('compiler_so', [tmp])
            self.compiler.set_executable('linker_so', [tmp])
            if platform.system() == "Windows":
                CUDA['extra_link_args'] += ['--shared', '-Xcompiler', '-MD']
                # pythonXX.lib must be in the linker paths
                # Is using sys.prefix\libs always correct ?
                CUDA['extra_link_args'].append('-L%s' % pjoin(sys.prefix, 'libs'))

            super().build_extension(ext)
            # Restore default linker and compiler
            self.compiler = default_compiler
        else:
            super().build_extension(ext)

    def get_export_symbols(self, ext):
        """ Hook to make sure we get the correct symbols for windows"""
        if ("opencl" in ext.name or "cuda" in ext.name) and platform.system() == "Windows":
            return ext.export_symbols
        return super().get_export_symbols(ext)

    def get_ext_filename(self, ext_name):
        """ Hook to make sure we keep the correct name (*.so) for windows"""
        if ("opencl" in ext_name or "cuda" in ext_name) and platform.system() == "Windows":
            return ext_name + '.so'
        return super().get_ext_filename(ext_name)


ext_modules = []
install_requires = ['numpy', 'psutil']
exclude_packages = ['examples']
CUDA = None
OPENCL = None

if "VKFFT_BACKEND" in os.environ:
    v = os.environ["VKFFT_BACKEND"]
    # Environment variable to manually select vkfft backends. useful e.g. if nvidia tools
    # are installed but not functional
    # e.g. use:
    #   VKFFT_BACKEND=opencl pip install pyvkfft
    #   VKFFT_BACKEND=cuda pip install .
    if 'opencl' not in v.lower():
        exclude_packages.append('opencl')
    if 'cuda' not in v.lower():
        exclude_packages.append('cuda')

if 'cuda' not in exclude_packages:
    try:
        CUDA = locate_cuda()
        vkfft_cuda_ext = Extension('pyvkfft._vkfft_cuda',
                                   sources=['src/vkfft_cuda.cu'],
                                   libraries=['nvrtc', 'cuda'],
                                   extra_compile_args=CUDA['extra_compile_args'],
                                   include_dirs=CUDA['include_dirs'] + ['src'],
                                   extra_link_args=CUDA['extra_link_args'],
                                   depends=['vkFFT.h']
                                   )
        ext_modules.append(vkfft_cuda_ext)
        # install_requires.append("pycuda")
        try:
            import pycuda

            has_pycuda = True
        except ImportError:
            has_pycuda = False
        try:
            import cupy
        except ImportError:
            if has_pycuda is False:
                print("Reminder: you need to install either PyCUDA or CuPy to use pyvkfft.cuda")
    except:
        exclude_packages.append('cuda')
        warnings.warn("CUDA not available ($CUDAHOME/$CUDA_PATH variables missing "
                      "and nvcc not in path. "
                      "Skipping pyvkfft.cuda module installation.", UserWarning)

if 'opencl' not in exclude_packages:
    OPENCL = locate_opencl()
    install_requires.append('pyopencl')

    # OpenCL extension
    vkfft_opencl_ext = Extension('pyvkfft._vkfft_opencl',
                                 sources=['src/vkfft_opencl.cpp'],
                                 extra_compile_args=OPENCL['extra_compile_args'],
                                 include_dirs=OPENCL['include_dirs'] + ['src'],
                                 libraries=OPENCL['libraries'],
                                 library_dirs=OPENCL['library_dirs'],
                                 extra_link_args=OPENCL['extra_link_args'],
                                 depends=['vkFFT.h']
                                 )

    ext_modules.append(vkfft_opencl_ext)

with open("README.rst", "r", encoding="utf-8") as fh:
    long_description = fh.read()


class bdist_egg_disabled(bdist_egg):
    """ Disabled bdist_egg, to prevent use of 'python setup.py install' """

    def run(self):
        sys.exit("Aborting building of eggs. Please use `pip install .` to install from source.")


# Console scripts, available e.g. as 'pyvkfft-test'
scripts = ['pyvkfft/scripts/pyvkfft_test.py', 'pyvkfft/scripts/pyvkfft_test_suite.py',
           'pyvkfft/scripts/pyvkfft_benchmark.py', 'pyvkfft/scripts/pyvkfft_info.py']

console_scripts = []
for s in scripts:
    s1 = os.path.splitext(os.path.split(s)[1])[0]
    s0 = os.path.splitext(s)[0]
    console_scripts.append("%s = %s:main" % (s1.replace('_', '-'), s0.replace('/', '.')))


class pyvkfft_sdist(su_sdist):
    """Hook to include git version of pyvkfft and VkFFT"""

    def make_release_tree(self, base_dir, files):
        super(pyvkfft_sdist, self).make_release_tree(base_dir, files)
        try:
            # Replace git_version_placeholder by real git version
            version_file = os.path.join(base_dir, "pyvkfft/version.py")
            vers = open(version_file).read()
            os.remove(version_file)
            with open(version_file, "w") as fh:
                vers = vers.replace("vkfft_git_version_placeholder", vkfft_git_version())
                vers = vers.replace("git_version_placeholder", git_version())
                fh.write(vers)
        except:
            print("sdist: replacing git_version failed")


class pyvkfft_install_lib(su_install_lib):
    """Hook to include git version of pyvkfft and VkFFT"""

    def run(self):
        super(pyvkfft_install_lib, self).run()
        try:
            # print(self.install_dir, self.build_dir)
            # Replace git_version_placeholder by real git version
            version_file = os.path.join(self.install_dir, "pyvkfft/version.py")
            vers = open(version_file).read()
            os.remove(version_file)
            with open(version_file, "w") as fh:
                vers = vers.replace("vkfft_git_version_placeholder", vkfft_git_version())
                vers = vers.replace("git_version_placeholder", git_version())
                fh.write(vers)
        except:
            print("install_lib: replacing git_version failed")


setup(name="pyvkfft",
      version=__version__,
      description="Python wrapper for the CUDA and OpenCL backends of VkFFT,"
                  "providing GPU FFT for PyCUDA, PyOpenCL and CuPy",
      long_description=long_description,
      ext_modules=ext_modules,
      packages=find_packages(exclude=exclude_packages),
      include_package_data=True,
      author="Vincent Favre-Nicolin",
      author_email="favre@esrf.fr",
      url="https://github.com/vincefn/pyvkfft",
      project_urls={
          "Source": "https://github.com/vincefn/pyvkfft",
          "Documentation": "https://pyvkfft.readthedocs.io/",
          "Bug Tracker": "https://github.com/vincefn/pyvkfft/issues",
          "VkFFT project": "https://github.com/DTolm/VkFFT",
      },
      classifiers=[
          "Programming Language :: Python :: 3",
          "License :: OSI Approved :: MIT License",
          "Operating System :: OS Independent",
          "Environment :: GPU",
      ],
      license='MIT License',
      cmdclass={'build_ext': build_ext_custom,
                'bdist_egg': bdist_egg if 'bdist_egg' in sys.argv else bdist_egg_disabled,
                'sdist': pyvkfft_sdist,
                'install_lib': pyvkfft_install_lib
                },
      install_requires=install_requires,
      extras_require={'doc': ['sphinx', 'nbsphinx', 'nbsphinx-link', 'sphinx-argparse',
                              'sphinx-rtd-theme']},
      test_suite="test",
      entry_points={'console_scripts': console_scripts},
      )