"""
A special directive for generating a matplotlib plot.

.. warning::

   This is a hacked version of plot_directive.py from Matplotlib.
   It's very much subject to change!


Usage
-----

Can be used like this::

    .. plot:: examples/example.py

    .. plot::

       import matplotlib.pyplot as plt
       plt.plot([1,2,3], [4,5,6])

    .. plot::

       A plotting example:

       >>> import matplotlib.pyplot as plt
       >>> plt.plot([1,2,3], [4,5,6])

The content is interpreted as doctest formatted if it has a line starting
with ``>>>``.

The ``plot`` directive supports the options

    format : {'python', 'doctest'}
        Specify the format of the input

    include-source : bool
        Whether to display the source code. Default can be changed in conf.py

and the ``image`` directive options ``alt``, ``height``, ``width``,
``scale``, ``align``, ``class``.

Configuration options
---------------------

The plot directive has the following configuration options:

    plot_include_source
        Default value for the include-source option

    plot_pre_code
        Code that should be executed before each plot.

    plot_basedir
        Base directory, to which plot:: file names are relative to.
        (If None or empty, file names are relative to the directoly where
        the file containing the directive is.)

    plot_formats
        File formats to generate. List of tuples or strings::

            [(suffix, dpi), suffix, ...]

        that determine the file format and the DPI. For entries whose
        DPI was omitted, sensible defaults are chosen.

    plot_html_show_formats
        Whether to show links to the files in HTML.

TODO
----

* Refactor Latex output; now it's plain images, but it would be nice
  to make them appear side-by-side, or in floats.

"""
from __future__ import division, absolute_import, print_function

import sys, os, glob, shutil, imp, warnings, re, textwrap, traceback
import sphinx

if sys.version_info[0] >= 3:
    from io import StringIO
else:
    from io import StringIO

import warnings

warnings.warn(
    "A plot_directive module is also available under "
    "matplotlib.sphinxext; expect this numpydoc.plot_directive "
    "module to be deprecated after relevant features have been "
    "integrated there.",
    FutureWarning,
    stacklevel=2,
)


# ------------------------------------------------------------------------------
# Registration hook
# ------------------------------------------------------------------------------


def setup(app):
    setup.app = app
    setup.config = app.config
    setup.confdir = app.confdir

    app.add_config_value("plot_pre_code", "", True)
    app.add_config_value("plot_include_source", False, True)
    app.add_config_value("plot_formats", ["png", "hires.png", "pdf"], True)
    app.add_config_value("plot_basedir", None, True)
    app.add_config_value("plot_html_show_formats", True, True)

    app.add_directive(
        "plot", plot_directive, True, (0, 1, False), **plot_directive_options
    )


# ------------------------------------------------------------------------------
# plot:: directive
# ------------------------------------------------------------------------------
from docutils.parsers.rst import directives
from docutils import nodes


def plot_directive(
    name,
    arguments,
    options,
    content,
    lineno,
    content_offset,
    block_text,
    state,
    state_machine,
):
    return run(arguments, content, options, state_machine, state, lineno)


plot_directive.__doc__ = __doc__


def _option_boolean(arg):
    if not arg or not arg.strip():
        # no argument given, assume used as a flag
        return True
    elif arg.strip().lower() in ("no", "0", "false"):
        return False
    elif arg.strip().lower() in ("yes", "1", "true"):
        return True
    else:
        raise ValueError('"%s" unknown boolean' % arg)


def _option_format(arg):
    return directives.choice(arg, ("python", "lisp"))


def _option_align(arg):
    return directives.choice(
        arg, ("top", "middle", "bottom", "left", "center", "right")
    )


plot_directive_options = {
    "alt": directives.unchanged,
    "height": directives.length_or_unitless,
    "width": directives.length_or_percentage_or_unitless,
    "scale": directives.nonnegative_int,
    "align": _option_align,
    "class": directives.class_option,
    "include-source": _option_boolean,
    "format": _option_format,
}

# ------------------------------------------------------------------------------
# Generating output
# ------------------------------------------------------------------------------

from docutils import nodes, utils

try:
    # Sphinx depends on either Jinja or Jinja2
    import jinja2

    def format_template(template, **kw):
        return jinja2.Template(template).render(**kw)


except ImportError:
    import jinja

    def format_template(template, **kw):
        return jinja.from_string(template, **kw)


