# pylint: disable=too-many-locals, too-many-arguments
# pylint: disable=too-many-branches
"""Plotting Library."""
import json
import warnings
from io import BytesIO
from typing import Any, Optional, Union

import numpy as np

from ._typing import PathLike
from .core import Booster, _deprecate_positional_args
from .sklearn import XGBModel

Axes = Any  # real type is matplotlib.axes.Axes
GraphvizSource = Any  # real type is graphviz.Source


@_deprecate_positional_args
def plot_importance(
    booster: Union[XGBModel, Booster, dict],
    *,
    ax: Optional[Axes] = None,
    height: float = 0.2,
    xlim: Optional[tuple] = None,
    ylim: Optional[tuple] = None,
    title: str = "Feature importance",
    xlabel: str = "Importance score",
    ylabel: str = "Features",
    fmap: PathLike = "",
    importance_type: str = "weight",
    max_num_features: Optional[int] = None,
    grid: bool = True,
    show_values: bool = True,
    values_format: str = "{v}",
    **kwargs: Any,
) -> Axes:
    """Plot importance based on fitted trees.

    Parameters
    ----------
    booster :
        Booster or XGBModel instance, or dict taken by Booster.get_fscore()
    ax : matplotlib Axes
        Target axes instance. If None, new figure and axes will be created.
    grid :
        Turn the axes grids on or off.  Default is True (On).
    importance_type :
        How the importance is calculated: either "weight", "gain", or "cover"

        * "weight" is the number of times a feature appears in a tree
        * "gain" is the average gain of splits which use the feature
        * "cover" is the average coverage of splits which use the feature
          where coverage is defined as the number of samples affected by the split
    max_num_features :
        Maximum number of top features displayed on plot. If None, all features will be
        displayed.
    height :
        Bar height, passed to ax.barh()
    xlim :
        Tuple passed to axes.xlim()
    ylim :
        Tuple passed to axes.ylim()
    title :
        Axes title. To disable, pass None.
    xlabel :
        X axis title label. To disable, pass None.
    ylabel :
        Y axis title label. To disable, pass None.
    fmap :
        The name of feature map file.
    show_values :
        Show values on plot. To disable, pass False.
    values_format :
        Format string for values. "v" will be replaced by the value of the feature
        importance.  e.g. Pass "{v:.2f}" in order to limit the number of digits after
        the decimal point to two, for each value printed on the graph.
    kwargs :
        Other keywords passed to ax.barh()

    Returns
    -------
    ax : matplotlib Axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError as e:
        raise ImportError("You must install matplotlib to plot importance") from e

    if isinstance(booster, XGBModel):
        importance = booster.get_booster().get_score(
            importance_type=importance_type, fmap=fmap
        )
    elif isinstance(booster, Booster):
        importance = booster.get_score(importance_type=importance_type, fmap=fmap)
    elif isinstance(booster, dict):
        importance = booster
    else:
        raise ValueError("tree must be Booster, XGBModel or dict instance")

    if not importance:
        raise ValueError(
            "Booster.get_score() results in empty.  "
            + "This maybe caused by having all trees as decision dumps."
        )

    tuples = [(k, importance[k]) for k in importance]
    if max_num_features is not None:
        # pylint: disable=invalid-unary-operand-type
        tuples = sorted(tuples, key=lambda _x: _x[1])[-max_num_features:]
    else:
        tuples = sorted(tuples, key=lambda _x: _x[1])
    labels, values = zip(*tuples)

    if ax is None:
        _, ax = plt.subplots(1, 1)

    ylocs = np.arange(len(values))
    ax.barh(ylocs, values, align="center", height=height, **kwargs)

    if show_values is True:
        for x, y in zip(values, ylocs):
            ax.text(x + 1, float(y), values_format.format(v=x), va="center")

    ax.set_yticks(ylocs)
    ax.set_yticklabels(labels)

    if xlim is not None:
        if not isinstance(xlim, tuple) or len(xlim) != 2:
            raise ValueError("xlim must be a tuple of 2 elements")
    else:
        xlim = (0, max(values) * 1.1)
    ax.set_xlim(xlim)

    if ylim is not None:
        if not isinstance(ylim, tuple) or len(ylim) != 2:
            raise ValueError("ylim must be a tuple of 2 elements")
    else:
        ylim = (-1, len(values))
    ax.set_ylim(ylim)

    if title is not None:
        ax.set_title(title)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)
    ax.grid(grid)
    return ax


@_deprecate_positional_args
def to_graphviz(
    booster: Union[Booster, XGBModel],
    *,
    fmap: PathLike = "",
    num_trees: Optional[int] = None,
    rankdir: Optional[str] = None,
    yes_color: Optional[str] = None,
    no_color: Optional[str] = None,
    condition_node_params: Optional[dict] = None,
    leaf_node_params: Optional[dict] = None,
    with_stats: bool = False,
    tree_idx: int = 0,
    **kwargs: Any,
) -> GraphvizSource:
    """Convert specified tree to graphviz instance. IPython can automatically plot
    the returned graphviz instance. Otherwise, you should call .render() method
    of the returned graphviz instance.

    Parameters
    ----------
    booster :
        Booster or XGBModel instance
    fmap :
       The name of feature map file
    num_trees :

        .. deprecated:: 3.0

        Specify the ordinal number of target tree

    rankdir :
        Passed to graphviz via graph_attr
    yes_color :
        Edge color when meets the node condition.
    no_color :
        Edge color when doesn't meet the node condition.
    condition_node_params :
        Condition node configuration for for graphviz.  Example:

        .. code-block:: python

            {'shape': 'box',
             'style': 'filled,rounded',
             'fillcolor': '#78bceb'}

    leaf_node_params :
        Leaf node configuration for graphviz. Example:

        .. code-block:: python

            {'shape': 'box',
             'style': 'filled',
             'fillcolor': '#e48038'}

    with_stats :

        .. versionadded:: 3.0

        Controls whether the split statistics should be included.

    tree_idx :

        .. versionadded:: 3.0

        Specify the ordinal index of target tree.

    kwargs :
        Other keywords passed to graphviz graph_attr, e.g. ``graph [ {key} = {value} ]``

    Returns
    -------
    graph: graphviz.Source

    """
    try:
        from graphviz import Source
    except ImportError as e:
        raise ImportError("You must install graphviz to plot tree") from e
    if isinstance(booster, XGBModel):
        booster = booster.get_booster()

    # squash everything back into kwargs again for compatibility
    parameters = "dot"
    extra = {}
    for key, value in kwargs.items():
        extra[key] = value

    if rankdir is not None:
        kwargs["graph_attrs"] = {}
        kwargs["graph_attrs"]["rankdir"] = rankdir
    for key, value in extra.items():
        if kwargs.get("graph_attrs", None) is not None:
            kwargs["graph_attrs"][key] = value
        else:
            kwargs["graph_attrs"] = {}
        del kwargs[key]

    if yes_color is not None or no_color is not None:
        kwargs["edge"] = {}
    if yes_color is not None:
        kwargs["edge"]["yes_color"] = yes_color
    if no_color is not None:
        kwargs["edge"]["no_color"] = no_color

    if condition_node_params is not None:
        kwargs["condition_node_params"] = condition_node_params
    if leaf_node_params is not None:
        kwargs["leaf_node_params"] = leaf_node_params

    if kwargs:
        parameters += ":"
        parameters += json.dumps(kwargs)

    if num_trees is not None:
        warnings.warn(
            "The `num_trees` parameter is deprecated, use `tree_idx` insetad. ",
            FutureWarning,
        )
        if tree_idx not in (0, num_trees):
            raise ValueError(
                "Both `num_trees` and `tree_idx` are used, prefer `tree_idx` instead."
            )
        tree_idx = num_trees

    tree = booster.get_dump(fmap=fmap, dump_format=parameters, with_stats=with_stats)[
        tree_idx
    ]
    g = Source(tree)
    return g


@_deprecate_positional_args
def plot_tree(
    booster: Union[Booster, XGBModel],
    *,
    fmap: PathLike = "",
    num_trees: Optional[int] = None,
    rankdir: Optional[str] = None,
    ax: Optional[Axes] = None,
    with_stats: bool = False,
    tree_idx: int = 0,
    **kwargs: Any,
) -> Axes:
    """Plot specified tree.

    Parameters
    ----------
    booster :
        Booster or XGBModel instance
    fmap: str (optional)
       The name of feature map file
    num_trees :

        .. deprecated:: 3.0

    rankdir : str, default "TB"
        Passed to graphviz via graph_attr
    ax : matplotlib Axes, default None
        Target axes instance. If None, new figure and axes will be created.

    with_stats :

        .. versionadded:: 3.0

        See :py:func:`to_graphviz`.

    tree_idx :

        .. versionadded:: 3.0

        See :py:func:`to_graphviz`.

    kwargs :
        Other keywords passed to :py:func:`to_graphviz`

    Returns
    -------
    ax : matplotlib Axes

    """
    try:
        from matplotlib import image
        from matplotlib import pyplot as plt
    except ImportError as e:
        raise ImportError("You must install matplotlib to plot tree") from e

    if ax is None:
        _, ax = plt.subplots(1, 1)

    g = to_graphviz(
        booster,
        fmap=fmap,
        num_trees=num_trees,
        rankdir=rankdir,
        with_stats=with_stats,
        tree_idx=tree_idx,
        **kwargs,
    )

    s = BytesIO()
    s.write(g.pipe(format="png"))
    s.seek(0)
    img = image.imread(s)

    ax.imshow(img)
    ax.axis("off")
    return ax
