bokeh.palettes

def _get_categorical_palette(factors: List[str]) -> Dict[str, str]: # TODO: Find a good way of doing this -- also taking continuous vs categorical into account.
    n = max(3, len(factors))
    if n   <   11:
        from bokeh.palettes import Category10
        palette = Category10[n]
    elif n  <  21:
        from bokeh.palettes import Category20
        palette = Category20[n]
    else:
        from bokeh.palettes import viridis
        palette = viridis(n)

    return CategoricalColorMapper(factors=factors, palette=palette)


def _get_scatter_plot_elements(sp: Plot, source_pd: pd.DataFrame, label_cols: List[str], colors: Dict[str, ColorMapper] = None):

def assign_colors(neurons_df, color_by='cellBodyFiber'):
    """
    Use a random colortable to assign a color to each row,
    according to the column given in ``color_by``.

    NaN values are always black.

    Works in-place.
    """
    from bokeh.palettes import Turbo256
    colors = list(Turbo256)
    colors[0] = '#000000'
    color_categories = np.sort(neurons_df[color_by].fillna('').unique())
    assert color_categories[0] == ''

    np.random.seed(0)
    np.random.shuffle(color_categories[1:])
    assert color_categories[0] == ''

    while len(colors)   <   len(color_categories):
        colors.extend(colors[1:])

    color_mapping = dict(zip(color_categories, colors))
    neurons_df['color'] = neurons_df[color_by].fillna('').map(color_mapping)

def plot_multi_hail_hist(hist_data: Dict[str, hl.Struct],
                         title: str = 'Plot',
                         log: bool = False,
                         fill_color: Dict[str, str] = None,
                         outlier_fill_color: Dict[str, str] = None,
                         line_color: str = '#033649',
                         hover_mode: str = 'mouse',
                         hide_zeros: bool = False,
                         alpha: float = None) -> bokeh.plotting.Figure:
    """
    Plots multiple histograms on the same plot.
    Each histogram can (and should) come straight from ht.aggregate(hl.agg.hist(ht.data, start, end, bins))

    Example usage:
    plot_multi_hail_hist(ht.aggregate(hl.agg.group_by(ht.pop, hl.agg.hist(ht.data, start, end, bins))))

    :param dict of str -> Struct hist_data: Data to plot
    :param str title: Plot title
    :param bool log: Whether the y-axis should be log
    :param dict of str ->str fill_color: Color to fill the histogram bars that fall within the hist boundaries
    :param dict of str -> str outlier_fill_color: Color to fill the histogram bars that fall outside the hist boundaries
    :param str line_color: Color of the lines around the histogram bars
    :param str hover_mode: Hover mode; one of 'mouse' (default), 'vline' or 'hline'
    :param bool hide_zeros: Remove hist bars with 0 count
    :param float alpha: Alpha value (if None, then 1.0/len(hist_data) is used)
    :return: Histogram plot
    :rtype: Figure
    """

    low = int(log)

    if alpha is None:
        alpha = 1.0/len(hist_data)

    if fill_color is None:
        from bokeh.palettes import d3
        color_palette = d3['Category10'][max(3, len(hist_data))]
        fill_color = {hist_name: color_palette[i] for i, hist_name in enumerate(hist_data.keys())}

    if outlier_fill_color is None:
        outlier_fill_color = fill_color

    p = figure(title=title, y_axis_type="log", tools=TOOLS) if log else figure(title=title, tools=TOOLS)
    hists = []
    for label, hist in hist_data.items():

        data = {}
        distance = abs(hist.bin_edges[0] - hist.bin_edges[1])
        data['top'] = [x + low for x in hist.bin_freq]
        data['left'] = hist.bin_edges[:-1]
        data['right'] = hist.bin_edges[1:]
        data['color'] = [fill_color[label]] * len(hist.bin_freq)
        if hist.n_smaller > 0:
            data['top'].insert(0, hist.n_smaller + low)
            data['left'].insert(0, hist.bin_edges[0] - distance)
            data['right'].insert(0, hist.bin_edges[0])
            data['color'].insert(0, outlier_fill_color[label])
        if hist.n_larger > 0:
            data['top'].append(hist.n_larger + low)
            data['left'].append(hist.bin_edges[-1])
            data['right'].append(hist.bin_edges[-1] + distance)
            data['color'].append(outlier_fill_color[label])

        data['bottom'] = [low] * len(data['top'])
        data['label'] = [label] * len(data['top'])

        hist_source = ColumnDataSource(data)

        view = CDSView(source=hist_source, filters=[BooleanFilter([top > 0 for top in hist_source.data['top']])]) if hide_zeros else CDSView(source=hist_source)
        hists.append((label, [p.quad(top='top', bottom='bottom', left='left', right='right', source=hist_source, view=view, fill_color='color', alpha=alpha, line_color=line_color)]))

    tooltips = [("bin", "$index"), ("bin_edges", "(@left, @right)"), ("freq", "@top")]
    if len(hist_data) > 1:
        tooltips.insert(0, ('label', '@label'))
        p.add_layout(Legend(items=hists, location=(0, 0), orientation='horizontal', click_policy='hide'), 'above')
    p.select_one(HoverTool).tooltips = tooltips
    p.select_one(HoverTool).mode = hover_mode
    num_data_points = sum([sum(x.bin_freq) for x in hist_data.values()])
    p.add_layout(Title(text=f'({num_data_points:,} data points)'), 'above')

    return p


