# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_plot.ipynb.

# %% auto 0
__all__ = ['SC_PARAMS', 'RC_PARAMS', 'reset_plt', 'setfigdir', 'makegif', 'suptitle', 'update_legend_texts', 'embplot3x',
           'boolplot', 'category_legend', 'figaxs', 'loc2fnc', 'getboxloc', 'textboxes', 'embplot', 'emphplot']

# %% ../nbs/03_plot.ipynb 6
import warnings
from pathlib import Path

# %% ../nbs/03_plot.ipynb 8
#| export


# %% ../nbs/03_plot.ipynb 10
try: import numpy as np, pandas as pd
except ImportError: warnings.warn("Some packages are not installed")

# %% ../nbs/03_plot.ipynb 11
try:
    import matplotlib as mpl, matplotlib.pyplot as plt, seaborn as sns
    from matplotlib.colors import LinearSegmentedColormap
    from matplotlib.patches import Patch
except ImportError:
    warnings.warn("Some packages are not installed")

# %% ../nbs/03_plot.ipynb 12
try: import scanpy as sc, scprep    
except ImportError: warnings.warn("Some packages are not installed")

# %% ../nbs/03_plot.ipynb 14
from nlit import (PROJECTION, COLOR_MAP, SUPTITLE)
from chck import isstr, isint, isnone, notnone
from atup import pack

# %% ../nbs/03_plot.ipynb 16
from xsc.attr import (
    getkey, getemb, getcategoryidx,
)
from .sens import (senspath)

# %% ../nbs/03_plot.ipynb 18
SC_PARAMS = dict(
    figsize        = (figsize := (4, 4)),
    dpi            = (dpi := 150),
    fontsize       = (fontsize := 10),
    ipython_format = (ipython_format := 'retina'),
    color_map      = (color_map := 'magma'),
    transparent    = (transparent := False),
)

RC_PARAMS = {
    'figure.figsize': figsize,
    'figure.dpi': dpi,
    'font.size': fontsize,

    'axes.labelsize':   fontsize - 2,
    'figure.labelsize': fontsize - 2,
    'axes.axisbelow': True, 'axes.grid': True,
    'figure.frameon': False,

    'legend.title_fontsize': 0,
    'legend.labelspacing': 1,
}

# %% ../nbs/03_plot.ipynb 19
def reset_plt(rc_params: dict = {}, sc_params: dict = {}):
    plt.rcdefaults()
    plt.rcParams.update(**{**RC_PARAMS, **rc_params})
    sc.set_figure_params(**{**SC_PARAMS, **sc_params})

def setfigdir(path: Path, sens: int = None):
    path = senspath(path, sens)
    sc.settings.figdir = str(path)
    if not path.is_dir():
        path.mkdir(exist_ok=True, parents=True)

# %% ../nbs/03_plot.ipynb 20
try:
    black_cmap = LinearSegmentedColormap.from_list('black', [(0, 0, 0), (0, 0, 0)])
    white_cmap = LinearSegmentedColormap.from_list('white', [(1, 1, 1), (1, 1, 1)])
except NameError:
    warnings.warn("LinearSegmentedColormap not imported from `matplotlib.colors`")

# %% ../nbs/03_plot.ipynb 21
def makegif(adata, embkey, catkey, file: str, outdir: Path, overwrite: bool = False):
    file = file or (outdir / f'{embkey}_{catkey}.gif')
    if file.is_file() and not overwrite: return
    return scprep.plot.rotate_scatter3d(adata.obsm[embkey], c=adata.obs[catkey], filename=str(file))

# %% ../nbs/03_plot.ipynb 22
def suptitle(title: str = None, ax = None, fig = None, **kwargs):    
    if fig is None and ax is not None: fig = ax.figure
    if title: fig.suptitle(title)
    return ax

def update_legend_texts(ax, **updates):
    legend_texts = ax.get_legend()
    if legend_texts is None: return
    for legend_text in legend_texts.get_texts():
        if (text := legend_texts.get_text()) in updates:
            legend_text.set_text(updates[text])

# %% ../nbs/03_plot.ipynb 23
def embplot3x(adata, basis, color, ax = None, *, size: int = 100, emph: int = 100, bord: int = 0.8, **kwargs):
    if ax is None: _, ax = plt.subplots()
    s = size = kwargs.pop('s', size)
    b = size + emph # black dots
    w = b * bord    # white dots    
    proj = kwargs.get(PROJECTION, '2d')
    cmap = kwargs.get(COLOR_MAP, 'magma')
    if (sup := kwargs.pop(SUPTITLE, None)): suptitle(sup, ax)

    sc.pl.embedding(
        adata, basis=basis, color=color, sort_order=True, ax=ax, size=b, s=b, show=False,
        color_map=black_cmap, projection=proj, legend_loc=None, colorbar_loc=None, 
        
    )
    sc.pl.embedding(
        adata, basis=basis, color=color, sort_order=True, ax=ax, size=w, s=w, show=False,        
        color_map=white_cmap, projection=proj, legend_loc=None, colorbar_loc=None,         
    )
    
    sc.pl.embedding(
        adata, basis=basis, color=color, sort_order=True, ax=ax, size=s, s=s,      
        color_map=cmap, projection=proj, **kwargs
    )

