from mpldatacursor import DataCursor
import matplotlib.colors as mplcolors
from matplotlib import _pylab_helpers as pylab_helpers
import numpy as np


def rectangle_props(event):
    """
    Returns the width, height, left, and bottom of a rectangle artist.

    Parameters
    -----------
    event : PickEvent
        The pick event to process

    Returns
    --------
    A dict with keys:
        `width` : The width of the rectangle
        `height` : The height of the rectangle
        `left` : The minimum x-coordinate of the rectangle
        `right` : The maximum x-coordinate of the rectangle
        `bottom` : The minimum y-coordinate of the rectangle
        `top` : The maximum y-coordinate of the rectangle
        `xcenter` : The mean x-coordinate of the rectangle
        `ycenter` : The mean y-coordinate of the rectangle
        `label` : The label for the rectangle or None
    """
    artist = event.artist
    color = mplcolors.to_hex(artist.get_facecolor())
    width, height = artist.get_width(), artist.get_height()
    left, bottom = artist.xy
    right, top = left + width, bottom + height
    xcenter = left + 0.5 * width
    ycenter = bottom + 0.5 * height

    label = artist.get_label()
    if label is None or label.startswith('_nolegend'):
        try:
            label = artist._mpldatacursor_label
        except AttributeError:
            label = None

    return dict(color=color, width=width, height=height, left=left, bottom=bottom,
                label=label, right=right, top=top,
                xcenter=xcenter, ycenter=ycenter)


class ColorDataCursor(DataCursor):
    def __init__(self, *args, **kwargs):
        self.color_label_map = kwargs.pop('color_label_map')

        super().__init__(*args, **kwargs)

    def event_info(self, event):
        """Get a dict of info for the artist selected by "event"."""
        x, y = event.mouseevent.xdata, event.mouseevent.ydata
        props = dict(x=x, y=y, label=event.artist.get_label(), event=event)
        props['ind'] = getattr(event, 'ind', None)
        props['point_label'] = self._point_label(event)

        funcs = [rectangle_props]

        for func in funcs:
            props.update(func(event))

        return props

    def update(self, event, annotation):
        """Update the specified annotation."""
        # Get artist-specific information about the pick event
        info = self.event_info(event)
        color = info['color'].upper()
        info['color_label'] = self.color_label_map[color]
        if self.props_override is not None:
            info = self.props_override(**info)

        # Update the xy position and text using the formatter function

        annotation.set_text(self.formatter(**info))
        annotation.xy = info['x'], info['y']

        # In case it's been hidden earlier...
        annotation.set_visible(True)

        if self.keep_inside:
            self._keep_annotation_inside(annotation)

        annotation._has_been_shown = True
        self._last_event = event
        self._last_annotation = annotation

        event.canvas.draw()


