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

# %% auto 0
__all__ = ['gene2idx', 'gene2val', 'genes2idxs', 'genes2vals', 'findgenes', 'gene2eid', 'getdata', 'subset_genes']

# %% ../nbs/00_util.ipynb 6
import warnings

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

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


# %% ../nbs/00_util.ipynb 12
#| export

# %% ../nbs/00_util.ipynb 14
from atyp import AnnData, StrQ, IterStr, IndexQ, Strs, NPArray, DataFrame

from psrc import dirdirs
from nlit import (GENE_SYMBOL, ENSEMBL_ID,)
from chck import isarr, islist, isany, isnone, isad, notnone, iscall, isstr
from xst.arr import asarr

# %% ../nbs/00_util.ipynb 16
def gene2idx(
    adata: AnnData, gene: str, gcol: str = GENE_SYMBOL
) -> IndexQ:
    genes = adata.var[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 = adata.var.loc[first, (vcol or gcol)]    
    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)

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

def subset_genes(adata: AnnData, genes: Strs, layer: StrQ = None, gcol: str = GENE_SYMBOL) -> DataFrame:
    genes = findgenes(adata, genes, gcol)
    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