# %% ../nbs/03_plot.ipynb 24
def boolplot(
    adata, basis, color, bidx, *, 
    figsize: tuple = (4, 4), dpi: int = 150, ncols: int = 1,
    size: int = 100, background: bool = True, emphasize: bool = True, 
    gridspec_kw = dict(wspace=1., hspace=1.),
    **kwargs
):    
    ncols = max(ncols, 1)
    nrows = len(color) // ncols
    figsize = (figsize[0] * ncols, figsize[1] * nrows)
    
    _, axs = plt.subplots(dpi=dpi, figsize=figsize, nrows=nrows, ncols=ncols, gridspec_kw=gridspec_kw)
    axs = np.array(axs)
    kwargs.update(show = False)
    
    for i, ax in enumerate(axs.flat):
        if background: sc.pl.embedding(adata[~bidx], basis=basis, ax=ax, size=size, show=False)
        plotfn = embplot3x if emphasize else sc.pl.embedding
        if (sup := kwargs.pop(SUPTITLE, None)): suptitle(sup, ax)
        plotfn(adata[bidx], basis=basis, color=color[i], ax=ax, size=size, **kwargs)

# %% ../nbs/03_plot.ipynb 25
def category_legend(
    categories, palette: str = 'husl', fmtfunc = None,
    patch_kw = dict(), legend_kw = dict(),
):
    palette = sns.color_palette(palette, len(categories)) if isstr(palette) else palette
    fmtfunc = str if fmtfunc is None else fmtfunc

    for k in ('facecolor', 'edgecolor', 'label'): patch_kw.pop(k, None)
    legend_elements = [
        Patch(facecolor=color, edgecolor=color, label=fmtfunc(cat), **patch_kw)
        for cat, color in zip(categories, palette)
    ]
    legdefs = dict(fontsize='small', bbox_to_anchor=(1.05, 0.5), loc='upper left', title='Categories',)
    for k, v in legdefs.items(): legend_kw.setdefault(k, v)
    legend_kw.update(handles=legend_elements)
    cat_legend = plt.legend(**legend_kw)
    return cat_legend


# %% ../nbs/03_plot.ipynb 26
# NOTE: never used
def scaleratio(ratio: float = 1.5, scale: float = 1) -> tuple[float, float]:
    return (ratio * scale, scale)

# %% ../nbs/03_plot.ipynb 27
def figaxs(
    dpi = 150, figsize = (3, 3), nrows = 1, ncols = 1, 
    subplot_kw = dict(projection='2d'), gridspec_kw = dict(wspace=0.2, hspace=0.2), 
    *, retone: bool = False, **kwargs
):
    gridspec_kw.setdefault('wspace', 0.2) 
    gridspec_kw.setdefault('hspace', 0.2)

    is3d = kwargs.pop('is3d', (subplot_kw.pop('projection', False)))
    proj = '3d' if (is3d) else '2d'
    if is3d: subplot_kw.setdefault('projection', proj)

    kwargs.update(dpi=dpi, figsize=figsize, nrows=nrows, ncols=ncols, gridspec_kw=gridspec_kw, subplot_kw=subplot_kw)
    fig, axs = plt.subplots(**kwargs)
    axs = np.array(axs).reshape(-1)

    if retone and len(axs) == 1: return fig, axs[0]
    return fig, axs

# %% ../nbs/03_plot.ipynb 28
def loc2fnc(loc: str):
    if loc is None: return
    pos = loc.lower().strip(' ')
    fmax = ('upper', 'top')
    fmin = ('lower', 'bottom', 'bot')
    fave = ('mid', 'middle', 'center', 'centre', 'centroid')
    if pos in fmax: return np.max
    elif pos in fmin: return np.min
    elif pos in fave: return np.mean
    elif pos in ('median', ): return np.median
    else: None

def getboxloc(data, x, y, loc: str = 'centroid'):    
    locn = '_'.join(loc.lower().split(' '))
    xarr, yarr = data.get(x), data.get(y)
    xstr, ystr = pack(*locn.split('_'), n=2)
    xfnc, yfnc = loc2fnc(xstr), loc2fnc(ystr)
    if yfnc is None: yfnc = xfnc
    return xfnc(xarr), yfnc(yarr)

