# (C) Copyright 2005-2021 Enthought, Inc., Austin, TX
# All rights reserved.
#
# This software is provided without warranty under the terms of the BSD
# license included in LICENSE.txt and may be redistributed only under
# the conditions described in the aforementioned license. The license
# is also available online at http://www.enthought.com/licenses/BSD.txt
#
# Thanks for using Enthought open source!

""" Defines the ColormappedSelectionOverlay class.
"""
import functools

from numpy import logical_and

# Enthought library imports
from traits.api import Any, Bool, Float, Instance, Property, Enum
from traits.observation.events import TraitChangeEvent

# Local imports
from chaco.abstract_overlay import AbstractOverlay
from chaco.plots.colormapped_scatterplot import ColormappedScatterPlot


class ColormappedSelectionOverlay(AbstractOverlay):
    """
    Overlays and changes a ColormappedScatterPlot to fade its non-selected
    points to a very low alpha.
    """

    #: The ColormappedScatterPlot that this overlay is listening to.
    #: By default, it looks at self.component
    plot = Property

    #: The amount to fade the unselected points.
    fade_alpha = Float(0.15)

    #: The minimum difference, in float percent, between the starting and
    #: ending selection values, if range selection mode is enabled
    minimum_delta = Float(0.01)

    #: Outline width for selected points.
    selected_outline_width = Float(1.0)
    #: Outline width for unselected points.
    unselected_outline_width = Float(0.0)

    #: The type of selection used by the data source.
    selection_type = Enum("range", "mask")

    _plot = Instance(ColormappedScatterPlot)

    _visible = Bool(False)

    _old_alpha = Float
    _old_outline_color = Any
    _old_line_width = Float(0.0)

    def __init__(self, component=None, **kw):
        super().__init__(**kw)
        self.component = component

    def overlay(self, component, gc, view_bounds=None, mode="normal"):
        """Draws this component overlaid on another component.

        Implements AbstractOverlay.
        """
        if not self._visible:
            return

        plot = self.plot
        datasource = plot.color_data

        if self.selection_type == "range":
            selections = datasource.metadata["selections"]

            if selections is not None and len(selections) == 0:
                return

            low, high = selections
            if abs(high - low) / abs(high + low) < self.minimum_delta:
                return

            # Mask the data with just the points falling within the data
            # range selected on the colorbar
            data_pts = datasource.get_data()
            mask = (data_pts >= low) & (data_pts <= high)

        elif self.selection_type == "mask":
            mask = functools.reduce(
                logical_and, datasource.metadata["selection_masks"]
            )
            if sum(mask) < 2:
                return

        datasource.set_mask(mask)

        # Store the current plot color settings before overwriting them
        fade_outline_color = plot.outline_color_

        # Overwrite marker outline color and fill alpha settings of
        # the plot, then manually invoke the plot to draw onto the GC.
        plot.outline_color = list(self._old_outline_color[:3]) + [1.0]
        plot.fill_alpha = 1.0
        plot.line_width = self.selected_outline_width
        plot._draw_plot(gc, view_bounds, mode)

        # Restore the plot's previous color settings and data mask.
        plot.fill_alpha = self.fade_alpha
        plot.outline_color = fade_outline_color
        plot.line_width = self.unselected_outline_width
        datasource.remove_mask()

    def _component_changed(self, old, new):
        if old:
            old.observe(
                self.datasource_change_handler, "color_data", remove=True
            )
        if new:
            new.observe(self.datasource_change_handler, "color_data")
            self._old_alpha = new.fill_alpha
            self._old_outline_color = new.outline_color
            self._old_line_width = new.line_width

            self.datasource_change_handler(
                TraitChangeEvent(
                    object=new, name="color_data", old=None, new=new.color_data
                )
            )

    def datasource_change_handler(self, event):
        old, new = (event.old, event.new)

        if old:
            old.observe(
                self.selection_change_handler, "metadata_changed", remove=True
            )
        if new:
            new.observe(self.selection_change_handler, "metadata_changed")
            self.selection_change_handler(
                TraitChangeEvent(
                    object=new,
                    name="metadata_changed",
                    old=None,
                    new=new.metadata,
                )
            )

    def selection_change_handler(self, event):
        new = event.new

        if self.selection_type == "range":
            selection_key = "selections"
        elif self.selection_type == "mask":
            selection_key = "selection_masks"

        if (
            type(new) == dict
            and new.get(selection_key, None) is not None
            and len(new[selection_key]) > 0
        ):
            if not self._visible:
                # We have a new selection, so replace the colors on the plot
                # with the faded alpha and colors
                plot = self.plot

                # Save the line width and set it to zero for the unselected
                # points
                self._old_line_width = plot.line_width
                plot.line_width = self.unselected_outline_width
                # Save the outline color and set it to the faded version
                self._old_outline_color = plot.outline_color_
                outline_color = list(plot.outline_color_)
                if len(outline_color) == 3:
                    outline_color += [self.fade_alpha]
                else:
                    outline_color[3] = self.fade_alpha
                plot.outline_color = outline_color

                # Save the alpha value and set it to a faded version
                self._old_alpha = plot.fill_alpha
                plot.fill_alpha = self.fade_alpha

            self.plot.invalidate_draw()
            self._visible = True
        else:
            self.plot.fill_alpha = self._old_alpha
            self.plot.outline_color = self._old_outline_color
            self.plot.line_width = self._old_line_width
            self.plot.invalidate_draw()
            self._visible = False

        self.plot.request_redraw()

    def _get_plot(self):
        if self._plot is not None:
            return self._plot
        else:
            return self.component

    def _set_plot(self, val):
        self._plot = val