import numpy as np

import bqplot.pyplot as plt
from bqplot import ColorScale, DateColorScale, OrdinalColorScale, ColorAxis

# setup data for plotting
np.random.seed(0)
n = 100
x_data = range(n)
y_data = np.cumsum(np.random.randn(n) * 100.0)

def create_fig(color_scale, color_data, fig_margin=None):
    # allow some margin on right for color bar
    if fig_margin is None:
        fig_margin = dict(top=50, bottom=70, left=50, right=100)
    fig = plt.figure(title="Up and Down", fig_margin=fig_margin)

    # setup color scale
    plt.scales(scales={"color": color_scale})

    # show color bar on right
    axes_options = {"color": {"orientation": "vertical", "side": "right"}}
    scat = plt.scatter(
        x_data, y_data, color=color_data, stroke="black", axes_options=axes_options
    )
    return fig

fig = create_fig(ColorScale(), y_data)
fig

"""
Explanation: ColorScale
The colors for the ColorScale can be defined one of two ways:
- Manually, by setting the scale's colors attribute to a list of css colors. They can be either:
   - html colors (link) 'white'
   - hex '#000000'
   - rgb 'rgb(0, 0, 0)'.
python
col_sc = ColorScale(colors=['yellow', 'red'])
- Using one of bqplot's color-schemes. As of now we support all the colorbrewer schemes (link), as well as the matplotlib schemes 'viridis', 'magma', 'inferno' and 'plasma'.
python
col_sc = ColorScale(scheme=['viridis'])
The color scale then linearly interpolates between its colors.
ColorAxis
A ColorAxis, like other Axis types, takes a color scale as input. It can then be displayed in a Figure.
python
ax_col = ColorAxis(scale=col_sc)
End of explanation
"""


color_scale = fig.marks[0].scales["color"]
color_scale.min = 0

color_scale.reverse = True

"""
Explanation: Attributes
ColorScales share attributes with the other Scale types:
- Their domain can be manually constrained with the min and max attributes
- They can be inversed by setting the reverse attribute to True
End of explanation
"""


color_scale.min = None
color_scale.mid = 0

"""
Explanation: Mid
In addition they also have a mid attribute, a value that will be mapped to the middle color. This is especially suited to diverging color schemes.
End of explanation
"""


import pandas as pd

fig_margin = dict(top=50, bottom=70, left=50, right=200)
date_col_sc = DateColorScale()
dates = pd.date_range(start="2015-01-01", periods=n)
create_fig(date_col_sc, dates, fig_margin=fig_margin)

date_col_sc.min = pd.datetime(2016, 2, 28)

"""
Explanation: DateColorScale
The DateColorScale is a color scale for dates. It works in every way like the regular ColorScale, except that its min, mid and max attributes — if defined — must be date elements (datetime, numpy or pandas).
End of explanation
"""


from bqplot.market_map import MarketMap

from ipywidgets import IntSlider, SelectionSlider, Dropdown
from ipywidgets import VBox, HBox, Layout
from traitlets import link

"""
Explanation: Color Schemes
Use the following widgets to browse through the available color schemes
End of explanation
"""


div_schemes = [
    "Spectral",
    "RdYlGn",
    "RdBu",
    "PiYG",
    "PRGn",
    "RdYlBu",
    "BrBG",
    "RdGy",
    "PuOr",
]

def scheme_inspector(fig, schemes, title=""):
    """
    Takes a Figure and a list of schemes and returns the Figure along with
    dropdown to go through the different schemes
    """

    # Get the color scale
    col_sc = fig.marks[0].scales["color"]

    # Create the widgets to select the colorscheme
    scheme_dd = Dropdown(description="Scheme", options=schemes)

    def update_scheme(*args):
        col_sc.scheme = scheme_dd.value

    scheme_dd.observe(update_scheme, "value")
    update_scheme()
    return VBox([scheme_dd, fig])


scheme_inspector(create_fig(ColorScale(), y_data), div_schemes)

"""
Explanation: Diverging schemes
End of explanation
"""


lin_schemes = [
    "OrRd",
    "PuBu",
    "BuPu",
    "Oranges",
    "BuGn",
    "YlOrBr",
    "YlGn",
    "Reds",
    "RdPu",
    "Greens",
    "YlGnBu",
    "Purples",
    "GnBu",
    "Greys",
    "YlOrRd",
    "PuRd",
    "Blues",
    "PuBuGn",
    "viridis",
    "plasma",
    "inferno",
    "magma",
]

scheme_inspector(
    create_fig(ColorScale(), y_data), lin_schemes, title="Non-diverging schemes"
)

"""
Explanation: Non-diverging schemes
End of explanation
"""


ord_schemes = {
    "Set2": 8,
    "Accent": 8,
    "Set1": 9,
    "Set3": 12,
    "Dark2": 8,
    "Paired": 12,
    "Pastel2": 8,
    "Pastel1": 9,
}

def partition(array, n_groups):
    n_elements = len(array)
    if n_groups > n_elements:
        return np.arange(n_elements)
    n_per_group = n_elements // n_groups + (n_elements % n_groups > 0)
    return np.tile(range(1, n_groups + 1), n_per_group)[:n_elements]

# Define the control widgets
n_groups_slider = IntSlider(description="n colors", min=3)
scheme_dd = Dropdown(description="Scheme", options=ord_schemes)


def update_scheme(*args):
    col_sc.scheme = scheme_dd.label
    ax_c.label = scheme_dd.label
    n_groups_slider.max = scheme_dd.value


def update_categories(*args):
    groups = partition(names, n_groups_slider.value)
    market_map.color = groups
    market_map.groups = groups


n_groups_slider.observe(update_categories, "value")
scheme_dd.observe(update_scheme)

# Define the bqplot marketmap
names = range(100)

col_sc = OrdinalColorScale()
ax_c = ColorAxis(scale=col_sc)

market_map = MarketMap(
    names=names,
    display_text=["" for _ in names],
    scales={"color": col_sc},
    axes=[ax_c],
    layout=Layout(min_width="800px", min_height="600px"),
)
update_scheme()
update_categories()

VBox([HBox([scheme_dd, n_groups_slider]), market_map])

"""
Explanation: OrdinalColorScale
The OrdinalColorScale is a color scale for categorical data, i.e. data that does not have an intrinsic order.
The scale colors may be specified by the user, or chosen from a set scheme. As of now, the supported color schemes are the colorbrewer categorical schemes, listed here along with their maximum number of colors.
End of explanation
"""


