"""MDTraj: a python library for loading, saving, and manipulating molecular dynamics trajectories.

MDTraj provides an easy to use python interface for manipulating MD trajectories.
It supports the reading and writing of molecular dynamics trajectories
in a variety of formats, including full support for PDB, DCD, XTC, TRR, binpos,
AMBER NetCDF, AMBER mdcrd and MDTraj HDF5. The package also provides a command line script
for converting trajectories between supported formats.
"""

DOCLINES = __doc__.split("\n")

import os
import sys
import shutil
import tempfile
import subprocess
from distutils.ccompiler import new_compiler
from setuptools import setup, Extension
from Cython.Distutils import build_ext
import numpy


##########################
VERSION = "0.4.1"
ISRELEASED = True
__version__ = VERSION
##########################


CLASSIFIERS = """\
Development Status :: 3 - Alpha
Intended Audience :: Science/Research
Intended Audience :: Developers
License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
Programming Language :: C
Programming Language :: Python
Topic :: Scientific/Engineering :: Bio-Informatics
Topic :: Scientific/Engineering :: Chemistry
Operating System :: POSIX
Operating System :: Unix
Operating System :: MacOS
"""

################################################################################
# Writing version control information to the module
################################################################################

def git_version():
    # Return the git revision as a string
    # copied from numpy setup.py
    def _minimal_ext_cmd(cmd):
        # construct minimal environment
        env = {}
        for k in ['SYSTEMROOT', 'PATH']:
            v = os.environ.get(k)
            if v is not None:
                env[k] = v
        # LANGUAGE is used on win32
        env['LANGUAGE'] = 'C'
        env['LANG'] = 'C'
        env['LC_ALL'] = 'C'
        out = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
        return out

    try:
        out = _minimal_ext_cmd(['git', 'rev-parse', 'HEAD'])
        GIT_REVISION = out.strip().decode('ascii')
    except OSError:
        GIT_REVISION = "Unknown"

    return GIT_REVISION


def write_version_py(filename='MDTraj/version.py'):
    cnt = """
# THIS FILE IS GENERATED FROM MDTRAJ SETUP.PY
short_version = '%(version)s'
version = '%(version)s'
full_version = '%(full_version)s'
git_revision = '%(git_revision)s'
release = %(isrelease)s

if not release:
    version = full_version
"""
    # Adding the git rev number needs to be done inside write_version_py(),
    # otherwise the import of numpy.version messes up the build under Python 3.
    FULLVERSION = VERSION
    if os.path.exists('.git'):
        GIT_REVISION = git_version()
    else:
        GIT_REVISION = "Unknown"

    if not ISRELEASED:
        FULLVERSION += '.dev-' + GIT_REVISION[:7]

    a = open(filename, 'w')
    try:
        a.write(cnt % {'version': VERSION,
                       'full_version': FULLVERSION,
                       'git_revision': GIT_REVISION,
                       'isrelease': str(ISRELEASED)})
    finally:
        a.close()

################################################################################
# Detection of compiler capabilities
################################################################################

def hasfunction(cc, funcname, include=None, extra_postargs=None):
    # From http://stackoverflow.com/questions/
    #            7018879/disabling-output-when-compiling-with-distutils
    tmpdir = tempfile.mkdtemp(prefix='hasfunction-')
    devnull = oldstderr = None
    try:
        try:
            fname = os.path.join(tmpdir, 'funcname.c')
            f = open(fname, 'w')
            if include is not None:
                f.write('#include %s\n' % include)
            f.write('int main(void) {\n')
            f.write('    %s;\n' % funcname)
            f.write('}\n')
            f.close()
            # Redirect stderr to /dev/null to hide any error messages
            # from the compiler.
            # This will have to be changed if we ever have to check
            # for a function on Windows.
            devnull = open('/dev/null', 'w')
            oldstderr = os.dup(sys.stderr.fileno())
            os.dup2(devnull.fileno(), sys.stderr.fileno())
            objects = cc.compile([fname], output_dir=tmpdir,
                                 extra_postargs=extra_postargs)
            cc.link_executable(objects, os.path.join(tmpdir, "a.out"))
        except Exception as e:
            return False
        return True
    finally:
        if oldstderr is not None:
            os.dup2(oldstderr, sys.stderr.fileno())
        if devnull is not None:
            devnull.close()
        shutil.rmtree(tmpdir)


def detect_openmp():
    "Does this compiler support OpenMP parallelization?"
    compiler = new_compiler()
    print "Attempting to autodetect OpenMP support...",
    hasopenmp = hasfunction(compiler, 'omp_get_num_threads()')
    needs_gomp = hasopenmp
    if not hasopenmp:
        compiler.add_library('gomp')
        hasopenmp = hasfunction(compiler, 'omp_get_num_threads()')
        needs_gomp = hasopenmp
    print
    if hasopenmp:
        print "Compiler supports OpenMP"
    else:
        print "Did not detect OpenMP support; parallel RMSD disabled"
    return hasopenmp, needs_gomp


