#!/usr/bin/env python

"""This is the setup and installation script for Asap.

To install, simply run
    python setup.py install --user

CUSTOMIZATION:
If you need to change compiler options or other stuff, 
do NOT EDIT THIS FILE.  Instead, you should edit 
customize.py or one of its friends.

Customization information is read from the first of
these files found:
    customize-local.py
    customize-hostname.domain.py
    customize.py
If setup.py is given the option --customize=myfile.py then
that file is used instead of the above.

Only one of these files is read.  See customize.py for further
instructions.  But remember: MOST USERS DO NOT NEED THIS!

"""

from __future__ import print_function
import distutils
import distutils.util
import distutils.spawn
import os
import os.path as op
import stat
import re
import sys
from distutils.command.build_ext import build_ext as _build_ext
from distutils.command.build_scripts import build_scripts as _build_scripts
#from distutils.command.sdist import sdist as _sdist
from distutils.core import setup, Extension
from distutils.sysconfig import get_config_vars
from glob import glob
import numpy
import platform
import subprocess
import recordversion   # In this project - records the compiler information into CPP

assert sys.version_info >= (2, 6)

description = "ASAP - classical potentials for MD with ASE."

long_description = """\
ASAP (Atomic SimulAtion Program or As Soon As Possible) is a
package for large-scale molecular dynamics within the Atomic
Simulation Environment (ASE).  It implements a number of 'classical'
potentials, most importantly the Effective Medium Theory, and also the
mechanisms for domain-decomposition of the atoms."""

folders = ['Basics', 'Potentials', 'Interface', 'Brenner', 'Tools',
           'PTM', 'PTM/qcprot', 'PTM/voronoi']
kim_folders = ['OpenKIMimport']
parallel_folders = ['Parallel', 'ParallelInterface']
exclude_files = ['Interface/AsapModule.cpp']
serial_only_files = ['Interface/AsapSerial.cpp']


libraries = ['m']
library_dirs = []
include_dirs = []
extra_link_args = []
extra_compile_args = []
remove_compile_args = []
runtime_library_dirs = []
extra_objects = []
define_macros = [] 
undef_macros = ['NDEBUG']

mpi_libraries = []
mpi_library_dirs = []
mpi_include_dirs = []
mpi_runtime_library_dirs = []
mpi_define_macros = []
mpi_undef_macros = []
mpi_compiler = ('mpicc', 'mpicxx')

use_intel_compiler = False  # False, True or 'auto' - enable in customize.py
intel_compiler = ('icc', 'icpc')
intel_mpi_compiler = ('env OMPI_MPICC=icc mpicc', 'env OMPI_MPICXX=icpc mpicxx')
intel_compile_args = ['-O3', '-g', '-xHost', '-fPIC']
intel_libraries = ['svml', 'm']
intel_link_args = []
# Some Intel compilers are broken, and cause Asap to fail.
intel_blacklist = ['15.0.1', '17.0.1', '17.0.2']

# Default values for compilation with MKL - only if MKL sets $MKLROOT or if
# defined manually in customize.py
mkl_supported = 'MKLROOT' in os.environ
mkl_define_macros = [('ASAP_MKL',1)]
mkl_compile_args = []
mkl_intel_compile_args = []
mkl_libraries = ['mkl_rt']
mkl_library_dirs = []
if mkl_supported:
    mkl_include_dirs = ['{MKLROOT}/include'.format(**os.environ)]
else:
    mkl_include_dirs = []
    
systemname, hostname, dummy1, dummy2, machinename, processorname = platform.uname()

#platform_id = ''  # Not changable for Asap (see gpaw/setup.py)

# Get the version number from Python/asap3/version.py
# Get the current version number:
with open('Python/asap3/version.py') as fd:
    version = re.search('__version__ = "(.*)"', fd.read()).group(1)
print("Asap version number:", version)

# Look for --customize arg
for i, arg in enumerate(sys.argv):
    if arg.startswith('--customize'):
        customize = sys.argv.pop(i).split('=')[1]
        break
else:
    # No --customize arg.
    for c in ('customize-local.py', 
              'customize-{0}.py'.format(hostname),
              'customize.py'):
        if os.path.exists(c):
            customize = c
            break
        elif '-' in c:
            print(c, "not found - this is OK.")
    else:
        raise RuntimeError("customize.py not found - broken source distribution.")

