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

# %% auto 0
__all__ = ['getpubattrs', 'iscallattr', 'namekeys', 'hasget', 'haskeys', 'hasnamekeys', 'iskeyed', 'getattrkeys', 'keyedattrs',
           'getkeys', 'getkey', 'getcol', 'getcatcol', 'getcategory', 'getcategoryidx', 'getcatval', 'getcatidx',
           'gene2idx', 'gene2val', 'genes2idxs', 'genes2vals', 'findgenes', 'gene2eid', 'genes2str', 'getdata',
           'clusteridx', 'conditionidx', 'subset_genes', 'numcols', 'keydims', 'axisname', 'axeschars', 'axesnames',
           '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, IterStr, IndexQ, Strs, NPArray, DataFrame
from chck import isad, isdf, notnone, ispublic, isseries

from nchr import U1
from nlit import GET, KEYS, SHAPE, COLUMNS, X_
from nlit import (GENE_SYMBOL, ENSEMBL_ID, CATEGORY, CLUSTER, CON)
from pstr import shiftchars
from asto import asarr
from calr import callattr

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


# %% ../nbs/00_attr.ipynb 18
def getpubattrs(obj) -> list: 
    return list(filter(ispublic, dir(obj)))

# %% ../nbs/00_attr.ipynb 20
def iscallattr(obj, attr: str) -> bool:
    return callable(getattr(obj, attr, None))

def namekeys(name: str) -> str: return f'{name}_{KEYS}'

def hasget(obj) -> bool: return iscallattr(obj, GET)

def haskeys(obj) -> bool: return iscallattr(obj, KEYS)

def hasnamekeys(obj, name: str) -> bool: return iscallattr(obj, namekeys(name))

def iskeyed(obj, name):    
    return haskeys(attr := getattr(obj, name, None)) or hasnamekeys(obj, name)

def getattrkeys(obj, name: str) -> list:
    attr = getattr(obj, name, None)    
    if haskeys(attr): return list(callattr(attr, KEYS))
    if hasnamekeys(obj, name): return list(callattr(obj, namekeys(attr)))
    return []

def keyedattrs(obj, **kwargs) -> list:
    keyed = [name for name in getpubattrs(obj, **kwargs) if iskeyed(obj, name)]
    return keyed
    
def getkeys(adata, locs = None) -> List[Tuple[str, str]]:
    with redirect_stderr(f_err := io.StringIO()):
        attrs = locs or getpubattrs(adata)
        keyfn = lambda a: zip(getattrkeys(adata, a), itertools.repeat(a))
        akeys = tuple(itertools.chain(*map(keyfn, attrs)))
    return akeys

# %% ../nbs/00_attr.ipynb 22
def getkey(
    obj, name: str, locs: list = None,
    obj1st: bool = False, useobj: bool = True, 
):
    res, errs = None, []

    if useobj & obj1st & hasget(obj):
        if notnone(res := obj.get(name)): return res
        errs.append(err := KeyError(f'Key {name} not found with `obj.get()`.'))
        
    
    known = getkeys(obj, locs)
    for (key, loc) in known:
        if key == name: break

    if key != name:
        errs.append(err := KeyError(f'Key {name} not found in obj attrs.'))
        raise err
    
    loc = getattr(obj, loc)
    if hasget(loc) and notnone(res := loc.get(name)): return res

    try: 
        if notnone(res := loc[name]): return res        
    except KeyError as err:
        errs.append(err := KeyError(f'Key {name} not found in attribute {loc}.'))        
    
    try:
        if useobj & (not obj1st) & hasget(obj):
            if notnone(res := obj.get(name)): return res

    except KeyError as err:
        errs.append(err)
        raise KeyError(f'Key {name} not found in attribute {loc}.')


# %% ../nbs/00_attr.ipynb 24
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 27
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 29
def getdata(adata: AnnData, layer: StrQ = None) -> NPArray:
    return adata.layers.get(layer, adata.X)

# %% ../nbs/00_attr.ipynb 30
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 31
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 34
def numcols(obj) -> int: 
    if hasattr(obj, SHAPE): return obj.shape[1]
    if hasattr(obj, COLUMNS): return len(obj.columns)
    return len(obj)

def keydims(adata, name: str) -> int: 
    return numcols(getkey(adata, name))

# %% ../nbs/00_attr.ipynb 36
def axisname(key: str, dim: int = None, modify = str.upper) -> int:
    key = str(key).removeprefix(X_)
    if notnone(dim): key = U1.join((key, str(dim)))
    return modify(key)

def axeschars(n: int, start: str = 'x') -> str:
    return shiftchars(start, string.ascii_lowercase)[:n]

def axesnames(adata, key: str, oneidx: bool = False, useint: bool = True): 
    dims = numcols(getkey(adata, key))
    adims = range(int(oneidx), dims + int(oneidx))
    if not useint: adims = axeschars('x')[:dims]
    return tuple(axisname(key, dim) for dim in adims)

# %% ../nbs/00_attr.ipynb 38
def getemb(adata, key, oneidx: bool = False, useint: bool = True):
    df_emb = pd.DataFrame(getkey(adata, key), columns=axesnames(adata, key, oneidx, useint))
    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
