# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_attr.ipynb.

# %% auto 0
__all__ = ['getdata', 'getcol', 'getcatcol', 'getcategory', 'getcategoryidx', 'getcatval', 'getcatidx', 'gene2idx', 'gene2val',
           'genes2idxs', 'genes2vals', 'findgenes', 'gene2eid', 'genes2str', 'clusteridx', 'conditionidx',
           'subset_genes', 'axeschars', 'labelaxis', 'labelaxes', 'getemb']

# %% ../nbs/00_attr.ipynb 6
import io, warnings, string, itertools
from contextlib import suppress, redirect_stderr

# %% ../nbs/00_attr.ipynb 8
from typing import List, Tuple

# %% ../nbs/00_attr.ipynb 10
try: from tqdm.auto import tqdm
except ImportError: warnings.warn("Some packages are not installed")

# %% ../nbs/00_attr.ipynb 11
try: import pandas as pd, numpy as np, scipy
except ImportError: warnings.warn("Some packages are not installed")

# %% ../nbs/00_attr.ipynb 12
try: import anndata as ad, scanpy as sc, scrublet as scr, scprep
except ImportError: warnings.warn("Some packages are not installed")

# %% ../nbs/00_attr.ipynb 13
try: import phate, magic, graphtools as gt
except ImportError: warnings.warn("Some packages are not installed")

# %% ../nbs/00_attr.ipynb 14
try: import matplotlib as mpl, matplotlib.pyplot as plt, seaborn as sns
except ImportError: warnings.warn("Some packages are not installed")

# %% ../nbs/00_attr.ipynb 16
from atyp import AnnData, Series, StrQ, CallQ,  IndexQ, Strs, NPArray, DataFrame
from chck import isdf, notnone, isseries, isnone

from nchr import U1
from nlit import X_
from nlit import (GENE_SYMBOL, ENSEMBL_ID, CATEGORY, CLUSTER, CON)
from pstr import shiftchars
from asto import asarr
from akey import getkey, numcols, keydims

# %% ../nbs/00_attr.ipynb 18
#| export


# %% ../nbs/00_attr.ipynb 20
#| export


# %% ../nbs/00_attr.ipynb 21
def getdata(adata: AnnData, layer: StrQ = None) -> NPArray:
    return adata.layers.get(layer, adata.X)

# %% ../nbs/00_attr.ipynb 23
def getcol(obj, key: str) -> Series:
    '''Get a categorical column from a dataframe'''
    col = obj.get(key) if isdf(obj) else getkey(obj, key)
    col = pd.Series(col) if not isseries(col) else col
    return col

def getcatcol(obj, key: str) -> Series:
    '''Get a categorical column from a dataframe'''
    return getcol(obj, key).astype(CATEGORY)

def getcategory(obj, key: str, cat: str) -> str:
    opt = (col := getcatcol(obj, key)).cat.categories
    loc = opt.astype(str).get_loc(str(cat))
    val = opt[loc]
    return val

def getcategoryidx(obj, key: str, cat: str) -> Series:
    col = getcatcol(obj, key)
    val = getcategory(obj, key, cat)
    return (col == val)

def getcatval(obj, key: str, cat: str) -> str:
    return getcategory(obj, key, cat)

def getcatidx(obj, key: str, cat: str) -> Series:
    return getcategoryidx(obj, key, cat)

# %% ../nbs/00_attr.ipynb 26
def gene2idx(
    adata: AnnData, gene: str, gcol: str = GENE_SYMBOL
) -> IndexQ:    
    genes = getkey(adata, gcol).str.casefold()
    found = genes.where(genes == gene.casefold()).dropna()
    if found.size == 0: return None
    return found.index    

def gene2val(
    adata: AnnData, gene: str, gcol: str = GENE_SYMBOL, vcol: StrQ = None
) -> StrQ:
    gidxs = gene2idx(adata, gene, gcol = gcol)
    if gidxs is None: return None
    first = gidxs[0]
    
    value = getkey(adata, (vcol or gcol)).loc[first]
    return value

def genes2idxs(
    adata: AnnData, genes: Strs, gcol: str = GENE_SYMBOL
) -> IndexQ:
    return list(map(lambda g: gene2idx(adata, g, gcol), genes))

def genes2vals(
    adata: AnnData, genes: Strs, gcol: str = GENE_SYMBOL, vcol: StrQ = None
) -> list:
    return list(map(lambda g: gene2val(adata, g, gcol, vcol), genes))

def findgenes(
    adata: AnnData, genes: Strs, gcol: str = GENE_SYMBOL, vcol: StrQ = None
) -> list:
    return list(filter(notnone, genes2vals(adata, genes, gcol=gcol, vcol=vcol)))

def gene2eid(
    adata: AnnData, gene: str, gcol: str = GENE_SYMBOL, ecol: StrQ = ENSEMBL_ID
) -> StrQ:
    return gene2val(adata, gene, col=ecol, gcol=gcol)

def genes2str(
    genes: Strs, 
    suffix: StrQ = None, prefix: StrQ = None, sep: str = U1, 
    modifier: CallQ = str.title, 
    modfix: bool = False
):
    gstr = sep.join(filter(notnone, genes))
    gstr = modifier(gstr) if notnone(modifier) else gstr
    gstr = sep.join(filter(notnone, (prefix, gstr, suffix)))
    gstr = modifier(gstr) if notnone(modifier) and modfix else gstr
    return gstr

# %% ../nbs/00_attr.ipynb 28
def clusteridx(adata, key: str, col: str = CLUSTER):
    return getcategoryidx(adata.obs, col, key)

def conditionidx(adata, key: str, col: str = CON):
    return getcategoryidx(adata.obs, col, key)

# %% ../nbs/00_attr.ipynb 29
def subset_genes(adata: AnnData, genes: Strs, layer: StrQ = None, gcol: str = GENE_SYMBOL) -> DataFrame:
    genes = findgenes(adata, genes, gcol)
    # getkey(adata, gcol).isin(genes)
    # getcol(adata, gcol).isin(genes)
    idx = adata.var[gcol].isin(genes)
    data = getdata(adata, layer)
    df = pd.DataFrame(asarr(data[:, idx]), index=adata.obs.index, columns=adata.var.index[idx])    
    return df

# %% ../nbs/00_attr.ipynb 31
def axeschars(n: int, start: str = 'x') -> str:
    return shiftchars(start, string.ascii_lowercase)[:n]

def labelaxis(s: str, dim: int | str = None, modify = str.upper, sep = U1) -> int:
    s = str(s).removeprefix(X_)
    if notnone(dim): s = sep.join((s, str(dim)))
    if callable(modify): s = modify(s)
    return modify(s)

def labelaxes(
    key: str, ndim: int = 0, oneidx: bool = False, useint: bool = True, 
    modify = str.upper, sep = U1, data: AnnData = None, 
) -> tuple: 
    if isnone(ndim) or notnone(data): ndim = numcols(getkey(data, key))
    dims = range(int(oneidx), ndim + int(oneidx))
    if not useint: dims = axeschars(ndim, 'x')
    return tuple(labelaxis(key, dim, modify, sep) for dim in dims)

# %% ../nbs/00_attr.ipynb 33
def getemb(adata, key, oneidx: bool = False, useint: bool = True):
    df_emb = pd.DataFrame(getkey(adata, key), columns=labelaxes(key, oneidx, useint, data=adata))
    nobs = df_emb.shape[0]
    if nobs == adata.obs.shape[0]: df_emb.index = adata.obs.index
    if nobs == adata.var.shape[0]: df_emb.index = adata.var.index
    return df_emb