print("Reading customization from", customize)
exec(open(customize).read())

# Check if command line arguments override the use of Intel compilers.
for i, arg in enumerate(sys.argv):
    if arg.lower() == '--with-intel':
        use_intel_compiler = True
        del sys.argv[i]
        break
    if arg.lower() == '--without-intel':
        use_intel_compiler = False
        del sys.argv[i]
        break

# Check if Intel compiler should be used iff present
if use_intel_compiler == 'auto':
    use_intel_compiler = bool(distutils.spawn.find_executable(intel_compiler[0]))
    print("Intel compiler autodetection: use Intel compiler =", use_intel_compiler)

# Check for command line argument asking for MKI
use_mkl = False
for i, arg in enumerate(sys.argv):
    if arg.lower() == '--with-mkl':
        print("Intel Math Kernel Library support enabled.")
        use_mkl = True
        del sys.argv[i]
        break

if use_mkl:
    if not mkl_supported:
        raise RuntimeError("You asked for Intel MKL support, but it is not available. Configure in customize.py or set $MKLROOT")
    extra_compile_args += mkl_compile_args
    intel_compile_args += mkl_intel_compile_args
    include_dirs += mkl_include_dirs
    libraries += mkl_libraries
    library_dirs += mkl_library_dirs
    define_macros += mkl_define_macros

# Check for the --clear-compiler-environment command line option
envirhack = False
for i, arg in enumerate(sys.argv):
    if arg.lower() == '--clear-compiler-environment':
        envirhack = True
        del sys.argv[i]
        break

if envirhack:
    print("Clearing environment variables (e.g. for EasyBuild).")
    for evar in ('CC', 'CXX', 'CFLAGS', 'CXXFLAGS', 'LDFLAGS'):
        if evar in os.environ:
            print("   Clearing {0}  (was '{1}')".format(evar,
                                                        os.environ[evar]))
            del os.environ[evar]
# Check if the MPI compiler is present.

if mpi_compiler:
    for mpicomp in mpi_compiler:
        if not distutils.spawn.find_executable(mpicomp):
            print("WARNING: No MPI compiler '{}': Not building parallel version.".format(mpicomp))
            mpi_compiler = None
            break

# If intel compiler selected, we need to change some compiler flags.
# This comes after the test for the MPI compiler, so the MPI compiler
# wrapper script is detected properly.
cfgDict = get_config_vars()
if use_intel_compiler:
    print("Using Intel compiler:")
    for var in ['BASECFLAGS', 'CFLAGS', 'OPT', 'PY_CFLAGS',
                'CCSHARED']: # 'CFLAGSFORSHARED', 'LINKFORSHARED']
        cfgDict[var] = ''  # Clear it
    cfgDict['BASECFLAGS'] = cfgDict['CFLAGS'] = ' '.join(intel_compile_args)
    libraries += intel_libraries
    for key in ['CC', 'CXX', 'LDSHARED', 'LDCXXSHARED']:
        if key in cfgDict:
            value = cfgDict[key].split()
            # first argument is the compiler/linker.
            value[0] = intel_compiler[1]  # C++ compiler
            cfgDict[key] = ' '.join(value)
            print("    Setting {0}={1}".format(key, cfgDict[key]))
    mpi_compiler = intel_mpi_compiler
    extra_compile_args = []  # Specified in intel_compile_args
    extra_link_args = intel_link_args
    # Check the version number of the intel compiler
    iccver = subprocess.check_output(intel_compiler[0] + ' -dumpversion', shell=True)
    if iccver in intel_blacklist:
        raise RuntimeError("Cannot compile with Intel compiler version {} due to a compiler bug.".format(iccver))


# Remove unneeded compiler arguments
if remove_compile_args:
    for fvar in ['BASECFLAGS', 'CFLAGS', 'OPT', 'PY_CFLAGS',
                  'CFLAGSFORSHARED']:
        if fvar in cfgDict:
            cflags = cfgDict[fvar].split()
            for arg in remove_compile_args:
                if arg in cflags:
                    cflags.remove(arg)
                    print("Removed compiler argument", arg)
            cfgDict[fvar] = ' '.join(cflags)
    