TEMPLATE = """
{{ source_code }}

{{ only_html }}

   {% if source_link or (html_show_formats and not multi_image) %}
   (
   {%- if source_link -%}
   `Source code <{{ source_link }}>`__
   {%- endif -%}
   {%- if html_show_formats and not multi_image -%}
     {%- for img in images -%}
       {%- for fmt in img.formats -%}
         {%- if source_link or not loop.first -%}, {% endif -%}
         `{{ fmt }} <{{ dest_dir }}/{{ img.basename }}.{{ fmt }}>`__
       {%- endfor -%}
     {%- endfor -%}
   {%- endif -%}
   )
   {% endif %}

   {% for img in images %}
   .. figure:: {{ build_dir }}/{{ img.basename }}.png
      {%- for option in options %}
      {{ option }}
      {% endfor %}

      {% if html_show_formats and multi_image -%}
        (
        {%- for fmt in img.formats -%}
        {%- if not loop.first -%}, {% endif -%}
        `{{ fmt }} <{{ dest_dir }}/{{ img.basename }}.{{ fmt }}>`__
        {%- endfor -%}
        )
      {%- endif -%}
   {% endfor %}

{{ only_latex }}

   {% for img in images %}
   .. image:: {{ build_dir }}/{{ img.basename }}.pdf
   {% endfor %}

"""


class ImageFile(object):
    def __init__(self, basename, dirname):
        self.basename = basename
        self.dirname = dirname
        self.formats = []

    def filename(self, format):
        return os.path.join(self.dirname, "%s.%s" % (self.basename, format))

    def filenames(self):
        return [self.filename(fmt) for fmt in self.formats]


def run(arguments, content, options, state_machine, state, lineno):
    if arguments and content:
        raise RuntimeError("plot:: directive can't have both args and content")

    document = state_machine.document
    config = document.settings.env.config

    options.setdefault("include-source", config.plot_include_source)

    # determine input
    rst_file = document.attributes["source"]
    rst_dir = os.path.dirname(rst_file)

    if arguments:
        if not config.plot_basedir:
            source_file_name = os.path.join(rst_dir, directives.uri(arguments[0]))
        else:
            source_file_name = os.path.join(
                setup.confdir, config.plot_basedir, directives.uri(arguments[0])
            )
        code = open(source_file_name, "r").read()
        output_base = os.path.basename(source_file_name)
    else:
        source_file_name = rst_file
        code = textwrap.dedent("\n".join(map(str, content)))
        counter = document.attributes.get("_plot_counter", 0) + 1
        document.attributes["_plot_counter"] = counter
        base, ext = os.path.splitext(os.path.basename(source_file_name))
        output_base = "%s-%d.py" % (base, counter)

    base, source_ext = os.path.splitext(output_base)
    if source_ext in (".py", ".rst", ".txt"):
        output_base = base
    else:
        source_ext = ""

    # ensure that LaTeX includegraphics doesn't choke in foo.bar.pdf filenames
    output_base = output_base.replace(".", "-")

    # is it in doctest format?
    is_doctest = contains_doctest(code)
    if "format" in options:
        if options["format"] == "python":
            is_doctest = False
        else:
            is_doctest = True

    # determine output directory name fragment
    source_rel_name = relpath(source_file_name, setup.confdir)
    source_rel_dir = os.path.dirname(source_rel_name)
    while source_rel_dir.startswith(os.path.sep):
        source_rel_dir = source_rel_dir[1:]

    # build_dir: where to place output files (temporarily)
    build_dir = os.path.join(
        os.path.dirname(setup.app.doctreedir), "plot_directive", source_rel_dir
    )
    if not os.path.exists(build_dir):
        os.makedirs(build_dir)

    # output_dir: final location in the builder's directory
    dest_dir = os.path.abspath(os.path.join(setup.app.builder.outdir, source_rel_dir))

    # how to link to files from the RST file
    dest_dir_link = os.path.join(
        relpath(setup.confdir, rst_dir), source_rel_dir
    ).replace(os.path.sep, "/")
    build_dir_link = relpath(build_dir, rst_dir).replace(os.path.sep, "/")
    source_link = dest_dir_link + "/" + output_base + source_ext

    # make figures
    try:
        results = makefig(code, source_file_name, build_dir, output_base, config)
        errors = []
    except PlotError as err:
        reporter = state.memo.reporter
        sm = reporter.system_message(
            2, "Exception occurred in plotting %s: %s" % (output_base, err), line=lineno
        )
        results = [(code, [])]
        errors = [sm]

    # generate output restructuredtext
    total_lines = []
    for j, (code_piece, images) in enumerate(results):
        if options["include-source"]:
            if is_doctest:
                lines = [""]
                lines += [row.rstrip() for row in code_piece.split("\n")]
            else:
                lines = [".. code-block:: python", ""]
                lines += ["    %s" % row.rstrip() for row in code_piece.split("\n")]
            source_code = "\n".join(lines)
        else:
            source_code = ""

        opts = [
            ":%s: %s" % (key, val)
            for key, val in list(options.items())
            if key in ("alt", "height", "width", "scale", "align", "class")
        ]

        only_html = ".. only:: html"
        only_latex = ".. only:: latex"

        if j == 0:
            src_link = source_link
        else:
            src_link = None

        result = format_template(
            TEMPLATE,
            dest_dir=dest_dir_link,
            build_dir=build_dir_link,
            source_link=src_link,
            multi_image=len(images) > 1,
            only_html=only_html,
            only_latex=only_latex,
            options=opts,
            images=images,
            source_code=source_code,
            html_show_formats=config.plot_html_show_formats,
        )

        total_lines.extend(result.split("\n"))
        total_lines.extend("\n")

    if total_lines:
        state_machine.insert_input(total_lines, source=source_file_name)

    # copy image files to builder's output directory
    if not os.path.exists(dest_dir):
        os.makedirs(dest_dir)

    for code_piece, images in results:
        for img in images:
            for fn in img.filenames():
                shutil.copyfile(fn, os.path.join(dest_dir, os.path.basename(fn)))

    # copy script (if necessary)
    if source_file_name == rst_file:
        target_name = os.path.join(dest_dir, output_base + source_ext)
        f = open(target_name, "w")
        f.write(unescape_doctest(code))
        f.close()

    return errors


