# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_plots.ipynb.

# %% auto 0
__all__ = ['set_rc_params', 'scale_figsize_for_grid', 'setup_grid_figure', 'mpl_cur_cmap', 'mpl_cat_val',
           'convert_cmap_to_categorical', 'make_categorical_legend', 'hide_axes', 'triple_scatter']

# %% ../nbs/04_plots.ipynb 4
import os
import numpy as np, pandas as pd

import matplotlib as mpl, matplotlib.pyplot as plt, seaborn as sns
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
from matplotlib.lines import Line2D

from typing import Optional, List, Union, Tuple

# %% ../nbs/04_plots.ipynb 5
from .types import ColorMap, ndarray
from .static import RC_PARAMS
from .utils import (filter_kwargs_for_func, filter_kwargs_for_class)

# %% ../nbs/04_plots.ipynb 6
def set_rc_params(params: Optional[dict] = None):
    param_dict = RC_PARAMS.copy()
    if params is not None:        
        param_dict.update(params)
    plt.rcParams.update(param_dict)

def scale_figsize_for_grid(
    figsize:Optional[Tuple[int, int]]=(12, 8), 
    n_cols:Optional[int]=1, n_rows:Optional[int]=1,
) -> Tuple[int, int]:
    grid_figsize = (figsize[0] * n_cols, figsize[1] * n_rows)
    return grid_figsize

def setup_grid_figure(
    ncols: Optional[int] = 1, nrows: Optional[int] = 1,
    figsize: Optional[Tuple[int, int]] = (12, 8), 
    dpi: Optional[int] = 300, hspace: Optional[float] = 0.3, wspace: Optional[float] = None,
    frameon: Optional[bool] = False, **kwargs
) -> Tuple[plt.Figure, plt.GridSpec]:
    
    grid_figsize = scale_figsize_for_grid(figsize, ncols, nrows)
    params = dict(
        ncols=ncols, nrows=nrows, figsize=grid_figsize,
        dpi=dpi, hspace=hspace, wspace=wspace, frameon=frameon
    )
    params.update(kwargs)
    
    f_params = filter_kwargs_for_func(plt.figure, **params)
    g_params = filter_kwargs_for_class(plt.GridSpec, **params)
    
    fig = plt.figure(None, **f_params)
    gspec = plt.GridSpec(**g_params)    
    return fig, gspec

def mpl_cur_cmap(cmap: Optional[ColorMap] = None) -> ColorMap:
    if cmap is None:
        fig = plt.gcf()
        cmap = fig.get_cmap() 
    return cmap

def mpl_cat_val(i: int, n: int) -> float:
    adj = 1 if n > 1 else 0
    return i / (n - adj)

def convert_cmap_to_categorical(values: List[str], cmap: Optional[ColorMap]) -> ListedColormap:
    unique = np.unique(values)
    n_vals = len(unique)
    
    # Generate a list of colors from the original colormap
    cmap = mpl_cur_cmap(cmap)
    colors = [cmap(mpl_cat_val(i, n_vals)) for i in range(n_vals)]
    
    # Create a new colormap from the list of colors
    new_cmap = ListedColormap(colors)
    
    return new_cmap

def make_categorical_legend(
    values: List[str], cmap: Optional[ColorMap], **kwargs
) -> mpl.legend.Legend:
    
    # NOTE: c will be the ground truth labels from pd.factorize (e.g. [0.0,0,0,1,0,1,1,1])    
    unique = sorted(np.unique(values))
    cmap = mpl_cur_cmap(cmap)
    cmap = convert_cmap_to_categorical(unique, cmap)

    line_opts = dict(marker='o', markersize=15)
    line_opts.update(kwargs)

    legend_elements = [        
        Line2D(
            [0], [0], color = cmap(i), markerfacecolor = cmap(i),
            label = label, **line_opts
        )
        for i, label in enumerate(unique)
    ]

    leg = plt.legend(handles=legend_elements)
    return leg

# %% ../nbs/04_plots.ipynb 7
def hide_axes(ax:mpl.axes.Axes) -> mpl.axes.Axes:
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.get_xaxis().get_major_formatter().set_scientific(False)
    ax.get_yaxis().get_major_formatter().set_scientific(False)
    kwargs = dict(bottom=False, left=False, labelbottom=False, labelleft=False)
    ax.tick_params(which='both', **kwargs)
    return ax

def triple_scatter(
    data: ndarray, ax: mpl.axes.Axes,
    c: Union[List[str], List[float]], s:List[float],
    cmap: Optional[ColorMap] = None,
    alpha: float = 1, marker: str = '.',
    point_size: float = 300, hide_axes: bool = True,
    frameon: bool = False,
):
    x, y = zip(*data)
    cmap = mpl_cur_cmap(cmap)
    
    # Handle size of points for bottom black layer, and middle white layer
    outline_width = (0.3, 0.05)
    bg_width, gap_width = outline_width
    
    point = np.sqrt(point_size)
    gap_size = (point + (point * gap_width) * 2) ** 2
    bg_size = (np.sqrt(gap_size) + (point * bg_width) * 2) ** 2

    scatter_params = dict(
        s = bg_size * np.sqrt(s),
        alpha = alpha, marker = marker, 
        linewidths = 0, edgecolors = None
    )

    scatter_params.update(c='black')
    ax.scatter(x, y, **scatter_params)

    scatter_params.update(c='white')
    ax.scatter(x, y, **scatter_params)

    scatter_params.update(
        c=c, s = point_size * np.sqrt(s),
        alpha = (0.7 * alpha),
        cmap = cmap,
    )    
    
    pnts = ax.scatter(x, y, **scatter_params)
    if hide_axes:
        ax = hide_axes(ax)

    ax.set_frame_on(frameon)
    ax.patch.set_alpha(0)
    return pnts
