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

except ImportError:
    warnings.warn("Some packages are not installed")

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

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


# %% ../nbs/00_attr.ipynb 12
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 14
#| export


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


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

# %% ../nbs/00_attr.ipynb 19
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 22
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 24
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 25
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 27
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 29
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