# ------------------------------------------------------------------------------
# Run code and capture figures
# ------------------------------------------------------------------------------

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.image as image
from matplotlib import _pylab_helpers

import exceptions


def contains_doctest(text):
    try:
        # check if it's valid Python as-is
        compile(text, "<string>", "exec")
        return False
    except SyntaxError:
        pass
    r = re.compile(r"^\s*>>>", re.M)
    m = r.search(text)
    return bool(m)


def unescape_doctest(text):
    """
    Extract code from a piece of text, which contains either Python code
    or doctests.

    """
    if not contains_doctest(text):
        return text

    code = ""
    for line in text.split("\n"):
        m = re.match(r"^\s*(>>>|\.\.\.) (.*)$", line)
        if m:
            code += m.group(2) + "\n"
        elif line.strip():
            code += "# " + line.strip() + "\n"
        else:
            code += "\n"
    return code


def split_code_at_show(text):
    """
    Split code at plt.show()

    """

    parts = []
    is_doctest = contains_doctest(text)

    part = []
    for line in text.split("\n"):
        if (not is_doctest and line.strip() == "plt.show()") or (
            is_doctest and line.strip() == ">>> plt.show()"
        ):
            part.append(line)
            parts.append("\n".join(part))
            part = []
        else:
            part.append(line)
    if "\n".join(part).strip():
        parts.append("\n".join(part))
    return parts


class PlotError(RuntimeError):
    pass


def run_code(code, code_path, ns=None):
    # Change the working directory to the directory of the example, so
    # it can get at its data files, if any.
    pwd = os.getcwd()
    old_sys_path = list(sys.path)
    if code_path is not None:
        dirname = os.path.abspath(os.path.dirname(code_path))
        os.chdir(dirname)
        sys.path.insert(0, dirname)

    # Redirect stdout
    stdout = sys.stdout
    sys.stdout = StringIO()

    # Reset sys.argv
    old_sys_argv = sys.argv
    sys.argv = [code_path]

    try:
        try:
            code = unescape_doctest(code)
            if ns is None:
                ns = {}
            if not ns:
                exec(setup.config.plot_pre_code, ns)
            exec(code, ns)
        except (Exception, SystemExit) as err:
            raise PlotError(traceback.format_exc())
    finally:
        os.chdir(pwd)
        sys.argv = old_sys_argv
        sys.path[:] = old_sys_path
        sys.stdout = stdout
    return ns


# ------------------------------------------------------------------------------
# Generating figures
# ------------------------------------------------------------------------------


def out_of_date(original, derived):
    """
    Returns True if derivative is out-of-date wrt original,
    both of which are full file paths.
    """
    return (
        not os.path.exists(derived)
        or os.stat(derived).st_mtime < os.stat(original).st_mtime
    )


