"""
Implementation of a plut using a custom overlay and tool
"""

from __future__ import with_statement

import numpy

from enthought.traits.api import HasTraits, Instance, Enum
from enthought.traits.ui.api import View, Item
from enthought.enable.api import ComponentEditor
from enthought.chaco.api import Plot, ArrayPlotData, AbstractOverlay
from enthought.enable.api import BaseTool
from enthought.enable.markers import DOT_MARKER, DotMarker

class BoxSelectTool(BaseTool):
    """ Tool for selecting all points within a box
    
        There are 2 states for this tool, normal and selecting. While the
        left mouse button is down the metadata on the datasources will be
        updated with the current selected bounds.
        
        Note that the tool does not actually store the selected point, but the
        bounds of the box. 
    """

    event_state = Enum("normal", "selecting")

    def normal_left_down(self, event):
        self.event_state = "selecting"
        self.selecting_mouse_move(event)
        
    def selecting_left_up(self, event):
        self.event_state = "normal"
        
    def selecting_mouse_move(self, event):
        x1, y1 = self.map_to_data(event.x-25, event.y-25)
        x2, y2 = self.map_to_data(event.x+25, event.y+25)
        
        index_datasource = self.component.index
        index_datasource.metadata['selections'] = (x1, x2)
        
        value_datasource = self.component.value
        value_datasource.metadata['selections'] = (y1, y2)
        
        self.component.request_redraw()
        
    def map_to_data(self, x, y):
        """ Returns the data space coordinates of the given x and y. 

        Takes into account orientation of the plot and the axis setting.
        """

        plot = self.component
        if plot.orientation == "h":
            index = plot.x_mapper.map_data(x)
            value = plot.y_mapper.map_data(y)
        else:
            index = plot.y_mapper.map_data(y)
            value = plot.x_mapper.map_data(x)
            
        return index, value
        

class XRayOverlay(AbstractOverlay):
    """ Overlay which draws scatter markers on top of plot data points.
    
        This overlay should be combined with a tool which updates the 
        datasources metadata with selection bounds.
    """

    marker = DotMarker()
    
    def overlay(self, component, gc, view_bounds=None, mode='normal'):
        x_range = self._get_selection_index_screen_range()
        y_range = self._get_selection_value_screen_range()
        
        if len(x_range) == 0:
            return
        
        x1, x2 = x_range
        y1, y2 = y_range

        with gc:
            gc.set_alpha(0.8)
            gc.set_fill_color((1.0,1.0,1.0))
            gc.rect(x1, y1, x2-x1, y2-y1)
            gc.draw_path()

        pts = self._get_selected_points()
        if len(pts) == 0:
            return
        screen_pts = self.component.map_screen(pts)
        if hasattr(gc, 'draw_marker_at_points'):
            gc.draw_marker_at_points(screen_pts, 3, DOT_MARKER)
        else:
            gc.save_state()
            for sx,sy in screen_pts:
                gc.translate_ctm(sx, sy)
                gc.begin_path()
                self.marker.add_to_path(gc, 3)
                gc.draw_path(self.marker.draw_mode)
                gc.translate_ctm(-sx, -sy)
            gc.restore_state()
        
    def _get_selected_points(self):
        """ gets all the points within the bounds defined in the datasources
            metadata
        """
        index_datasource = self.component.index
        index_selection = index_datasource.metadata['selections']
        index = index_datasource.get_data()
        
        value_datasource = self.component.value
        value_selection = value_datasource.metadata['selections']
        value = value_datasource.get_data()
        
        x_indices = numpy.where((index > index_selection[0]) & (index < index_selection[-1]))
        y_indices = numpy.where((value > value_selection[0]) & (value < value_selection[-1]))
        
        indices = list(set(x_indices[0]) & set(y_indices[0]))
        
        sel_index = index[indices]
        sel_value = value[indices]
        
        return zip(sel_index, sel_value)
        
    def _get_selection_index_screen_range(self):
        """ maps the selected bounds which were set by the tool into screen
            space. The screen space points can be used for drawing the overlay
        """
        index_datasource = self.component.index
        index_mapper = self.component.index_mapper
        index_selection = index_datasource.metadata['selections']
        return tuple(index_mapper.map_screen(numpy.array(index_selection)))

    def _get_selection_value_screen_range(self):
        """ maps the selected bounds which were set by the tool into screen
            space. The screen space points can be used for drawing the overlay
        """
        value_datasource = self.component.value
        value_mapper = self.component.value_mapper
        value_selection = value_datasource.metadata['selections']
        return tuple(value_mapper.map_screen(numpy.array(value_selection)))
        
class PlotExample(HasTraits):

    plot = Instance(Plot)
    
    traits_view = View(Item('plot', editor=ComponentEditor()),
                       width=600, height=600)
                       
    def __init__(self, index, value, *args, **kw):
        super(PlotExample, self).__init__(*args, **kw)
        
        plot_data = ArrayPlotData(index=index)
        plot_data.set_data('value', value)
        
        self.plot = Plot(plot_data)
        line = self.plot.plot(('index', 'value'))[0]
        
        line.overlays.append(XRayOverlay(line))
        line.tools.append(BoxSelectTool(line))
        
index = numpy.arange(0, 25, 0.25)
value = numpy.sin(index) + numpy.arange(0, 10, 0.1)

example = PlotExample(index, value)
example.configure_traits()