#! /usr/bin/env python
# System imports
from setuptools import setup, Extension
import os
import subprocess
import platform
import sys
import warnings
import glob
import tempfile
import textwrap
import subprocess

try:
    from Cython.Build import cythonize
except ImportError:
    USE_CYTHON = False
    warnings.warn(
        "Cython was not found. Slow reflectivity calculations will be used."
    )
else:
    USE_CYTHON = True


# versioning
MAJOR = 0
MINOR = 1
MICRO = 60
ISRELEASED = True
VERSION = f"{MAJOR}.{MINOR}.{MICRO}"


# are we on windows, darwin, etc?
platform = sys.platform
# packages = find_packages()
# try:
#     idx = packages.index("benchmarks")
#     if idx >= 0:
#         packages.pop(idx)
#     idx = packages.index("benchmarks.benchmarks")
#     if idx >= 0:
#         packages.pop(idx)
#     idx = packages.index("motofit")
#     if idx >= 0:
#         packages.pop(idx)
# except ValueError:
#     pass


# Return the git revision as a string
def git_version():
    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 get_version_info():
    # Adding the git rev number needs to be done inside
    # write_version_py(), otherwise the import of refnx.version messes
    # up the build under Python 3.
    FULLVERSION = VERSION
    if os.path.exists(".git"):
        GIT_REVISION = git_version()
    elif os.path.exists("refnx/version.py"):
        # must be a source distribution, use existing version file
        # load it as a separate module to not load refnx/__init__.py
        import importlib.util

        spec = importlib.util.spec_from_file_location(
            "version", "refnx/version.py"
        )
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        GIT_REVISION = module.git_revision
    else:
        GIT_REVISION = "Unknown"

    if not ISRELEASED:
        FULLVERSION += ".dev0+" + GIT_REVISION[:7]

    return FULLVERSION, GIT_REVISION


def write_version_py(filename="refnx/version.py"):
    cnt = """
# THIS FILE IS GENERATED FROM REFNX 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
"""
    FULLVERSION, GIT_REVISION = get_version_info()

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


###############################################################################
"""
Is openMP usable?
"""
CCODE = textwrap.dedent(
    """\
    #include <omp.h>
    #include <stdio.h>
    int main(void) {
    #pragma omp parallel
    printf("nthreads=%d\\n", omp_get_num_threads());
    return 0;
    }
    """
)


def get_openmp_flag(compiler):
    if hasattr(compiler, "compiler"):
        compiler = compiler.compiler[0]
    else:
        compiler = compiler.__class__.__name__

    if sys.platform == "win32" and ("icc" in compiler or "icl" in compiler):
        return ["/Qopenmp"], []
    elif sys.platform == "win32":
        return ["/openmp"], []
    elif sys.platform == "darwin" and ("icc" in compiler or "icl" in compiler):
        return ["-openmp"], []
    elif sys.platform == "darwin":
        # default for macOS, assuming Apple-clang
        # -fopenmp can't be passed as compile flag when using Apple-clang.
        # OpenMP support has to be enabled during preprocessing.
        #
        # it may be possible that someone builds with a different/updated
        # compiler (don't know how to check for that).
        #
        # set the following environment variables, assumes that llvm openmp
        # has been built and installed by the user.
        #
        # brew install libomp
        # export CC=clang
        # export CXX =clang++
        # export CXXFLAGS="$CXXFLAGS -Xpreprocessor -fopenmp"
        # export CFLAGS="$CFLAGS -I/usr/local/opt/libomp/include"
        # export CXXFLAGS="$CXXFLAGS -I/usr/local/opt/libomp/include"
        # export LDFLAGS="$LDFLAGS -L/usr/local/opt/libomp/lib -lomp"
        # export DYLD_LIBRARY_PATH =/usr/local/opt/libomp/lib
        return ["-fopenmp"], ["-lomp"]
    # Default flag for GCC and clang:
    return ["-fopenmp"], []


def check_openmp_support():
    """Check whether OpenMP test code can be compiled and run"""

    try:
        from setuptools._distutils.ccompiler import new_compiler
        from setuptools._distutils.sysconfig import customize_compiler

        # from numpy.distutils.ccompiler import new_compiler
        # from distutils.sysconfig import customize_compiler
        from distutils.errors import CompileError, LinkError
    except ImportError:
        return False

    ccompiler = new_compiler()
    customize_compiler(ccompiler)

    start_dir = os.path.abspath(".")

    with tempfile.TemporaryDirectory() as tmp_dir:
        try:
            os.chdir(tmp_dir)

            # Write test program
            with open("test_openmp.c", "w") as f:
                f.write(CCODE)

            os.mkdir("objects")

            # Compile, test program
            openmp_flags = get_openmp_flag(ccompiler)

            ccompiler.compile(
                ["test_openmp.c"],
                output_dir="objects",
                extra_postargs=openmp_flags[0],
            )

            # Link test program
            extra_preargs = os.getenv("LDFLAGS", None)
            if extra_preargs is not None:
                extra_preargs = extra_preargs.split(" ")
            else:
                extra_preargs = []

            objects = glob.glob(
                os.path.join("objects", "*" + ccompiler.obj_extension)
            )
            ccompiler.link_executable(
                objects,
                "test_openmp",
                extra_preargs=extra_preargs,
                extra_postargs=openmp_flags[1],
            )

            # Run test program
            output = subprocess.check_output("./test_openmp")
            output = output.decode(sys.stdout.encoding or "utf-8").splitlines()

            # Check test program output
            if "nthreads=" in output[0]:
                nthreads = int(output[0].strip().split("=")[1])
                openmp_supported = len(output) == nthreads
            else:
                openmp_supported = False

        except (CompileError, LinkError, subprocess.CalledProcessError):
            openmp_supported = False

        finally:
            os.chdir(start_dir)

    return openmp_supported