def makefig(code, code_path, output_dir, output_base, config):
    """
    Run a pyplot script *code* and save the images under *output_dir*
    with file names derived from *output_base*

    """

    # -- Parse format list
    default_dpi = {"png": 80, "hires.png": 200, "pdf": 50}
    formats = []
    for fmt in config.plot_formats:
        if isinstance(fmt, str):
            formats.append((fmt, default_dpi.get(fmt, 80)))
        elif type(fmt) in (tuple, list) and len(fmt) == 2:
            formats.append((str(fmt[0]), int(fmt[1])))
        else:
            raise PlotError('invalid image format "%r" in plot_formats' % fmt)

    # -- Try to determine if all images already exist

    code_pieces = split_code_at_show(code)

    # Look for single-figure output files first
    all_exists = True
    img = ImageFile(output_base, output_dir)
    for format, dpi in formats:
        if out_of_date(code_path, img.filename(format)):
            all_exists = False
            break
        img.formats.append(format)

    if all_exists:
        return [(code, [img])]

    # Then look for multi-figure output files
    results = []
    all_exists = True
    for i, code_piece in enumerate(code_pieces):
        images = []
        for j in range(1000):
            img = ImageFile("%s_%02d_%02d" % (output_base, i, j), output_dir)
            for format, dpi in formats:
                if out_of_date(code_path, img.filename(format)):
                    all_exists = False
                    break
                img.formats.append(format)

            # assume that if we have one, we have them all
            if not all_exists:
                all_exists = j > 0
                break
            images.append(img)
        if not all_exists:
            break
        results.append((code_piece, images))

    if all_exists:
        return results

    # -- We didn't find the files, so build them

    results = []
    ns = {}

    for i, code_piece in enumerate(code_pieces):
        # Clear between runs
        plt.close("all")

        # Run code
        run_code(code_piece, code_path, ns)

        # Collect images
        images = []
        fig_managers = _pylab_helpers.Gcf.get_all_fig_managers()
        for j, figman in enumerate(fig_managers):
            if len(fig_managers) == 1 and len(code_pieces) == 1:
                img = ImageFile(output_base, output_dir)
            else:
                img = ImageFile("%s_%02d_%02d" % (output_base, i, j), output_dir)
            images.append(img)
            for format, dpi in formats:
                try:
                    figman.canvas.figure.savefig(img.filename(format), dpi=dpi)
                except exceptions.BaseException as err:
                    raise PlotError(traceback.format_exc())
                img.formats.append(format)

        # Results
        results.append((code_piece, images))

    return results


# ------------------------------------------------------------------------------
# Relative pathnames
# ------------------------------------------------------------------------------

try:
    from os.path import relpath
except ImportError:
    # Copied from Python 2.7
    if "posix" in sys.builtin_module_names:

        def relpath(path, start=os.path.curdir):
            """Return a relative version of a path"""
            from os.path import sep, curdir, join, abspath, commonprefix, pardir

            if not path:
                raise ValueError("no path specified")

            start_list = abspath(start).split(sep)
            path_list = abspath(path).split(sep)

            # Work out how much of the filepath is shared by start and path.
            i = len(commonprefix([start_list, path_list]))

            rel_list = [pardir] * (len(start_list) - i) + path_list[i:]
            if not rel_list:
                return curdir
            return join(*rel_list)

    elif "nt" in sys.builtin_module_names:

        def relpath(path, start=os.path.curdir):
            """Return a relative version of a path"""
            from os.path import (
                sep,
                curdir,
                join,
                abspath,
                commonprefix,
                pardir,
                splitunc,
            )

            if not path:
                raise ValueError("no path specified")
            start_list = abspath(start).split(sep)
            path_list = abspath(path).split(sep)
            if start_list[0].lower() != path_list[0].lower():
                unc_path, rest = splitunc(path)
                unc_start, rest = splitunc(start)
                if bool(unc_path) ^ bool(unc_start):
                    raise ValueError(
                        "Cannot mix UNC and non-UNC paths (%s and %s)" % (path, start)
                    )
                else:
                    raise ValueError(
                        "path is on drive %s, start on drive %s"
                        % (path_list[0], start_list[0])
                    )
            # Work out how much of the filepath is shared by start and path.
            for i in range(min(len(start_list), len(path_list))):
                if start_list[i].lower() != path_list[i].lower():
                    break
            else:
                i += 1

            rel_list = [pardir] * (len(start_list) - i) + path_list[i:]
            if not rel_list:
                return curdir
            return join(*rel_list)

    else:
        raise RuntimeError("Unsupported platform (no relpath available!)")