# -*- coding: utf-8 -*-
"""
=====================
Cython related magics
=====================

Magic command interface for interactive work with Cython

.. note::

  The ``Cython`` package needs to be installed separately. It
  can be obtained using ``easy_install`` or ``pip``.

Usage
=====

To enable the magics below, execute ``%load_ext cythonmagic``.

``%%cython``

{CYTHON_DOC}

``%%cython_inline``

{CYTHON_INLINE_DOC}

``%%cython_pyximport``

{CYTHON_PYXIMPORT_DOC}

Author:
* Brian Granger

Code moved from IPython and adapted by:
* Martín Gaitán

Parts of this code were taken from Cython.inline.
"""
#-----------------------------------------------------------------------------
# Copyright (C) 2010-2011, IPython Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file ipython-COPYING.rst, distributed with this software.
#-----------------------------------------------------------------------------

from __future__ import absolute_import, print_function

import imp
import io
import os
import re
import sys
import time

try:
    reload
except NameError:   # Python 3
    from imp import reload

try:
    import hashlib
except ImportError:
    import md5 as hashlib

from distutils.core import Distribution, Extension
from distutils.command.build_ext import build_ext

from IPython.core import display
from IPython.core import magic_arguments
from IPython.core.magic import Magics, magics_class, cell_magic
from IPython.utils import py3compat
try:
    from IPython.paths import get_ipython_cache_dir
except ImportError:
    # older IPython version
    from IPython.utils.path import get_ipython_cache_dir
from IPython.utils.text import dedent

from ..Shadow import __version__ as cython_version
from ..Compiler.Errors import CompileError
from .Inline import cython_inline
from .Dependencies import cythonize