# do you want to parallelise things with openmp?
HAS_OPENMP = check_openmp_support()
# HAS_OPENMP = False
###############################################################################


# refnx setup
info = {
    "include_package_data": True,
}


####################################################################
# this is where setup starts
####################################################################
def setup_package():

    # Rewrite the version file every time
    write_version_py()
    info["version"] = get_version_info()[0]
    print(info["version"])

    if USE_CYTHON:
        # Obtain the numpy include directory.  This logic works across numpy
        # versions.
        ext_modules = []
        HAS_NUMPY = True

        try:
            import numpy as np
        except:
            info["setup_requires"] = ["numpy"]
            HAS_NUMPY = False

        if HAS_NUMPY:
            try:
                numpy_include = np.get_include()
            except AttributeError:
                numpy_include = np.get_numpy_include()

            _cevent = Extension(
                name="refnx.reduce._cevent",
                sources=["src/_cevent.pyx"],
                include_dirs=[numpy_include],
                language="c++",
                define_macros=[
                    ("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")
                ],
                # libraries=
                # extra_compile_args = "...".split(),
            )
            ext_modules.append(_cevent)

            _cutil = Extension(
                name="refnx._lib._cutil",
                sources=["src/_cutil.pyx"],
                include_dirs=[numpy_include],
                language="c",
                # libraries=
                # extra_compile_args = "...".split(),
            )
            ext_modules.append(_cutil)

            # creflect extension module
            # Compile reflectivity calculator to object with C compiler
            # first.
            # It's not possible to do this in an Extension object because
            # the `-std=c++11` compile arg and C99 C code are incompatible
            # (at least on Darwin).
            from setuptools._distutils.ccompiler import new_compiler
            from setuptools._distutils.sysconfig import customize_compiler

            # from numpy.distutils.ccompiler import new_compiler
            # from distutils.sysconfig import customize_compiler

            ccompiler = new_compiler()
            customize_compiler(ccompiler)
            ccompiler.verbose = True
            extra_preargs = [
                "-O3",
            ]

            if sys.platform == "win32":
                # use the C++ code on Windows. The C++ code uses the
                # std::complex<double> object for its arithmetic.
                f = ["src/refcalc.cpp"]
            else:
                # and C code on other machines. The C code uses C99 complex
                # arithmetic which is 10-20% faster.
                # the CMPLX macro was only standardised in C11
                extra_preargs.extend(
                    [
                        "-std=c17",
                    ]
                )
                f = ["src/refcalc.c"]
            refcalc_obj = ccompiler.compile(f, extra_preargs=extra_preargs)
            # print(refcalc_obj)

            _creflect = Extension(
                name="refnx.reflect._creflect",
                sources=[
                    "src/_creflect.pyx",
                    "src/refcaller.cpp",
                    "src/pnr/magnetic.cc",
                ],
                include_dirs=[numpy_include],
                language="c++",
                extra_compile_args=["-std=c++11"],
                define_macros=[
                    ("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")
                ],
                extra_objects=refcalc_obj,
            )
            ext_modules.append(_creflect)

            # if we have openmp use pure cython version
            # openmp should be present on windows, linux
            #
            # However, it's not present in Apple Clang. Therefore one has to
            # jump through hoops to enable it.
            # It's probably easier to install OpenMP on macOS via homebrew
            #
            # brew install llvm
            # export PATH="/opt/homebrew/opt/llvm/bin:$PATH"
            # export CPPFLAGS="-I/opt/homebrew/opt/llvm/include"
            # export LDFLAGS="-L/opt/homebrew/opt/llvm/lib/c++ -Wl,-rpath,/opt/homebrew/opt/llvm/lib/c++ -L/opt/homebrew/opt/llvm/lib"
            print(f"{HAS_OPENMP=}, {refcalc_obj=}")
            if HAS_OPENMP:
                # cyreflect extension module
                _cyreflect = Extension(
                    name="refnx.reflect._cyreflect",
                    sources=[
                        "src/_cyreflect.pyx",
                        "src/refcaller.cpp",
                        "src/pnr/magnetic.cc",
                    ],
                    include_dirs=[numpy_include],
                    language="c++",
                    extra_compile_args=[],
                    extra_link_args=[],
                    define_macros=[],
                    extra_objects=refcalc_obj,
                    # extra_compile_args = "...".split(),
                )
                openmp_flags = get_openmp_flag(ccompiler)
                _cyreflect.extra_compile_args += openmp_flags[0]
                _cyreflect.extra_link_args += openmp_flags[1]

                ext_modules.append(_cyreflect)

            # specify min deployment version for macOS
            # if platform == "darwin":
            #     for mod in ext_modules:
            #         mod.extra_compile_args.append("-mmacosx-version-min=10.9")

            info["ext_modules"] = cythonize(ext_modules)
            info["zip_safe"] = False

    try:
        setup(**info)
    except ValueError:
        # there probably wasn't a C-compiler (windows). Try removing extension
        # compilation
        print("")
        print("*****WARNING*****")
        print(
            "You didn't try to build the Reflectivity calculation extension."
            " Calculation will be slow, falling back to pure python."
            " To compile extension install cython. If installing in windows you"
            " should then install from Visual Studio command prompt (this makes"
            " C compiler available"
        )
        print("*****************")
        print("")
        info.pop("cmdclass")
        info.pop("ext_modules")
        setup(**info)


if __name__ == "__main__":
    setup_package()