# coding=utf-8
import os
import pathlib
import subprocess
import sys
import textwrap

from setuptools import Command, Extension, setup


def define_extensions(use_openmp):
    compile_args = []
    if not os.environ.get("LIGHTFM_NO_CFLAGS"):
        compile_args += ["-ffast-math", "-fno-finite-math-only"]

        if sys.platform.startswith("darwin"):
            compile_args += []
        else:
            compile_args += ["-march=native"]

    no_openmp_ext = Extension(
        "lightfm._lightfm_fast_no_openmp",
        ["lightfm/_lightfm_fast_no_openmp.c"],
        extra_compile_args=compile_args,
    )

    if not use_openmp:
        print("Compiling without OpenMP support.")
        return [no_openmp_ext]
    else:
        openmp_ext = Extension(
            "lightfm._lightfm_fast_openmp",
            ["lightfm/_lightfm_fast_openmp.c"],
            extra_link_args=["-fopenmp"],
            extra_compile_args=compile_args + ["-fopenmp"],
        )
        print("Compiling with OpenMP support and no-OpenMP fallback.")
        return [no_openmp_ext, openmp_ext]


class Cythonize(Command):
    """
    Compile the extension .pyx files.
    """

    user_options = []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def generate_pyx(self):
        openmp_import = textwrap.dedent(
            """
             from cython.parallel import parallel, prange
             cimport openmp
        """
        )

        lock_init = textwrap.dedent(
            """
             cdef openmp.omp_lock_t THREAD_LOCK
             openmp.omp_init_lock(&THREAD_LOCK)
        """
        )

        params = (
            (
                "no_openmp",
                dict(
                    openmp_import="",
                    nogil_block="with nogil:",
                    range_block="range",
                    thread_num="0",
                    lock_init="",
                    lock_acquire="",
                    lock_release="",
                ),
            ),
            (
                "openmp",
                dict(
                    openmp_import=openmp_import,
                    nogil_block="with nogil, parallel(num_threads=num_threads):",
                    range_block="prange",
                    thread_num="openmp.omp_get_thread_num()",
                    lock_init=lock_init,
                    lock_acquire="openmp.omp_set_lock(&THREAD_LOCK)",
                    lock_release="openmp.omp_unset_lock(&THREAD_LOCK)",
                ),
            ),
        )

        file_dir = os.path.join(os.path.dirname(__file__), "lightfm")

        with open(os.path.join(file_dir, "_lightfm_fast.pyx.template"), "r") as fl:
            template = fl.read()

        for variant, template_params in params:
            with open(
                os.path.join(file_dir, "_lightfm_fast_{}.pyx".format(variant)), "w"
            ) as fl:
                fl.write(template.format(**template_params))

    def run(self):
        from Cython.Build import cythonize

        self.generate_pyx()

        cythonize(
            [
                Extension(
                    "lightfm._lightfm_fast_no_openmp",
                    ["lightfm/_lightfm_fast_no_openmp.pyx"],
                ),
                Extension(
                    "lightfm._lightfm_fast_openmp",
                    ["lightfm/_lightfm_fast_openmp.pyx"],
                    extra_link_args=["-fopenmp"],
                ),
            ],
            compiler_directives={'language_level' : "3"}
        )


class Clean(Command):
    """
    Clean build files.
    """

    user_options = [("all", None, "(Compatibility with original clean command)")]

    def initialize_options(self):
        self.all = False

    def finalize_options(self):
        pass

    def run(self):
        pth = os.path.dirname(os.path.abspath(__file__))

        subprocess.call(["rm", "-rf", os.path.join(pth, "build")])
        subprocess.call(["rm", "-rf", os.path.join(pth, "lightfm.egg-info")])
        subprocess.call(["find", pth, "-name", "lightfm*.pyc", "-type", "f", "-delete"])
        subprocess.call(["rm", os.path.join(pth, "lightfm", "_lightfm_fast.so")])


def read_version():
    mod = {}
    path = os.path.join(
        os.path.dirname(__file__),
        "lightfm",
        "version.py",
    )
    with open(path) as fd:
        exec(fd.read(), mod)
    return mod["__version__"]


use_openmp = not sys.platform.startswith("darwin") and not sys.platform.startswith(
    "win"
)

long_description = pathlib.Path(__file__).parent.joinpath("README.md").read_text()

setup(
    cmdclass={"cythonize": Cythonize, "clean": Clean},
    ext_modules=define_extensions(use_openmp),
)
