import os
import sys
import setuptools
import pathlib
from setuptools import Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.install import install
import subprocess

# VERSION is in the same path as setup.py
here = os.path.abspath(os.path.dirname(__file__))
with open(os.path.join(here, 'VERSION')) as version_file:
    version_ = version_file.read().strip()

def check_cmake():
    try:
        out = subprocess.check_output(['cmake', '--version'])
        return True
    except OSError:
        return False

# Extend the command line options available to the install phase.
# These arguments must come after `install` on the command line, e.g.:
#    python3 setup.py install --mpi --arch=skylake
#    pip3 install --install-option '--mpi' --install-option '--arch=skylake' .
class install_command(install):
    user_options = install.user_options + [
        ('mpi',   None, 'enable mpi support (requires MPI library)'),
        ('gpu',   None, 'enable nvidia cuda support (requires cudaruntime and nvcc)'),
        ('vec',   None, 'enable vectorization'),
        ('arch=', None, 'cpu architecture, e.g. haswell, skylake, armv8-a'),
    ]

    def initialize_options(self):
        install.initialize_options(self)
        self.mpi  = None
        self.gpu  = None
        self.arch = None
        self.vec  = None

    def finalize_options(self):
        install.finalize_options(self)

    def run(self):
        # The options are stored in a dictionary cl_opt, with the following keys:
        #   'mpi'  : build with MPI support (boolean).
        #   'gpu'  : build with CUDA support (boolean).
        #   'vec'  : generate SIMD vectorized kernels for CPU micro-architecture (boolean).
        #   'arch' : target CPU micro-architecture (string).
        global cl_opt
        cl_opt = {
            'mpi' : self.mpi is not None,
            'gpu' : self.gpu is not None,
            'vec' : self.vec is not None,
            'arch': "native" if self.arch is None else self.arch
        }
        install.run(self)

class cmake_extension(Extension):
    def __init__(self, name):
        Extension.__init__(self, name, sources=[])

class cmake_build(build_ext):
    def run(self):
        if not check_cmake():
            raise RuntimeError('CMake is not available. CMake 3.12 is required.')

        # cl_opt contains the command line arguments passed to install, via
        # --install-option if using pip.
        # pip skips building wheels when --install-option flags are set.
        # However, when no --install-options are passed, it runs build_ext
        # without running install_command, required to create and set cl_opt.
        # This hack works around this. I think that the upshot of this is
        # that only wheels built with default configuration will be possible.

        if 'cl_opt' not in globals():
            cl_opt = {
                    'mpi': False,
                    'gpu': False,
                    'vec': False,
                    'arch': 'native'
            }

        # The path where CMake will be configured and Arbor will be built.
        build_directory = os.path.abspath(self.build_temp)
        # The path where the package will be copied after building.
        lib_directory = os.path.abspath(self.build_lib)
        # The path where the Python package will be compiled.
        source_path = build_directory + '/python/arbor'
        # Where to copy the package after it is built, so that whatever the next phase is
        # can copy it into the target 'prefix' path.
        dest_path = lib_directory + '/arbor'

        cmake_args = [
            '-DARB_WITH_PYTHON=on',
            '-DPYTHON_EXECUTABLE=' + sys.executable,
            '-DARB_WITH_MPI={}'.format('on' if cl_opt['mpi'] else 'off'),
            '-DARB_WITH_GPU={}'.format('on' if cl_opt['gpu'] else 'off'),
            '-DARB_VECTORIZE={}'.format('on' if cl_opt['vec'] else 'off'),
            '-DARB_ARCH={}'.format(cl_opt['arch']),
        ]

        print('-'*5, 'command line options: {}'.format(cl_opt))
        print('-'*5, 'cmake arguments: {}'.format(cmake_args))

        cfg = 'Debug' if self.debug else 'Release'
        build_args = ['--config', cfg]

        cmake_args += ['-DCMAKE_BUILD_TYPE=' + cfg]

        # Assuming Makefiles
        build_args += ['--', '-j4']

        self.build_args = build_args

        env = os.environ.copy()
        env['CXXFLAGS'] = '{}'.format(env.get('CXXFLAGS', ''))
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)

        # CMakeLists.txt is in the same directory as this setup.py file
        cmake_list_dir = os.path.abspath(os.path.dirname(__file__))
        print('-'*20, 'Configure CMake')
        subprocess.check_call(['cmake', cmake_list_dir] + cmake_args,
                              cwd=self.build_temp, env=env)

        print('-'*20, 'Build')
        cmake_cmd = ['cmake', '--build', '.'] + self.build_args
        subprocess.check_call(cmake_cmd,
                              cwd=self.build_temp)

        # Copy from build path to some other place from whence it will later be installed.
        # ... or something like that
        # ... setuptools is an enigma monkey patched on a mystery
        if not os.path.exists(dest_path):
            os.makedirs(dest_path, exist_ok=True)
        self.copy_tree(source_path, dest_path)

setuptools.setup(
    name='arbor',
    version=version_,
    python_requires='>=3.6',

    install_requires=[],
    setup_requires=[],
    zip_safe=False,
    ext_modules=[cmake_extension('arbor')],
    cmdclass={
        'build_ext': cmake_build,
        'install':   install_command,
    },

    author='The lovely Arbor devs.',
    url='https://github.com/arbor-sim/arbor',
    description='High performance simulation of networks of multicompartment neurons.',
    long_description='',
    classifiers=[
        'Development Status :: 4 - Beta', # Upgrade to "5 - Production/Stable" on release of v0.3
        'Intended Audience :: Science/Research',
        'Topic :: Scientific/Engineering :: Build Tools',
        'License :: OSI Approved :: BSD License'
        'Programming Language :: Python :: 3.6',
        'Programming Language :: Python :: 3.7',
        'Programming Language :: Python :: 3.8',
    ],
    project_urls={
        'Source': 'https://github.com/arbor-sim/arbor',
        'Documentation': 'https://arbor.readthedocs.io',
        'Bug Reports': 'https://github.com/arbor-sim/arbor/issues',
    },
)