def colordatacursor(artists=None, axes=None, **kwargs):
    """
    Create an interactive data cursor for the specified artists or specified
    axes. The data cursor displays information about a selected artist in a
    "popup" annotation box.
    If a specific sequence of artists is given, only the specified artists will
    be interactively selectable.  Otherwise, all manually-plotted artists in
    *axes* will be used (*axes* defaults to all axes in all figures).
    Parameters
    -----------
    artists : a matplotlib artist or sequence of artists, optional
        The artists to make selectable and display information for. If this is
        not specified, then all manually plotted artists in `axes` will be
        used.
    axes : a matplotlib axes of sequence of axes, optional
        The axes to selected artists from if a sequence of artists is not
        specified.  If `axes` is not specified, then all available axes in all
        figures will be used.
    tolerance : number, optional
        The radius (in points) that the mouse click must be within to select
        the artist. Default: 5 points.
    formatter : callable, optional
        A function that accepts arbitrary kwargs and returns a string that will
        be displayed with annotate. Often, it is convienent to pass in the
        format method of a template string, e.g.
        ``formatter="{label}".format``.
        Keyword arguments passed in to the `formatter` function:
            `x`, `y` : floats
                The x and y data coordinates of the clicked point
            `event` : a matplotlib ``PickEvent``
                The pick event that was fired (note that the selected
                artist can be accessed through ``event.artist``).
            `label` : string or None
                The legend label of the selected artist.
            `ind` : list of ints or None
                If the artist has "subitems" (e.g. points in a scatter or
                line plot), this will be a list of the item(s) that were
                clicked on.  If the artist does not have "subitems", this
                will be None. Note that this is always a list, even when
                a single item is selected.
        Some selected artists may supply additional keyword arguments that
        are not always present, for example:
            `z` : number
                The "z" (usually color or array) value, if present. For an
                ``AxesImage`` (as created by ``imshow``), this will be the
                uninterpolated array value at the point clicked. For a
                ``PathCollection`` (as created by ``scatter``) this will be the
                "c" value if an array was passed to "c".
            `i`, `j` : ints
                The row, column indicies of the selected point for an
                ``AxesImage`` (as created by ``imshow``)
            `s` : number
                The size of the selected item in a ``PathCollection`` if a size
                array is specified.
            `c` : number
                The array value displayed as color for a ``PathCollection``
                if a "c" array is specified (identical to "z").
            `point_label` : list
                If `point_labels` is given when the data cursor is initialized
                and the artist has "subitems", this will be a list of the items
                of `point_labels` that correspond to the selected artists.
                Note that this is always a list, even when a single artist is
                selected.
            `width`, `height`, `top`, `bottom` : numbers
                The parameters for ``Rectangle`` artists (e.g. bar plots).
    point_labels : sequence or dict, optional
        For artists with "subitems" (e.g. Line2D's), the item(s) of
        `point_labels` corresponding to the selected "subitems" of the artist
        will be passed into the formatter function as the "point_label" kwarg.
        If a single sequence is given, it will be used for all artists with
        "subitems". Alternatively, a dict of artist:sequence pairs may be given
        to match an artist to the correct series of point labels.
    display : {"one-per-axes", "single", "multiple"}, optional
        Controls whether more than one annotation box will be shown.
        Default: "one-per-axes"
    draggable : boolean, optional
        Controls whether or not the annotation box will be interactively
        draggable to a new location after being displayed. Defaults to False.
    hover : boolean, optional
        If True, the datacursor will "pop up" when the mouse hovers over an
        artist.  Defaults to False.  Enabling hover also sets
        `display="single"` and `draggable=False`.
    props_override : function, optional
        If specified, this function customizes the parameters passed into the
        formatter function and the x, y location that the datacursor "pop up"
        "points" to.  This is often useful to make the annotation "point" to a
        specific side or corner of an artist, regardless of the position
        clicked. The function is passed the same kwargs as the `formatter`
        function and is expected to return a dict with at least the keys "x"
        and "y" (and probably several others).
        Expected call signature: `props_dict = props_override(**kwargs)`
    keybindings : boolean or dict, optional
        By default, the keys "d" and "t" will be bound to deleting/hiding all
        annotation boxes and toggling interactivity for datacursors,
        respectively.  If keybindings is False, the ability to hide/toggle
        datacursors interactively will be disabled. Alternatively, a dict of
        the form {'hide':'somekey', 'toggle':'somekey'} may specified to
        customize the keyboard shortcuts.
    date_format : string, optional
        The strftime-style formatting string for dates. Used only if the x or y
        axes have been set to display dates. Defaults to "%x %X".
    display_button: int, optional
        The mouse button that will triggers displaying an annotation box.
        Defaults to 1, for left-clicking. (Common options are 1:left-click,
        2:middle-click, 3:right-click)
    hide_button: int or None, optional
        The mouse button that triggers hiding the selected annotation box.
        Defaults to 3, for right-clicking. (Common options are 1:left-click,
        2:middle-click, 3:right-click, None:hiding disabled)
    keep_inside : boolean, optional
        Whether or not to adjust the x,y offset to keep the text box inside the
        figure. This option has no effect on draggable datacursors. Defaults to
        True. Note: Currently disabled on OSX and NbAgg/notebook backends.
    **kwargs : additional keyword arguments, optional
        Additional keyword arguments are passed on to annotate.
    Returns
    -------
    dc : A ``mpldatacursor.DataCursor`` instance
    """
    def plotted_artists(ax):
        artists = (ax.lines + ax.patches + ax.collections
                   + ax.images + ax.containers)
        return artists

    # If no axes are specified, get all axes.
    if axes is None:
        managers = pylab_helpers.Gcf.get_all_fig_managers()
        figs = [manager.canvas.figure for manager in managers]
        axes = [ax for fig in figs for ax in fig.axes]

    if not np.iterable(axes):
        axes = [axes]

    # If no artists are specified, get all manually plotted artists in all of
    # the specified axes.
    if artists is None:
        artists = [artist for ax in axes for artist in plotted_artists(ax)]

    return ColorDataCursor(artists, **kwargs)