# 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']

# %% ../nbs/03_plot.ipynb 6
import warnings, string, itertools
from itertools import chain
from pathlib import Path
try:
    import pandas as pd, numpy as np, scipy
    import anndata as ad, scanpy as sc, scrublet as scr, scprep
    import phate, magic, graphtools as gt

    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 8
from typing import Any, NamedTuple, Optional

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


# %% ../nbs/03_plot.ipynb 12
from atyp import AnnData, Series, StrQ, CallQ, IterStr, IndexQ, Strs, NPArray, DataFrame

from nlit import (PROJECTION, COLOR_MAP, SUPTITLE)
from chck import isstr, isint, isnone

from atup import pack
from asto import asarr

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

# %% ../nbs/03_plot.ipynb 16
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 17
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 18
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 19
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 20
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 21
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 22
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 23
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 24
def figaxs(
    dpi = 150, figsize = (3, 3), nrows = 1, ncols = 1,
    gridspec_kw = dict(wspace=0.2, hspace=0.2), **kwargs
):
    kwargs.update(dpi=dpi, figsize=figsize, nrows=nrows, ncols=ncols, gridspec_kw=gridspec_kw)
    fig, axs = plt.subplots(**kwargs)
    axs = np.array(axs).reshape(-1)
    return fig, axs

# %% ../nbs/03_plot.ipynb 25
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 26
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 27
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