def plot_hail_hist_cumulative(hist_data: hl.Struct, title: str = 'Plot', normalize: bool = True,

def pair_plot(
        data: pd.DataFrame,
        label_col: str = None,
        colors: Union[List[str], Dict[str, str]] = None,
        tools: str = "save,pan,box_zoom,reset,wheel_zoom,box_select,lasso_select,help",
        tooltip_cols: List[str] = None
) -> Column:
    """

    Plots each column of `data` against each other and returns a grid of plots.
    The diagonal contains a histogram of each column, or a density plot if labels are provided.
    The lower diagonal contains scatter plots of each column against each other.
    The upper diagonal is empty.
    All columns should be numerical with the exception of the `label_col` if provided.
    If a color dict containing provided mapping labels to specific colors can be specified using `color_dict`

    :param DataFrame data: Dataframe to plot
    :param str label_col: Column of the DataFrame containing the labels
    :param list of str or dict of str -> str colors: RGB hex colors. If a dict is provided, it should contain the mapping of label to colors.
    :param str tools: Tools for the resulting plots
    :param list of str tooltip_cols: Additional columns that should be displayed in tooltip
    :return: Grid of plots (column of rows)
    :rtype: Column
    """
    from bokeh.palettes import viridis

    if tooltip_cols is None:
        tooltip_cols = [] if label_col is None else [label_col]
    elif label_col not in tooltip_cols:
        tooltip_cols.append(label_col)

    if label_col is None and colors is not None:
        logger.warn('`colors_dict` ignored since no `label_col` specified')

    colors_col = '__pair_plot_color'

    colors_dict = {}
    if label_col is None:
        data[colors_col] = viridis(1) * len(data)
    else:
        if not isinstance(colors, dict):
            labels = set(data[label_col])
            color_palette = viridis(len(labels)) if colors is None else colors
            colors_dict = {l: color_palette[i] for i, l in enumerate(labels)}
        else:
            colors_dict = colors
        data[colors_col] = [colors_dict.get(l, 'gray') for l in data[label_col]]
        tools = 'hover,' + tools

    data_cols = [c for c in data.columns if c not in [colors_col, label_col] + tooltip_cols]
    data_ranges = [DataRange1d(start=rmin - (abs(rmin - rmax) * 0.05), end=rmax + (abs(rmin - rmax) * 0.05)) for rmin, rmax in zip(data[data_cols].min(axis=0), data[data_cols].max(axis=0))]
    data_source = ColumnDataSource(data={c: data[c] for c in data.columns})

    n_cols = len(data_cols)

    plot_grid = []
    for i in range(n_cols):
        row = [None] * n_cols
        for j in range(i + 1):
            p = figure(
                x_axis_label=data_cols[j] if i == n_cols - 1 else '',
                y_axis_label=data_cols[i] if j == 0 else '',
                tools=tools
            )
            p.x_range = data_ranges[j]

            if i == j:
                if label_col is None:
                    hist, edges = np.histogram(data[data_cols[i]], density=False, bins=50)
                    p.quad(top=hist, bottom=0, left=edges[:-1], right=edges[1:])
                else:
                    density_data = data[[label_col, data_cols[i]]].groupby(label_col).apply(lambda df: np.histogram(df[data_cols[i]], density=True, bins=20))
                    for label, (hist, edges) in density_data.iteritems():
                        line_source = ColumnDataSource({'edges': edges[:-1], 'hist': hist, label_col: [label] * len(hist)})
                        p.line('edges', 'hist', color=colors_dict[label], legend=label_col, source=line_source)
                        p.select_one(HoverTool).tooltips = [(label_col, f'@{label_col}')]
            else:
                p.y_range = data_ranges[i]
                if label_col is not None:
                    p.circle(data_cols[j], data_cols[i], source=data_source, color=colors_col, legend=label_col)
                else:
                    p.circle(data_cols[j], data_cols[i], source=data_source, color=colors_col)
                if tooltip_cols:
                    p.select_one(HoverTool).tooltips = [(x, f'@{x}') for x in tooltip_cols]

            row[j] = p
        plot_grid.append(row)

    return gridplot(plot_grid, toolbar_location='left')