from matplotlib.patches import RegularPolygon, Rectangle
from matplotlib.widgets import RectangleSelector


class SquareSelector(RectangleSelector):
    # smaller minspan makes little sense due to interpolation radius
    def __init__(self, ax, onselect=None, useblit=True, rectprops=None,
                 dflt_size=50, minspan=15, maxselections=4):
        if rectprops is None:
            rectprops = dict(facecolor='red', edgecolor='black',
                             alpha=0.5, fill=True)
        self.sq_onselect = onselect
        RectangleSelector.__init__(self, ax=ax, onselect=lambda x, y: True,
                                   useblit=useblit, rectprops=rectprops,
                                   minspanx=None, minspany=None,
                                   spancoords='pixels', drawtype='box')
        self.selections = []
        self.dflt_size = dflt_size
        self.minspan = minspan
        self.maxselections = maxselections
        self.cross = None

    def update(self):
        if self.useblit:
            if self.background is not None:
                self.canvas.restore_region(self.background)
            self.ax.draw_artist(self.to_draw)
            if self.cross is not None:
                self.ax.draw_artist(self.cross)
            for a in self.selections:
                self.ax.draw_artist(a[0])
                self.ax.draw_artist(a[1])
            self.canvas.blit(self.ax.bbox)
        else:
            self.canvas.draw_idle()
        return False

    def press(self, event):
        # Is the correct button pressed within the correct axes?
        if self.ignore(event):
            return
        if event.button == 3:
            rem = []
            for s in reversed(self.selections):
                if s[1].contains(event)[0]:
                    x, y = s[1].xy
                    h = s[1].get_height() / 2
                    if self.sq_onselect is not None:
                        self.sq_onselect([x + h, y + h, h], False)
                    self.ax.patches.remove(s[0])
                    self.ax.patches.remove(s[1])
                    rem.append(s)
                    break
            for r in rem:
                self.selections.remove(r)
            self.canvas.draw_idle()
            return
        if len(self.selections) >= self.maxselections:
            return

        # make the drawed box/line visible get the click-coordinates,
        # button, ...
        self.to_draw.set_visible(self.visible)
        self.eventpress = event

        self.cross = RegularPolygon((event.xdata, event.ydata), numVertices=4,
                                    radius=1, facecolor='white')
        s = self.dflt_size
        self.to_draw.set_width(s * 2)
        self.to_draw.set_height(s * 2)
        self.to_draw.set_x(event.xdata - s)
        self.to_draw.set_y(event.ydata - s)
        self.ax.add_patch(self.cross)
        self.eventpress.data = (event.xdata, event.ydata, s)
        self.update()
        return False

    def release(self, event):
        if not (self.eventpress is None or self.ignore(event)) and \
           event.button == 1 and len(self.selections) < self.maxselections:
            c, r = draw_selection(*self.eventpress.data)
            self.ax.add_patch(r)
            self.ax.add_patch(c)
            self.selections.append((c, r))
            self.sq_onselect(self.eventpress.data, True)

        if self.cross is not None:
            self.ax.patches.remove(self.cross)
            self.cross = None

        if self.ignore(event):
            self.to_draw.set_visible(False)
            self.canvas.draw_idle()

        RectangleSelector.release(self, event)

    def onmove(self, event):
        if self.eventpress is None or self.ignore(event):
            return
        x, y = event.xdata, event.ydata

        bx, by = self.eventpress.xdata, self.eventpress.ydata
        s = max(abs(by - y), abs(bx - x))
        if s < self.minspan:
            s = self.dflt_size

        self.to_draw.set_width(s * 2)
        self.to_draw.set_height(s * 2)

        self.to_draw.set_x(bx - s)
        self.to_draw.set_y(by - s)

        self.eventpress.data = (bx, by, s)
        self.update()
        return False

    def get_selections(self):
        res = []
        for s in self.selections:
            x, y = s[1].xy
            w = s[1].get_height()
            res.append((x + w / 2, y + w / 2, w))
        return res


def draw_selection(x, y, s, rectprops=None):
    if rectprops is None:
        rectprops = dict(facecolor='white', edgecolor='black',
                         alpha=0.5, fill=True)

    cross = RegularPolygon((x, y), numVertices=4, radius=1, facecolor='white')
    rect = Rectangle((x - s, y - s), s * 2, s * 2, visible=True, **rectprops)
    return cross, rect
