"""
Script to autogenerate pyplot wrappers.

pyplot.py consists of two parts: a hand-written part at the top, and an
automatically generated part at the bottom, starting with the comment

  ### REMAINING CONTENT GENERATED BY boilerplate.py ###

This script generates the automatically generated part of pyplot.py. It
consists of colormap setter functions and wrapper functions for methods
of Figure and Axes. Whenever the API of one of the wrapped methods changes,
this script has to be rerun to keep pyplot.py up to date.

The test ``lib/matplotlib/test_pyplot.py::test_pyplot_up_to_date`` checks
that the autogenerated part of pyplot.py is up to date. It will fail in the
case of an API mismatch and remind the developer to rerun this script.
"""

# Although it is possible to dynamically generate the pyplot functions at
# runtime with the proper signatures, a static pyplot.py is simpler for static
# analysis tools to parse.

import ast
from enum import Enum
import functools
import inspect
from inspect import Parameter
from pathlib import Path
import sys
import subprocess


# This line imports the installed copy of matplotlib, and not the local copy.
import numpy as np
from matplotlib import _api, mlab
from matplotlib.axes import Axes
from matplotlib.figure import Figure


# This is the magic line that must exist in pyplot, after which the boilerplate
# content will be appended.
PYPLOT_MAGIC_HEADER = (
    "################# REMAINING CONTENT GENERATED BY boilerplate.py "
    "##############\n")

AUTOGEN_MSG = """

# Autogenerated by boilerplate.py.  Do not edit as changes will be lost."""

AXES_CMAPPABLE_METHOD_TEMPLATE = AUTOGEN_MSG + """
@_copy_docstring_and_deprecators(Axes.{called_name})
def {name}{signature}:
    __ret = gca().{called_name}{call}
    {sci_command}
    return __ret
"""

AXES_METHOD_TEMPLATE = AUTOGEN_MSG + """
@_copy_docstring_and_deprecators(Axes.{called_name})
def {name}{signature}:
    {return_statement}gca().{called_name}{call}
"""

FIGURE_METHOD_TEMPLATE = AUTOGEN_MSG + """
@_copy_docstring_and_deprecators(Figure.{called_name})
def {name}{signature}:
    {return_statement}gcf().{called_name}{call}
"""

CMAP_TEMPLATE = '''
def {name}() -> None:
    """
    Set the colormap to {name!r}.

    This changes the default colormap as well as the colormap of the current
    image if there is one. See ``help(colormaps)`` for more information.
    """
    set_cmap({name!r})
'''  # Colormap functions.


class value_formatter:
    """
    Format function default values as needed for inspect.formatargspec.
    The interesting part is a hard-coded list of functions used
    as defaults in pyplot methods.
    """

    def __init__(self, value):
        if value is mlab.detrend_none:
            self._repr = "mlab.detrend_none"
        elif value is mlab.window_hanning:
            self._repr = "mlab.window_hanning"
        elif value is np.mean:
            self._repr = "np.mean"
        elif value is _api.deprecation._deprecated_parameter:
            self._repr = "_api.deprecation._deprecated_parameter"
        elif isinstance(value, Enum):
            # Enum str is Class.Name whereas their repr is <Class.Name: value>.
            self._repr = f'{type(value).__name__}.{value.name}'
        else:
            self._repr = repr(value)

    def __repr__(self):
        return self._repr


class direct_repr:
    """
    A placeholder class to destringify annotations from ast
    """
    def __init__(self, value):
        self._repr = value

    def __repr__(self):
        return self._repr