@magics_class
class CythonMagics(Magics):

    def __init__(self, shell):
        super(CythonMagics, self).__init__(shell)
        self._reloads = {}
        self._code_cache = {}
        self._pyximport_installed = False

    def _import_all(self, module):
        mdict = module.__dict__
        if '__all__' in mdict:
            keys = mdict['__all__']
        else:
            keys = [k for k in mdict if not k.startswith('_')]

        for k in keys:
            try:
                self.shell.push({k: mdict[k]})
            except KeyError:
                msg = "'module' object has no attribute '%s'" % k
                raise AttributeError(msg)

    @cell_magic
    def cython_inline(self, line, cell):
        """Compile and run a Cython code cell using Cython.inline.

        This magic simply passes the body of the cell to Cython.inline
        and returns the result. If the variables `a` and `b` are defined
        in the user's namespace, here is a simple example that returns
        their sum::

            %%cython_inline
            return a+b

        For most purposes, we recommend the usage of the `%%cython` magic.
        """
        locs = self.shell.user_global_ns
        globs = self.shell.user_ns
        return cython_inline(cell, locals=locs, globals=globs)

    @cell_magic
    def cython_pyximport(self, line, cell):
        """Compile and import a Cython code cell using pyximport.

        The contents of the cell are written to a `.pyx` file in the current
        working directory, which is then imported using `pyximport`. This
        magic requires a module name to be passed::

            %%cython_pyximport modulename
            def f(x):
                return 2.0*x

        The compiled module is then imported and all of its symbols are
        injected into the user's namespace. For most purposes, we recommend
        the usage of the `%%cython` magic.
        """
        module_name = line.strip()
        if not module_name:
            raise ValueError('module name must be given')
        fname = module_name + '.pyx'
        with io.open(fname, 'w', encoding='utf-8') as f:
            f.write(cell)
        if 'pyximport' not in sys.modules or not self._pyximport_installed:
            import pyximport
            pyximport.install(reload_support=True)
            self._pyximport_installed = True
        if module_name in self._reloads:
            module = self._reloads[module_name]
            reload(module)
        else:
            __import__(module_name)
            module = sys.modules[module_name]
            self._reloads[module_name] = module
        self._import_all(module)

    @magic_arguments.magic_arguments()
    @magic_arguments.argument(
        '-3', dest='language_level', action='store_const', const=3, default=None,
        help="Select Python 3 syntax."
    )
    @magic_arguments.argument(
        '-2', dest='language_level', action='store_const', const=2, default=None,
        help="Select Python 2 syntax."
    )
    @magic_arguments.argument(
        '-c', '--compile-args', action='append', default=[],
        help="Extra flags to pass to compiler via the `extra_compile_args` "
             "Extension flag (can be specified  multiple times)."
    )
    @magic_arguments.argument(
        '--link-args', action='append', default=[],
        help="Extra flags to pass to linker via the `extra_link_args` "
             "Extension flag (can be specified  multiple times)."
    )
    @magic_arguments.argument(
        '-l', '--lib', action='append', default=[],
        help="Add a library to link the extension against (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-n', '--name',
        help="Specify a name for the Cython module."
    )
    @magic_arguments.argument(
        '-L', dest='library_dirs', metavar='dir', action='append', default=[],
        help="Add a path to the list of library directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-I', '--include', action='append', default=[],
        help="Add a path to the list of include directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-S', '--src', action='append', default=[],
        help="Add a path to the list of src files (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-+', '--cplus', action='store_true', default=False,
        help="Output a C++ rather than C file."
    )
    @magic_arguments.argument(
        '-f', '--force', action='store_true', default=False,
        help="Force the compilation of a new module, even if the source has been "
             "previously compiled."
    )
    @magic_arguments.argument(
        '-a', '--annotate', action='store_true', default=False,
        help="Produce a colorized HTML version of the source."
    )
    @cell_magic
    def cython(self, line, cell):
        """Compile and import everything from a Cython code cell.

        The contents of the cell are written to a `.pyx` file in the
        directory `IPYTHONDIR/cython` using a filename with the hash of the
        code. This file is then cythonized and compiled. The resulting module
        is imported and all of its symbols are injected into the user's
        namespace. The usage is similar to that of `%%cython_pyximport` but
        you don't have to pass a module name::

            %%cython
            def f(x):
                return 2.0*x

        To compile OpenMP codes, pass the required  `--compile-args`
        and `--link-args`.  For example with gcc::

            %%cython --compile-args=-fopenmp --link-args=-fopenmp
            ...
        """
        args = magic_arguments.parse_argstring(self.cython, line)
        code = cell if cell.endswith('\n') else cell + '\n'
        lib_dir = os.path.join(get_ipython_cache_dir(), 'cython')
        quiet = True
        key = code, line, sys.version_info, sys.executable, cython_version

        if not os.path.exists(lib_dir):
            os.makedirs(lib_dir)

        if args.force:
            # Force a new module name by adding the current time to the
            # key which is hashed to determine the module name.
            key += time.time(),

        if args.name:
            module_name = py3compat.unicode_to_str(args.name)
        else:
            module_name = "_cython_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()
        module_path = os.path.join(lib_dir, module_name + self.so_ext)

        have_module = os.path.isfile(module_path)
        need_cythonize = not have_module

        if args.annotate:
            html_file = os.path.join(lib_dir, module_name + '.html')
            if not os.path.isfile(html_file):
                need_cythonize = True

        if need_cythonize:
            c_include_dirs = args.include
            c_src_files = list(map(str, args.src))
            if 'numpy' in code:
                import numpy
                c_include_dirs.append(numpy.get_include())
            pyx_file = os.path.join(lib_dir, module_name + '.pyx')
            pyx_file = py3compat.cast_bytes_py2(pyx_file, encoding=sys.getfilesystemencoding())
            with io.open(pyx_file, 'w', encoding='utf-8') as f:
                f.write(code)
            extension = Extension(
                name=module_name,
                sources=[pyx_file] + c_src_files,
                include_dirs=c_include_dirs,
                library_dirs=args.library_dirs,
                extra_compile_args=args.compile_args,
                extra_link_args=args.link_args,
                libraries=args.lib,
                language='c++' if args.cplus else 'c',
            )
            build_extension = self._get_build_extension()
            try:
                opts = dict(
                    quiet=quiet,
                    annotate=args.annotate,
                    force=True,
                )
                if args.language_level is not None:
                    assert args.language_level in (2, 3)
                    opts['language_level'] = args.language_level
                elif sys.version_info[0] > 2:
                    opts['language_level'] = 3
                build_extension.extensions = cythonize([extension], **opts)
            except CompileError:
                return

        if not have_module:
            build_extension.build_temp = os.path.dirname(pyx_file)
            build_extension.build_lib = lib_dir
            build_extension.run()
            self._code_cache[key] = module_name

        module = imp.load_dynamic(module_name, module_path)
        self._import_all(module)

        if args.annotate:
            try:
                with io.open(html_file, encoding='utf-8') as f:
                    annotated_html = f.read()
            except IOError as e:
                # File could not be opened. Most likely the user has a version
                # of Cython before 0.15.1 (when `cythonize` learned the
                # `force` keyword argument) and has already compiled this
                # exact source without annotation.
                print('Cython completed successfully but the annotated '
                      'source could not be read.', file=sys.stderr)
                print(e, file=sys.stderr)
            else:
                return display.HTML(self.clean_annotated_html(annotated_html))

    @property
    def so_ext(self):
        """The extension suffix for compiled modules."""
        try:
            return self._so_ext
        except AttributeError:
            self._so_ext = self._get_build_extension().get_ext_filename('')
            return self._so_ext

    def _clear_distutils_mkpath_cache(self):
        """clear distutils mkpath cache

        prevents distutils from skipping re-creation of dirs that have been removed
        """
        try:
            from distutils.dir_util import _path_created
        except ImportError:
            pass
        else:
            _path_created.clear()

    def _get_build_extension(self):
        self._clear_distutils_mkpath_cache()
        dist = Distribution()
        config_files = dist.find_config_files()
        try:
            config_files.remove('setup.cfg')
        except ValueError:
            pass
        dist.parse_config_files(config_files)
        build_extension = build_ext(dist)
        build_extension.finalize_options()
        return build_extension

    @staticmethod
    def clean_annotated_html(html):
        """Clean up the annotated HTML source.

        Strips the link to the generated C or C++ file, which we do not
        present to the user.
        """
        r = re.compile('<p>Raw output: <a href="(.*)">(.*)</a>')
        html = '\n'.join(l for l in html.splitlines() if not r.match(l))
        return html

__doc__ = __doc__.format(
    # rST doesn't see the -+ flag as part of an option list, so we
    # hide it from the module-level docstring.
    CYTHON_DOC=dedent(CythonMagics.cython.__doc__\
                                  .replace('-+, --cplus', '--cplus    ')),
    CYTHON_INLINE_DOC=dedent(CythonMagics.cython_inline.__doc__),
    CYTHON_PYXIMPORT_DOC=dedent(CythonMagics.cython_pyximport.__doc__),
)
