from .Panel import *

__all__ = ["BubblePanel"]

default_size = plt.matplotlib.rcParams["lines.markersize"] ** 2


class BubblePanel(Panel):
    """
    BubblePanel is a general wrapper for making scatter plots
    where planets are represented as bubbles that can have
    informative sizes and/or colors.
    """

    def __init__(
        self,
        xaxis=None,
        yaxis=None,
        size=None,
        size_normalization=None,
        color=None,
        cmap="plasma",
        vmin=None,
        vmax=None,
        color_normalization=None,
        **kw,
    ):
        """
        Initialize a plotting panel.

        Parameters
        ----------
        size : PlottableAxis, str, float, None
            What should the sizes of points be or encode?
        size_normalization : float
            If sizes depend on quantities,
            how should they be normalized?
        color : PlottableAxis, str, float, None
            What should the colors of points be or encode?
        cmap : str, cmap from plt.matplotlib.cm
            If the colors depend on quantities,
            what cmap should be used for them?
        vmin : float, astropy.units.quantity.Quantity
            If the colors depend on quantities,
            what should the bottom of the cmap be?
        vmax : float, astropy.units.quantity.Quantity
            If the colors depend on quantities,
            what should the top of the cmap be?
        color_normalization : matplotlib.colors.Normalize
            If color depend on quantities, how should
            the values be normalized. If color_normalization
            is defined, any values provided here for
            vmin and vmax will be ignored.
        **kw : dict
            Other keywords will be passed on to *all*
            Panel/Plottable initializations (which may
            include x, y, size, and color). If you need
            more fine-grained control over which axis
            gets which keyword, consider initializing
            those panels one-by-one.
        """

        # initialize the basics of the panel with the plottable axes
        Panel.__init__(self, xaxis=xaxis, yaxis=yaxis, **kw)

        # set up how we should scale the sizes of points
        size = clean_axis(size)
        try:
            # try to make a variable size axis
            self.plottable["size"] = size(panel=self, **kw)
            default_size_normalization = self.plottable["size"].size_normalization

        except TypeError:
            # otherwise, use a single size for all points
            self.plottable["size"] = size
            default_size_normalization = 1

        # self.plottable['x'].panel = self
        # self.plottable['y'].panel = self

        # make sure a size normalization has been defined
        self.size_normalization = size_normalization or default_size_normalization

        # set up how we should set the colors of points
        color = clean_axis(color)
        try:
            # try to make a variable color axis
            self.plottable["color"] = color(panel=self, **kw)
            default_lim = self.plottable["color"].lim
        except TypeError:
            # otherwise, use a single color for all points
            self.plottable["color"] = color
            default_lim = [None, None]

        # if an actual cmap was provided, use it
        if isinstance(cmap, plt.matplotlib.colors.Colormap):
            self.cmap = cmap
        # otherwise, treat the cmap as a string key
        else:
            self.cmap = plt.matplotlib.colormaps[cmap]

        # make sure the color map limits are set
        self.vmin = vmin or default_lim[0]
        self.vmax = vmax or default_lim[1]

        # if a custom normalization is used, reset vmin + vmax
        self.color_normalization = color_normalization
        if isinstance(self.color_normalization, plt.matplotlib.colors.Normalize):
            # pull the normalization's min/max for information
            self.vmin = color_normalization.vmin
            self.vmax = color_normalization.vmax

        # apply (x,y) axis labels, scales, limits appropriately
        for axis in "xy":
            for attribute in ["label", "scale", "lim"]:
                setattr(
                    self, f"{axis}{attribute}", getattr(self.plottable[axis], attribute)
                )

        # DEBUG
        self.summarize()

    def get_sizes(self):
        """
        The sizes of the bubbles.

        Returns
        -------
        s : an input for plt.scatter
            Either a single scalar, or an array with variable
            sizes for each bubble according to some quantity.
        """

        # should we ignore any variable size instructions?
        if self.pop.respond_to_size == False:
            size = self.pop._plotkw.get("s", None)
        # if desired, set variable sizes
        elif isinstance(self.plottable["size"], PlottableAxis):
            # get the raw values for the sizes
            x = self.plottable["size"].value()
            # calculate the normalized size
            size = default_size * x / self.size_normalization
        # otherwise, set a single size
        else:
            # get default, first from pop and then from panel
            size = self.pop._plotkw.get("s", self.plottable["size"])

        # return a valid input to plt.scatter(s=...)
        return size

    def get_colors(self):
        """
        The colors of the bubbles.

        Returns
        -------
        c : an input for plt.scatter
            Either a single color, or an array with variable
            colors for each bubble according to some quantity.
        """

        # should we ignore any variable color instructions?
        if self.pop.respond_to_color == False:
            color = self.pop.color
        # should we use a variable color?
        elif isinstance(self.plottable["color"], PlottableAxis):
            # get the raw values to go into the color
            x = self.plottable["color"].value()

            # FIXME - make sure to check vmin/vmax are valid
            # if (self.vmin is None) or (self.vmax is None):
            #    raise AtlasError(f'''
            #    It looks like you're trying to use
            #    {self.plottable['color']} to set variable
            #    colors for bubbles. To do so, please make
            #    sure it has finite values defined for its
            #    .vmin and .vmax attributes.
            #    ''')

            # make sure we have *some* normalizer defined
            f = plt.matplotlib.colors.Normalize
            self.color_normalization = self.color_normalization or f(
                vmin=self.vmin, vmax=self.vmax
            )

            normalized = self.color_normalization(x)
            color = self.cmap(normalized)
        # finally, should we just use a default color?
        else:
            # get default, first from pop and then from panel
            color = self.pop.color
            if color is None:
                color = self.plottable["color"]

        # return a valid input to any one of the following:
        #   plt.scatter(c=...)
        #   plt.scatter(edgecolors=...)
        #   plt.scatter(facecolors=...)
        return color

    def kw(self, key=None, **kwargs):
        """
        Do a little decision-making about the plotting keyword
        arguments, pulling defaults from each population where
        needed.

        Parameter
        ---------
        key : str
            The population for which we should pull keywords.
            If None, go with the current population.
        **kwargs : dict
            All other keywords will be directed toward
            overwriting individual population defaults.
        """

        # identify the population we're working with
        if key is None:
            key = self.key
        # else:
        self.point_at(key)

        # define some default keywords, which can be over-written
        default = dict(
            s=self.get_sizes(),
            marker=self.pop.marker,
            linewidth=self.pop.linewidth,
            alpha=self.pop.alpha,
            zorder=self.pop.zorder,
            label=self.pop.label,
        )

        # sort out whether faces and/or edges should get color
        c = self.get_colors()
        if self.pop.filled:
            default["facecolors"] = c
        else:
            default["facecolors"] = "none"
        if self.pop.outlined:
            default["edgecolors"] = c
        else:
            default["edgecolors"] = "none"

        # if any other keywords are provided, overwrite these defaults
        for k, v in kwargs.items():
            default[k] = v

        return default

    def plot(self, key, ax=None, labelkw={}, **kwargs):
        """
        Add the points for a particular population to this panel.

        Parameters
        ----------
        key : str
            The population (as an item in the self.pops dictionary) to add.
        ax :
            Into what ax should we place this plot?
            If None, use default.
        labelkw : dict
            Keywords for labeling the planet names.
        **kwargs : dict
            Any extra keywords will be passed on to `scatter`
        """

        # focus attention on that population
        self.point_at(key)

        # make sure we're plotting into the appropriate axes
        try:
            plt.sca(self.ax)
        except AttributeError:
            self.setup(ax=ax)

        # add the scattered points
        self.scattered[key] = self.ax.scatter(self.x, self.y, **self.kw(key, **kwargs))

        # set the scales, limits, labels
        self.finish_plot(labelkw=labelkw)