def runcmd(cmd, verb=False):
    if verb:
        print(cmd)
    x = os.system(cmd)
    if x:
        raise RuntimeError("Command failed: "+cmd)

def build_interpreter(compiler, common_src, parallel_src,
                      libraries, library_dirs, include_dirs, 
                      extra_link_args, extra_compile_args,
                      runtime_library_dirs, extra_objects, define_macros,
                      undef_macros):
    """Build the asap-python executable."""
    print("\n\n*** Building the asap-python executable ***\n\n")
    (c_compiler, cpp_compiler) = compiler

    # Get configuration variable
    cfgDict = get_config_vars()
    plat = distutils.util.get_platform() + '-' + sys.version[0:3]
    if hasattr(sys, 'gettotalrefcount'):
        plat += '-pydebug'

    # Split sources into C and C++, and find object names.
    c_sources = {}
    cpp_sources = {}
    common_sources = {} # Will not be compiled, don't care about C/C++
    for s in parallel_src:
        dir = 'build/temp.{0}/{1}'.format(plat, os.path.dirname(s))
        if not os.path.exists(dir):
            os.mkdir(dir)   # Create folder for object file.
        if s.endswith('.cpp'):
            objname = 'build/temp.{0}/{1}.o'.format(plat, s[:-4])
            cpp_sources[s] = objname
        elif s.endswith('.c'):
            objname = 'build/temp.{0}/{1}.o'.format(plat, s[:-2])
            c_sources[s] = objname
        else:
            raise RuntimeError("Source file {0} is neither C nor C++.".format(s))
    for s in common_src:
        if s.endswith('.cpp'):
            objname = 'build/temp.{0}/{1}.o'.format(plat, s[:-4])
            common_sources[s] = objname
        elif s.endswith('.c'):
            objname = 'build/temp.{0}/{1}.o'.format(plat, s[:-2])
            common_sources[s] = objname
        else:
            raise RuntimeError("Source file {0} is neither C nor C++.".format(s))        
    exefile = "build/bin.{0}/asap-python".format(plat)
    if not os.path.exists("build/bin.{0}".format(plat)):
        os.mkdir("build/bin.{0}".format(plat))
        
    # Now construct the compilation command lines
    macros = ' '.join(['-D%s=%s' % x for x in define_macros if x[0].strip()])
    macros += ' '.join(['-U%s' % (x,) for x in undef_macros if x.strip()])
    
    include_dirs = include_dirs + [cfgDict['INCLUDEPY'],
                                   cfgDict['CONFINCLUDEPY']]
    includes = ' '.join(['-I' + incdir for incdir in include_dirs])

    library_dirs = library_dirs + [cfgDict['LIBPL']]
    lib_dirs = ' '.join(['-L' + lib for lib in library_dirs])

    libs = ' '.join(['-l' + lib for lib in libraries if lib.strip()])
    libs += ' ' + cfgDict.get('BLDLIBRARY', '-lpython%s' % cfgDict['VERSION'])
    libs += ' ' + ' '.join([libs, cfgDict['LIBS'], cfgDict['LIBM']])
    if sys.platform == 'darwin':
        runtime_libs = ''   # No -rpath on macOS
    else:
        runtime_lib_option = '-Wl,-rpath='
        runtime_libs = ' '.join([runtime_lib_option + lib
                                for lib in runtime_library_dirs])
    extra_link_args = extra_link_args + [cfgDict['LDFLAGS'],
                                         cfgDict['MODLIBS'],
                                         cfgDict['LINKFORSHARED']]
    # Compile
    all_objects = []
    cflags = cfgDict['CFLAGS'].split()
    compile_args_c = cflags + list(extra_compile_args)
    compile_args_cpp = cflags + list(extra_compile_args)
    if '-std=c++11' in compile_args_c:
        compile_args_c.remove('-std=c++11')
    for compiler, sources, compargs in (\
                (c_compiler, c_sources, compile_args_c),
                                        (cpp_compiler, cpp_sources, compile_args_cpp)):
        print("*** COMPILING: ", list(sorted(sources.keys())))
        for src in sorted(sources.keys()):
            obj = sources[src]
            all_objects.append(obj)
            cmd = ('%s %s %s %s -o %s -c %s ') % \
                  (compiler,
                   ' '.join(compargs),
                   macros,
                   includes,
                   obj,
                   src)
            runcmd(cmd, True)
    # Link
    all_objects = list(common_sources.values()) + all_objects
    cmd = ('%s -o %s %s %s %s %s %s %s') % \
          (cpp_compiler,
           exefile,
           ' '.join(all_objects),
           ' '.join(extra_objects),
           lib_dirs,
           libs,
           runtime_libs,
           ' '.join(extra_link_args))
    runcmd(cmd, True)