def detect_sse3():
    "Does this compiler support SSE3 intrinsics?"
    compiler = new_compiler()
    return hasfunction(compiler, '__m128 v; _mm_hadd_ps(v,v)',
                       include='<pmmintrin.h>',
                       extra_postargs=['-msse3'])

def detect_sse4():
    "Does this compiler support SSE4 intrinsics?"
    compiler = new_compiler()
    return hasfunction(compiler, '__m128 v; _mm_round_ps(v,0x00)',
                       include='<smmintrin.h>',
                       extra_postargs=['-msse4'])

################################################################################
# Declaration of the compiled extension modules (cython + c)
################################################################################


xtc = Extension('mdtraj.xtc',
                sources=['MDTraj/xtc/src/xdrfile.c',
                         'MDTraj/xtc/src/xdrfile_xtc.c',
                         'MDTraj/xtc/xtc.pyx'],
                include_dirs=["MDTraj/xtc/include/",
                              'MDTraj/xtc/', numpy.get_include()])

trr = Extension('mdtraj.trr',
                sources=['MDTraj/xtc/src/xdrfile.c',
                         'MDTraj/xtc/src/xdrfile_trr.c',
                         'MDTraj/xtc/trr.pyx'],
                include_dirs=["MDTraj/xtc/include/",
                              'MDTraj/xtc/', numpy.get_include()])

dcd = Extension('mdtraj.dcd',
                sources=["MDTraj/dcd/src/dcdplugin.c", "MDTraj/dcd/dcd.pyx"],
                libraries=['m'],
                include_dirs=["MDTraj/dcd/include/",
                              'MDTraj/dcd/', numpy.get_include()])

binpos = Extension('mdtraj.binpos',
                   sources=['MDTraj/binpos/src/binposplugin.c',
                            'MDTraj/binpos/binpos.pyx'],
                   include_dirs=["MDTraj/binpos/include/",
                                 'MDTraj/binpos/', numpy.get_include()])


def rmsd_extension():
    openmp_enabled, needs_gomp = detect_openmp()
    compiler_args = ['-msse2' if not detect_sse3() else '-mssse3',
                     '--std=gnu99', '-O3', '-funroll-loops']
    if openmp_enabled:
        compiler_args.append('-fopenmp')
    compiler_libraries = ['gomp'] if needs_gomp else []
    #compiler_defs = [('USE_OPENMP', None)] if openmp_enabled else []

    rmsd = Extension('mdtraj._rmsd',
                     sources=[
                         'MDTraj/rmsd/src/theobald_rmsd.c', 'MDTraj/rmsd/_rmsd.pyx'],
                     include_dirs=[
                         "MDTraj/rmsd/include", numpy.get_include()],
                     extra_compile_args=compiler_args,
                     # define_macros=compiler_defs,
                     libraries=compiler_libraries)
    return rmsd


def libgeometry():
    if not detect_sse3():
        return None

    extra_compile_args = ['-mssse3']
    compiler_defs = []
    if detect_sse4():
        extra_compile_args.append('-msse4')

    return Extension('mdtraj.geometry.libgeometry',
                     sources=['MDTraj/geometry/src/geometry.c',
                              'MDTraj/geometry/src/sasa.c'],
                     include_dirs=['MDTraj/geometry/include'],
                     extra_compile_args=extra_compile_args)

extensions = [xtc, trr, dcd, binpos, rmsd_extension()]
ext = libgeometry()
if ext is not None:
    extensions.append(ext)

write_version_py()
setup(name='mdtraj',
      author='Robert McGibbon',
      author_email='rmcgibbo@gmail.com',
      description=DOCLINES[0],
      long_description="\n".join(DOCLINES[2:]),
      version=__version__,
      license='GPLv3+',
      url="http://rmcgibbo.github.io/mdtraj",
      platforms=["Linux", "Mac OS-X", "Unix"],
      classifiers=CLASSIFIERS.splitlines(),
      packages=['mdtraj', 'mdtraj.pdb', 'mdtraj.testing', 'mdtraj.utils',
                'mdtraj.reporters', 'mdtraj.geometry', 'mdtraj.tests'],
      package_dir={'mdtraj': 'MDTraj'},
      install_requires=['numpy', 'cython', 'nose', 'nose-exclude'],
      zip_safe=False,
      scripts=['scripts/mdconvert', 'scripts/mdinspect'],
      ext_modules=extensions,
      cmdclass={'build_ext': build_ext},
      package_data={'mdtraj.pdb': ['data/*'],
                    'mdtraj.testing': ["reference/*"]})
