from .base import Plugin
from ..utils import ClearColormap


class OverlayPlugin(Plugin):
    """Plugin for ImageViewer that displays an overlay on top of main image.

    Attributes
    ----------
    overlay : array
        Overlay displayed on top of image. This overlay defaults to a color map
        with alpha values varying linearly from 0 to 1.
    """
    colors = {'red': (1, 0, 0),
              'yellow': (1, 1, 0),
              'green': (0, 1, 0),
              'cyan': (0, 1, 1)}

    def __init__(self, **kwargs):
        super(OverlayPlugin, self).__init__(**kwargs)
        self._overlay_plot = None
        self._overlay = None
        self.cmap = None
        self.color_names = self.colors.keys()

    def attach(self, image_viewer):
        super(OverlayPlugin, self).attach(image_viewer)
        #TODO: `color` doesn't update GUI widget when set manually.
        self.color = 0

    @property
    def overlay(self):
        return self._overlay

    @overlay.setter
    def overlay(self, image):
        self._overlay = image
        ax = self.image_viewer.ax
        if image is None:
            ax.images.remove(self._overlay_plot)
            self._overlay_plot = None
        elif self._overlay_plot is None:
            self._overlay_plot = ax.imshow(image, cmap=self.cmap)
        else:
            self._overlay_plot.set_array(image)
        self.image_viewer.redraw()

    def display_filtered_image(self, image):
        """Display image over image in viewer."""
        self.overlay = image

    def closeEvent(self, event):
        self.overlay = None
        super(OverlayPlugin, self).closeEvent(event)

    @property
    def color(self):
        return self._color

    @color.setter
    def color(self, index):
        # Update colormap whenever color is changed.
        name = self.color_names[index]
        self._color = name
        rgb = self.colors[name]
        self.cmap = ClearColormap(rgb)

        if self._overlay_plot is not None:
            self._overlay_plot.set_cmap(self.cmap)
        self.image_viewer.redraw()

@@ -161,16 +161,223 @@
 image.%0A%0A
+    The base Plugin class displays the filtered image directly on the viewer.%0A    OverlayPlugin will instead overlay an image with a transparent colormap.%0A%0A    See base Plugin class for additional details.%0A%0A
     Attr
@@ -554,16 +554,58 @@
 0 to 1.%0A
+    color : int%0A        Color of overlay.%0A
     %22%22%22%0A
@@ -1713,18 +1713,46 @@
 lay 
-image over
+filtered image as an overlay on top of
  ima
@@ -1831,16 +1831,66 @@
 event):%0A
+        # clear overlay from ImageViewer on close%0A
         
@@ -2125,16 +2125,189 @@
 hanged.%0A
+        if isinstance(index, basestring) and index not in self.color_names:%0A            raise ValueError(%22%25s not defined in OverlayPlugin.colors%22 %25 index)%0A        else:%0A    
         

allow color to be set by name