def generate_function(name, called_fullname, template, **kwargs):
    """
    Create a wrapper function *pyplot_name* calling *call_name*.

    Parameters
    ----------
    name : str
        The function to be created.
    called_fullname : str
        The method to be wrapped in the format ``"Class.method"``.
    template : str
        The template to be used. The template must contain {}-style format
        placeholders. The following placeholders are filled in:

        - name: The function name.
        - signature: The function signature (including parentheses).
        - called_name: The name of the called function.
        - call: Parameters passed to *called_name* (including parentheses).

    **kwargs
        Additional parameters are passed to ``template.format()``.
    """
    # Get signature of wrapped function.
    class_name, called_name = called_fullname.split('.')
    class_ = {'Axes': Axes, 'Figure': Figure}[class_name]

    meth = getattr(class_, called_name)
    decorator = _api.deprecation.DECORATORS.get(meth)
    # Generate the wrapper with the non-kwonly signature, as it will get
    # redecorated with make_keyword_only by _copy_docstring_and_deprecators.
    if decorator and decorator.func is _api.make_keyword_only:
        meth = meth.__wrapped__

    annotated_trees = get_ast_mro_trees(class_)
    signature = get_matching_signature(meth, annotated_trees)

    # Replace self argument.
    params = list(signature.parameters.values())[1:]
    has_return_value = str(signature.return_annotation) != 'None'
    signature = str(signature.replace(parameters=[
        param.replace(default=value_formatter(param.default))
        if param.default is not param.empty else param
        for param in params]))
    # How to call the wrapped function.
    call = '(' + ', '.join((
           # Pass "intended-as-positional" parameters positionally to avoid
           # forcing third-party subclasses to reproduce the parameter names.
           '{0}'
           if param.kind in [
               Parameter.POSITIONAL_OR_KEYWORD]
              and param.default is Parameter.empty else
           # Only pass the data kwarg if it is actually set, to avoid forcing
           # third-party subclasses to support it.
           '**({{"data": data}} if data is not None else {{}})'
           if param.name == "data" else
           '{0}={0}'
           if param.kind in [
               Parameter.POSITIONAL_OR_KEYWORD,
               Parameter.KEYWORD_ONLY] else
           '{0}'
           if param.kind is Parameter.POSITIONAL_ONLY else
           '*{0}'
           if param.kind is Parameter.VAR_POSITIONAL else
           '**{0}'
           if param.kind is Parameter.VAR_KEYWORD else
           None).format(param.name)
       for param in params) + ')'
    return_statement = 'return ' if has_return_value else ''
    # Bail out in case of name collision.
    for reserved in ('gca', 'gci', 'gcf', '__ret'):
        if reserved in params:
            raise ValueError(
                f'Method {called_fullname} has kwarg named {reserved}')

    return template.format(
        name=name,
        called_name=called_name,
        signature=signature,
        call=call,
        return_statement=return_statement,
        **kwargs)


def boilerplate_gen():
    """Generator of lines for the automated part of pyplot."""

    _figure_commands = (
        'figimage',
        'figtext:text',
        'gca',
        'gci:_gci',
        'ginput',
        'subplots_adjust',
        'suptitle',
        'tight_layout',
        'waitforbuttonpress',
    )

    # These methods are all simple wrappers of Axes methods by the same name.
    _axes_commands = (
        'acorr',
        'angle_spectrum',
        'annotate',
        'arrow',
        'autoscale',
        'axhline',
        'axhspan',
        'axis',
        'axline',
        'axvline',
        'axvspan',
        'bar',
        'barbs',
        'barh',
        'bar_label',
        'boxplot',
        'broken_barh',
        'clabel',
        'cohere',
        'contour',
        'contourf',
        'csd',
        'ecdf',
        'errorbar',
        'eventplot',
        'fill',
        'fill_between',
        'fill_betweenx',
        'grid',
        'grouped_bar',
        'hexbin',
        'hist',
        'stairs',
        'hist2d',
        'hlines',
        'imshow',
        'legend',
        'locator_params',
        'loglog',
        'magnitude_spectrum',
        'margins',
        'minorticks_off',
        'minorticks_on',
        'pcolor',
        'pcolormesh',
        'phase_spectrum',
        'pie',
        'pie_label',
        'plot',
        'psd',
        'quiver',
        'quiverkey',
        'scatter',
        'semilogx',
        'semilogy',
        'specgram',
        'spy',
        'stackplot',
        'stem',
        'step',
        'streamplot',
        'table',
        'text',
        'tick_params',
        'ticklabel_format',
        'tricontour',
        'tricontourf',
        'tripcolor',
        'triplot',
        'violinplot',
        'vlines',
        'xcorr',
        # pyplot name : real name
        'sci:_sci',
        'title:set_title',
        'xlabel:set_xlabel',
        'ylabel:set_ylabel',
        'xscale:set_xscale',
        'yscale:set_yscale',
    )

    cmappable = {
        'contour': (
            'if __ret._A is not None:  # type: ignore[attr-defined]\n'
            '        sci(__ret)'
        ),
        'contourf': (
            'if __ret._A is not None:  # type: ignore[attr-defined]\n'
            '        sci(__ret)'
        ),
        'hexbin': 'sci(__ret)',
        'scatter': 'sci(__ret)',
        'pcolor': 'sci(__ret)',
        'pcolormesh': 'sci(__ret)',
        'hist2d': 'sci(__ret[-1])',
        'imshow': 'sci(__ret)',
        'spy': (
            'if isinstance(__ret, _ColorizerInterface):\n'
            '        sci(__ret)'
        ),
        'quiver': 'sci(__ret)',
        'specgram': 'sci(__ret[-1])',
        'streamplot': 'sci(__ret.lines)',
        'tricontour': (
            'if __ret._A is not None:  # type: ignore[attr-defined]\n'
            '        sci(__ret)'
        ),
        'tricontourf': (
            'if __ret._A is not None:  # type: ignore[attr-defined]\n'
            '        sci(__ret)'
        ),
        'tripcolor': 'sci(__ret)',
    }

    for spec in _figure_commands:
        if ':' in spec:
            name, called_name = spec.split(':')
        else:
            name = called_name = spec
        yield generate_function(name, f'Figure.{called_name}',
                                FIGURE_METHOD_TEMPLATE)

    for spec in _axes_commands:
        if ':' in spec:
            name, called_name = spec.split(':')
        else:
            name = called_name = spec

        template = (AXES_CMAPPABLE_METHOD_TEMPLATE if name in cmappable else
                    AXES_METHOD_TEMPLATE)
        yield generate_function(name, f'Axes.{called_name}', template,
                                sci_command=cmappable.get(name))

    cmaps = (
        'autumn',
        'bone',
        'cool',
        'copper',
        'flag',
        'gray',
        'hot',
        'hsv',
        'jet',
        'pink',
        'prism',
        'spring',
        'summer',
        'winter',
        'magma',
        'inferno',
        'plasma',
        'viridis',
        "nipy_spectral"
    )
    # add all the colormaps (autumn, hsv, ....)
    for name in cmaps:
        yield AUTOGEN_MSG
        yield CMAP_TEMPLATE.format(name=name)


