# (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 ArrayPlotData.
"""
from numpy import array, ndarray

# Enthought library imports
from traits.api import Dict

# Local, relative imports
from .abstract_plot_data import AbstractPlotData
from .abstract_data_source import AbstractDataSource


class ArrayPlotData(AbstractPlotData):
    """A PlotData implementation class that handles a list of Numpy arrays
    (or a 2-D Numpy array).

    By default, it doesn't allow its input data to be modified by downstream
    Chaco components or interactors.
    """

    # -------------------------------------------------------------------------
    # Public traits
    # -------------------------------------------------------------------------

    #: Map of names to arrays.  Although there is no restriction on the array
    #: dimensions, each array must correspond to a single plot item; that
    #: is, a single name must not map to a multi-dimensional array unless
    #: the array is being used for an image plot or for something that can handle
    #: multi-dimensional input data.
    arrays = Dict

    #: Consumers can write data to this object (overrides AbstractPlotData).
    writable = True

    def __init__(self, *data, **kw):
        """ArrayPlotData can be constructed by passing in arrays.

        Keyword arguments can be used to give certain arrays specific names;
        unnamed arrays are given a generic name of the format 'seriesN', where
        N is its position in the argument list.

        For example::

            ArrayPlotData(array1, array2, index=array3, foo=array4)

        This call results in the creation of four entries in self.arrays::

            'series1' -> array1
            'series2' -> array2
            'index'   -> array3
            'foo'     -> array4

        If any names in the keyword parameter list collide with the
        auto-generated positional names "series1", "series2", etc., then those
        arrays are replaced.

        Note that this factor means that keyword traits are *not* set using the
        keyword parameters in the constructor. This strategy defies some
        conventions, but was it chosen for convenience, since the raison d'etre
        of this class is convenience.
        """
        super().__init__()
        self._update_data(kw)
        data = dict(zip(self._generate_names(len(data)), data))
        self._update_data(data)

    # ------------------------------------------------------------------------
    # AbstractPlotData Interface
    # ------------------------------------------------------------------------

    def list_data(self):
        """Returns a list of the names of the arrays managed by this instance."""
        return list(self.arrays.keys())

    def get_data(self, name):
        """Returns the array associated with *name*.

        Implements AbstractDataSource.
        """
        return self.arrays.get(name, None)

    def del_data(self, name):
        """Deletes the array specified by *name*, or raises a KeyError if
        the named array does not exist.
        """
        if not self.writable:
            return None

        if name in self.arrays:
            del self.arrays[name]
            self.data_changed = {"removed": [name]}
        else:
            raise KeyError("Data series '%s' does not exist." % name)

    def set_data(self, name, new_data, generate_name=False):
        """Sets the specified array as the value for either the specified
        name or a generated name.

        If the instance's `writable` attribute is True, then this method sets
        the data associated with the given name to the new value, otherwise it
        does nothing.

        Parameters
        ----------
        name : string
            The name of the array whose value is to be set.
        new_data : array
            The array to set as the value of *name*.
        generate_name : Boolean
            If True, a unique name of the form 'seriesN' is created for the
            array, and is used in place of *name*. The 'N' in 'seriesN' is
            one greater the largest N already used.

        Returns
        -------
        The name under which the array was set.

        See Also
        --------
        update_data: Use if needing to set multiple ArrayPlotData entries at
            once, for example because new arrays' dimensions change and
            updating one at a time would break an existing Plot.
        """
        if not self.writable:
            return None

        if generate_name:
            names = self._generate_names(1)
            name = names[0]

        self.update_data({name: new_data})
        return name

    def update_data(self, *args, **kwargs):
        """Updates any number of arrays before triggering a `data_changed`
        event.

        Useful to set multiple ArrayPlotData entries at once, for example
        because new arrays' dimensions change and updating one at a time would
        break an existing Plot.

        Note: Implements AbstractPlotData's update_data() method.  This method
        has the same signature as the dictionary update() method.

        See Also
        --------
        set_data: Simpler interface to set only 1 entry at a time.
        """
        if not self.writable:
            return None

        data = dict(*args, **kwargs)
        event = {}
        for name in data:
            if name in self.arrays:
                event.setdefault("changed", []).append(name)
            else:
                event.setdefault("added", []).append(name)

        self._update_data(data)
        self.data_changed = event

    def set_selection(self, name, selection):
        """Overrides AbstractPlotData to do nothing and not raise an error."""
        pass

    # ------------------------------------------------------------------------
    # Private methods
    # ------------------------------------------------------------------------

    def _generate_names(self, n):
        """Generate n new names"""
        max_index = max(self._generate_indices())
        names = [
            "series{0:d}".format(n)
            for n in range(max_index + 1, max_index + n + 1)
        ]
        return names

    def _generate_indices(self):
        """Generator that yields all integers that match "series%d" in keys"""
        yield 0  # default minimum
        for name in self.list_data():
            if name.startswith("series"):
                try:
                    v = int(name[6:])
                except ValueError:
                    continue
                yield v

    def _update_data(self, data):
        """Update the array, ensuring that data is an array"""
        # note that this call modifies data, but that's OK since the callers
        # all create the dictionary that they pass in
        for name, value in list(data.items()):
            if not isinstance(value, (ndarray, AbstractDataSource)):
                data[name] = array(value)
            else:
                data[name] = value

        self.arrays.update(data)