# Check if we are making a source distribution.  In that case, we
# should remove the VersionInfo_autogen folder and all files in it.
versiondir = "VersionInfo_autogen"
is_making_distro = 'sdist' in sys.argv
if is_making_distro and os.path.isdir(versiondir):
    print("Clearing", versiondir)
    for f in os.listdir(versiondir):
        os.remove(os.path.join(versiondir, f))
    os.rmdir(versiondir)

# Create the version.cpp file
try:
    host = os.uname()[1]
except:
    host = 'unknown'
if not is_making_distro:
    versioninfo_s = "{0}/version_{1}_s.cpp".format(versiondir, host)
    versioninfo_p = "{0}/version_{1}_p.cpp".format(versiondir, host)
    if not os.path.exists(versiondir):
        os.mkdir("VersionInfo_autogen")
    t = os.stat("Python/asap3/version.py")[stat.ST_MTIME]
    if (not os.path.exists(versioninfo_s) or t > os.stat(versioninfo_s)[stat.ST_MTIME] or
        not os.path.exists(versioninfo_p) or t > os.stat(versioninfo_p)[stat.ST_MTIME]):
        print("Recording version info into VersionInfo_autogen")
        print("  CC =", cfgDict['CC'])
        print("  CFLAGS =", cfgDict['CFLAGS'])
        serialcomp = "distutils with {0} {1}".format(cfgDict['CC'], cfgDict['CFLAGS'])
        with open(versioninfo_s, "wt") as versioncpp:
            versioncpp.write(recordversion.contents % (version, 'serial', host, serialcomp, version))
        if mpi_compiler is not None:
            parallelcomp = "distutils with {0} {1}".format(mpi_compiler[1], 
                                                           " ".join([cfgDict['CFLAGS']]
                                                                    +extra_compile_args))
            with open(versioninfo_p, "wt") as versioncpp:
                versioncpp.write(recordversion.contents % (version, 'parallel', host, parallelcomp, version))


if kim_folders and 'KIM_HOME' in os.environ:
    kh = os.environ['KIM_HOME']
    include_dirs += [os.path.join(kh, 'include', 'kim-api')]
    include_dirs += kim_folders
    library_dirs += [os.path.join(kh, 'lib64'), os.path.join(kh, 'lib'),]
    runtime_library_dirs += [os.path.join(kh, 'lib64'), os.path.join(kh, 'lib')]
    libraries += ['kim-api']
    define_macros += [('WITH_OPENKIM', '1')]
    folders.extend(kim_folders)
else:
    # Try to use pkgconfig to locate OpenKIM
    failed = subprocess.call("pkg-config --exists libkim-api",
                             shell=True)
    if not failed:
        # pkg-config is installed and so is OpenKIM.
        kimincl = subprocess.check_output(
            "pkg-config --cflags-only-I libkim-api", shell=True).decode()
        include_dirs += [f.strip() for f in kimincl.split('-I') if f]
        include_dirs += kim_folders
        kimlibd = subprocess.check_output(
            "pkg-config --libs-only-L libkim-api", shell=True).decode()
        kimliblist = [f.strip() for f in kimlibd.split('-L') if f]
        library_dirs += kimliblist
        runtime_library_dirs += kimliblist
        kimlibl = subprocess.check_output(
            "pkg-config --libs-only-l libkim-api", shell=True).decode()
        libraries += [f.strip() for f in kimlibl.split('-l') if f]
        define_macros += [('WITH_OPENKIM', '1')]
        folders.extend(kim_folders)
        
include_dirs += folders
include_dirs.append(numpy.get_include())
mpi_include_dirs += parallel_folders