# %% ../nbs/03_plot.ipynb 29
def textboxes(coords: dict, palette: str = None, *, text_kw: dict = dict(), bbox_kw: dict = dict(),  **kwargs):
    (text := dict(fontsize=12, ha='center', va='center', )).update(text_kw)
    (bbox := dict(facecolor='white', boxstyle='round,pad=0.5', alpha=0.5)).update(bbox_kw, **text.pop('bbox', {}))

    palette = sns.color_palette(palette, len(coords)) if isstr(palette) else palette
    for i, (key, (x, y, *rest)) in enumerate(coords.items()):
        color = palette[i] if isnone(rest) or len(rest) == 0 else rest[0]
        bbox.update(edgecolor=color)
        plt.text(x, y, str(key), color=color, **text, bbox=bbox)

# %% ../nbs/03_plot.ipynb 30
def embplot(
    adata, embkey: str, catkey: str, huekey: str, 
    x = 0, y = 1, palette: str = 'husl',
    *, fig_kw: dict = {}, ax_kw: dict = {}, kde_kw: dict = {}, 
    text_kw: dict = {}, bbox_kw: dict = {},
    patch_kw: dict = {}, legend_kw: dict = {},
    kde_textboxes: bool = False, **kwargs
):
    fig, axs = figaxs(**fig_kw)
    df_emb = getemb(adata, embkey)
    axcols = df_emb.columns
    x = axcols.values[x] if isint(x) else axcols.get_loc(x)
    y = axcols.values[y] if isint(y) else axcols.get_loc(y)

    categories = getkey(adata, catkey).astype('category').cat.categories.unique().values    
    palette = sns.color_palette(palette, len(categories)) if isstr(palette) else palette
    
    defs = dict(fill=True, alpha=0.3)
    for k, v in defs.items(): kde_kw.setdefault(k, v)
    
    defs = dict(alpha=kde_kw.get('alpha'))
    for k, v in defs.items(): bbox_kw.setdefault(k, v)
                   
    centroids = {}
    for idx, cat in enumerate(categories):
        bidx = getcategoryidx(adata.obs, catkey, cat)
        boxx, boxy = getboxloc(df_emb[bidx], x, y, loc='upper right')    
        centroids[cat] = (boxx, boxy)
        color = palette[idx]
        sns.kdeplot(df_emb[bidx], x=x, y=y, color=color, label=cat, ax=axs[0], **kde_kw)

    df_emb.loc[:, huekey] = getkey(adata, huekey)
    # return df_emb
    sns.scatterplot(df_emb, x=x, y=y, hue=huekey, edgecolor='w', ax=axs[0])
  
    hue_legend = plt.legend(title=huekey, loc='upper left', bbox_to_anchor=(1.05, 1))
    if kde_textboxes: textboxes(centroids, palette=palette, text_kw=text_kw, bbox_kw=bbox_kw)
    cat_legend = category_legend(categories, palette=palette, patch_kw=patch_kw, legend_kw=legend_kw)
    plt.gca().add_artist(hue_legend)
    plt.tight_layout(rect=[0, 0, 0.85, 1])
    return fig, axs


# %% ../nbs/03_plot.ipynb 31
def emphplot(
    adata, 
    *points: tuple,
    basis: str = 'X_phate',
    x: int = 0, y: int = 1, z: int = None,
    fig_kw: dict = {}, 
    size: int = 10,
    emphasize: int = 250,
    alpha: float = 0.5,
    emphaalpha: int = 0.8,
    palette: str = 'tab10',
    oneidx: bool = True,
    **kwargs
):    
    df_emb = getemb(adata, basis, oneidx=oneidx).reset_index(drop=True)    
    df_emb = df_emb.sort_index(key = lambda i: np.isin(i, points), ascending = False)
    
    axcols = df_emb.columns
    xyz = list(filter(notnone, (x, y, z)))
    xyz = [axcols.values[a] if isint(a) else axcols.get_loc(a) for a in xyz]
    
    fig, ax = figaxs(**fig_kw, is3d = len(xyz) > 2, retone=True)
    
    p, n = len(points), len(df_emb)
    palette = sns.color_palette(palette, 2) if isstr(palette) else palette

    s = np.vstack((np.ones((p, 1)) * emphasize,  np.ones((n - p, 1)) * size))
    c = np.vstack((np.ones((p, 3)) * palette[-1], np.ones((n - p, 3)) * palette[0]))
    
    for a, label in zip('xyz', xyz):
        if isnone(setlabel := getattr(ax, f'set_{a}label', None)): continue
        setlabel(label, fontsize=8)
        
    ax.scatter(*tuple(df_emb[p:][dim] for dim in xyz), c=c[p:], s=s[p:], alpha=alpha, **kwargs)
    ax.scatter(*tuple(df_emb[:p][dim] for dim in xyz), c=c[:p], s=s[:p], alpha=emphaalpha, **kwargs)
    
    ax.tick_params(labelsize=6)
    return fig, ax