def build_pyplot(pyplot_path):
    pyplot_orig = pyplot_path.read_text().splitlines(keepends=True)
    try:
        pyplot_orig = pyplot_orig[:pyplot_orig.index(PYPLOT_MAGIC_HEADER) + 1]
    except IndexError as err:
        raise ValueError('The pyplot.py file *must* have the exact line: %s'
                         % PYPLOT_MAGIC_HEADER) from err

    with pyplot_path.open('w') as pyplot:
        pyplot.writelines(pyplot_orig)
        pyplot.writelines(boilerplate_gen())

    # Run black to autoformat pyplot
    subprocess.run(
        [sys.executable, "-m", "black", "--line-length=88", pyplot_path],
        check=True
    )


### Methods for retrieving signatures from pyi stub files

def get_ast_tree(cls):
    path = Path(inspect.getfile(cls))
    stubpath = path.with_suffix(".pyi")
    path = stubpath if stubpath.exists() else path
    tree = ast.parse(path.read_text())
    for item in tree.body:
        if isinstance(item, ast.ClassDef) and item.name == cls.__name__:
            return item
    raise ValueError(f"Cannot find {cls.__name__} in ast")


@functools.lru_cache
def get_ast_mro_trees(cls):
    return [get_ast_tree(c) for c in cls.__mro__ if c.__module__ != "builtins"]


def get_matching_signature(method, trees):
    sig = inspect.signature(method)
    for tree in trees:
        for item in tree.body:
            if not isinstance(item, ast.FunctionDef):
                continue
            if item.name == method.__name__:
                return update_sig_from_node(item, sig)
    # The following methods are implemented outside of the mro of Axes
    # and thus do not get their annotated versions found with current code
    #     stackplot
    #     streamplot
    #     table
    #     tricontour
    #     tricontourf
    #     tripcolor
    #     triplot

    # import warnings
    # warnings.warn(f"'{method.__name__}' not found")
    return sig


def update_sig_from_node(node, sig):
    params = dict(sig.parameters)
    args = node.args
    allargs = (
        *args.posonlyargs,
        *args.args,
        args.vararg,
        *args.kwonlyargs,
        args.kwarg,
    )
    for param in allargs:
        if param is None:
            continue
        if param.annotation is None:
            continue
        annotation = direct_repr(ast.unparse(param.annotation))
        params[param.arg] = params[param.arg].replace(annotation=annotation)

    if node.returns is not None:
        return inspect.Signature(
            params.values(),
            return_annotation=direct_repr(ast.unparse(node.returns))
        )
    else:
        return inspect.Signature(params.values())


if __name__ == '__main__':
    # Write the matplotlib.pyplot file.
    if len(sys.argv) > 1:
        pyplot_path = Path(sys.argv[1])
    else:
        mpl_path = (Path(__file__).parent / ".." /"lib"/"matplotlib").resolve()
        pyplot_path = mpl_path / "pyplot.py"
        for cls in [Axes, Figure]:
            if mpl_path not in  Path(inspect.getfile(cls)).parents:
                raise RuntimeError(
                    f"{cls.__name__} import path is not {mpl_path}.\n"
                    "Please make sure your Matplotlib installation "
                    "is from the same source checkout as boilerplate.py"
                )
    build_pyplot(pyplot_path)