print("Identifying source files")
# Find source files for serial compilation
common_src_files = []
for d in folders:
    for f in os.listdir(d):
        if f.endswith('.cpp'):
            fn = os.path.join(d,f)
            if fn not in (exclude_files + serial_only_files):
                common_src_files.append(fn)
                #print("  ", fn)
serial_src_files = common_src_files + serial_only_files
if not is_making_distro:
    serial_src_files.append(versioninfo_s)
serial_src_files.sort()
parallel_src_files = []
for d in parallel_folders:
    for f in os.listdir(d):
        if f.endswith('.cpp') or f.endswith('.c'):
            fn = os.path.join(d,f)
            parallel_src_files.append(fn)
if not is_making_distro:
    parallel_src_files.append(versioninfo_p)

print("Identifying Python submodules")
packages = []
for dirname, dirnames, filenames in os.walk('Python/asap3'):
    if '__init__.py' in filenames:
        dname = dirname.split('/')[1:]  # Remove leading Python/
        packages.append('.'.join(dname))

extensions = [Extension('_asap_p{0}'.format(sys.version_info[0]),
                        serial_src_files,
                        libraries=libraries,
                        library_dirs=library_dirs,
                        include_dirs=include_dirs,
                        define_macros=define_macros,
                        undef_macros=undef_macros,
                        extra_link_args=extra_link_args,
                        extra_compile_args=extra_compile_args,
                        runtime_library_dirs=runtime_library_dirs,
                        )]

# Scripts
scripts = ['scripts/asap-qsub', 'scripts/asap-sbatch']
if mpi_compiler:
    plat = distutils.util.get_platform() + '-' + sys.version[0:3]
    if hasattr(sys, 'gettotalrefcount'):
        plat += '-pydebug'
    scripts.append("build/bin.{0}/asap-python".format(plat))
    
class build_ext(_build_ext):
    def run(self):
        _build_ext.run(self)
        if mpi_compiler:
            # Also build asap-python:
            build_interpreter(
                mpi_compiler,
                common_src_files,
                parallel_src_files,
                libraries=libraries+mpi_libraries,
                library_dirs=library_dirs+mpi_library_dirs,
                include_dirs=include_dirs+mpi_include_dirs,
                extra_link_args=extra_link_args,
                extra_compile_args=extra_compile_args,
                runtime_library_dirs=runtime_library_dirs,
                extra_objects=extra_objects,
                define_macros=define_macros+mpi_define_macros,
                undef_macros=undef_macros+mpi_undef_macros)

class build_scripts(_build_scripts):
    # Python 3 will try to read the asap-python executable, handle
    # that one specially.
    def copy_scripts(self):
        executables = []
        for i, s in enumerate(self.scripts):
            if 'asap-python' in s:
                print("Protecting", s, "from ajustment")
                executables.append(s)
        for s in executables:
            self.scripts.remove(s)
        assert len(executables) <= 1
        x =_build_scripts.copy_scripts(self)
        if x is None:
            # Python 2
            outfiles = []
            updated_files = []
        else:
            outfiles, updated_files = x
        for script in executables:
            outfile = op.join(self.build_dir, op.basename(script))
            outfiles.append(outfile)
            updated_files.append(outfile)
            self.copy_file(script, outfile)
        return outfiles, updated_files
            
setup(name="asap3",
      version=version,
      description=description,
      long_description=long_description,
      maintainer="Jakob Schiotz et. al.",
      maintainer_email="schiotz@fysik.dtu.dk",
      url="https://wiki.fysik.dtu.dk/asap",
      packages=packages,
      package_dir={'asap3': 'Python/asap3'},
      ext_modules=extensions,
      license='LGPLv3',
      platforms=['unix'],
      scripts=scripts,
      cmdclass={'build_ext': build_ext,
                'build_scripts': build_scripts},
      classifiers=[
          'Development Status :: 5 - Production/Stable',
          'License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)',
          'Operating System :: POSIX',
          'Operating System :: MacOS :: MacOS X',
                    'Programming Language :: Python :: 2',
          'Programming Language :: Python :: 2.6',
          'Programming Language :: Python :: 2.7',
          'Programming Language :: Python :: 3',
          'Programming Language :: Python :: 3.5',
          'Topic :: Scientific/Engineering :: Physics'
      ])

