"""utilities module for `anndash`"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_utils.ipynb.

# %% auto 0
__all__ = ['fullpath', 'default_env', 'env_dependencies', 'ispublic', 'isproperty', 're_imatch', 'str_match', 'str_imatch',
           'map_match', 'map_imatch', 'filter_matches', 'filter_imatches', 'first_match', 'first_imatch',
           'str_cased_eq', 'col_cased_eq', 'check_columns', 'subset_column_names', 'subset_columns', 'var_index',
           'get_basis_keys', 'get_gene_info', 'get_ndim', 'configure_anndash_keys', 'safe_colors', 'flatten_obs_array',
           'get_basis_dataframe', 'make_bases_dataframe', 'make_embed_dataframe', 'make_color_dataframe']

# %% ../nbs/02_utils.ipynb 4
import os, re, io, yaml
from enum import StrEnum, auto
from pathlib import Path
from inspect import isfunction, ismethod
from functools import reduce
from importlib_resources import files

try: 
    from typing import Any, Sequence, TypeAlias
except ImportError:
    from typing_extensions import Any, Sequence, TypeAlias
        
import numpy as np, pandas as pd, anndata as ad

try: from sklearn.preprocessing import MinMaxScaler
except ImportError: ...

# %% ../nbs/02_utils.ipynb 5
from anndash.hints import (
    Strings, DataLike, Pattern, PatternQ,
    Flags, MatchQ, Columns, REGEX_BASIS
)

# %% ../nbs/02_utils.ipynb 7
def fullpath(path: str | Path) -> Path:
    '''Resolves user and symlinks.'''
    return Path(path).expanduser().absolute()

# %% ../nbs/02_utils.ipynb 9
def default_env(stem: str = 'env', resource_path: str = '.') -> Path:
    '''Try and return the path a template `env` file provided by this package.'''
    try:
        import anndash
    except ImportError as err:
        err.add_note('Error raised from `default_env` function call attempting to load the default environment. Do you have `anndash` installed?')
        raise err
    path = files(anndash).joinpath(resource_path, stem).with_suffix('.yml')
    return path

def env_dependencies(file: str | Path | None = None) -> list[str]:
    '''Try and return the list of dependencies used within the yaml file.'''
    env_file = file or default_env()
    env_data = yaml.safe_load(io.open(env_file))
    return env_data.get('dependencies', [])

# %% ../nbs/02_utils.ipynb 11
def ispublic(member):
    name = getattr(member, '__name__', '_')
    return name.startswith('_')

def isproperty(member):    
    return (not ismethod(member) and not isfunction(member)) and ispublic(member)

# %% ../nbs/02_utils.ipynb 13
def re_imatch(pattern: Pattern, string: str) -> MatchQ:
    '''Invoke`re.match` with case-insensitive flag i.e. `flags=re.RegexFlag.IGNORECASE`.'''
    return re.match(pattern, string, re.RegexFlag.IGNORECASE)

def str_match(string: str, pattern: PatternQ = None, flags: Flags = re.RegexFlag.NOFLAG) -> str | None:
    '''Returns `string` if `re.match` passes else `None`.'''
    if re.match(pattern or r'*.', string, flags): return string

def str_imatch(string: str, pattern: PatternQ = None) -> str | None:
    '''Returns `string` if `re.match` with case-insensitive flag i.e.`flags=re.RegexFlag.IGNORECASE` passes else `None`.'''
    return str_match(string, pattern, re.RegexFlag.IGNORECASE)

def map_match(strings: Strings, pattern: PatternQ = None, flags: Flags = re.RegexFlag.NOFLAG) -> list[str]:
    '''Maps `re.match(pattern, string,flags=flags)` across `strings`.'''
    match_with_flags = lambda s: str_match(s, pattern, flags)
    return map(match_with_flags, strings)

def map_imatch(strings: Strings, pattern: PatternQ = None) -> list[str]:
    '''Maps `re.match(pattern, string,flags=flags)` across `strings` with case-insensitive flag i.e.`flags=re.RegexFlag.IGNORECASE`.'''
    return map_match(strings, pattern, re.RegexFlag.IGNORECASE)

def filter_matches(strings: Strings, pattern: PatternQ = None, flags: Flags = re.RegexFlag.NOFLAG) -> list[str]:
    '''Filter mapped matches across strings to just those that passes.'''
    return list(filter(None, map_match(strings, pattern, flags)))

def filter_imatches(strings: Strings, pattern: PatternQ = None) -> list[str]:
    '''Filter mapped matches across strings to just those that passes with case-insensitive flag i.e.`flags=re.RegexFlag.IGNORECASE`.'''
    return filter_matches(strings, pattern, re.RegexFlag.IGNORECASE)

def first_match(strings: Strings, pattern: PatternQ = None, flags: Flags = re.RegexFlag.NOFLAG) -> str | None:
    '''Returns the first string to match with `pattern` otherwise `None`.'''
    return next(iter(filter_matches(strings, pattern, flags)), None)

def first_imatch(strings: Strings, pattern: PatternQ = None) -> str | None:
    '''Returns the first string to match with `pattern` with case-insensitive flag i.e.`flags=re.RegexFlag.IGNORECASE` otherwise `None`.'''
    return first_match(strings, pattern, re.RegexFlag.IGNORECASE)

# %% ../nbs/02_utils.ipynb 14
def str_cased_eq(s: str, t: str, i: bool = True):
    if i: return str(s).casefold() == str(t).casefold()
    return str(s) == str(t)

def col_cased_eq(c: pd.Series, s: str, i: bool = True):
    if i and (c.dtype in {'object', 'category'} or hasattr(c, 'str')):
        return c.str.casefold() == str(s).casefold()
    return c == s

# %% ../nbs/02_utils.ipynb 16
def check_columns(
    df: pd.DataFrame, 
    value: Any, 
    subset: Columns | None = None, 
    insensitive: bool = True
) -> pd.Series | None:
    '''
    Search dataframe `df`'s columns (or subset thereof) for a column 
    named `value` case `insensitive` as specified by the user.
    '''
    cols = df.columns
    if subset is not None:
        cols = [col for col in df.columns if col in subset]
    
    flags = re.RegexFlag.IGNORECASE if insensitive else re.RegexFlag.NOFLAG
    first = first_match(cols, str(value), flags)
    if first:
        return df[first]
    return None

def subset_column_names(
    df: pd.DataFrame, 
    subset: Columns | None = None
) -> list[str]:
    if subset is None: return list(df.columns)
    return [str_imatch(col, sub) for sub in subset for col in df.columns]

def subset_columns(
    df: pd.DataFrame, 
    subset: Columns | None = None
) -> pd.DataFrame:
    if subset is None: return df
    subset = subset_column_names(df, subset)
    return df[subset]


# %% ../nbs/02_utils.ipynb 18
def var_index(
    adata: ad.AnnData, col: str, val: str | None
):
    var = adata.var[col]
    ixs = var[(adata.var.index == val) | (var == val)]
    idx = ixs.first_valid_index()
    return idx

# %% ../nbs/02_utils.ipynb 19
def get_basis_keys(adata: ad.AnnData, pattern: Pattern = REGEX_BASIS):
    return filter_matches(pattern, list(adata.obsm.keys()))

# %% ../nbs/02_utils.ipynb 20
# def get_gene_info(adata: ad.AnnData, gene: str, gene_col: str = 'gene_symbol') -> tuple[int, str | int, str]:
#     mask = adata.var.apply(lambda col: col == gene)
#     gidx = mask.any(axis=1)
#     gloc = adata.var[mask]
#     info = gloc[gene_col]
#     ensb = info.index[0]
#     symb = info.values[0]
#     gidx = adata.var.index.get_loc(ensb)
#     return gidx, ensb, symb

def get_gene_info(adata: ad.AnnData, gene: str, gene_col: str | None = 'gene_symbol') -> tuple[int, str | int, str]:
    # the anndata.var dataframe
    mask = adata.var
    if gene_col in mask: 
        # subset to just the `gene_col` containing the entries to search
        mask = mask[gene_col]
    
    # any entry in the `mask` dataframe that matches `gene` as safety fallback
    mask = mask.apply(lambda col: col == gene)
    
    # boolean of rows where any column matches `gene`
    gene_bool = mask.any(axis=1)
    
    # just the rows matching `gene`
    gene_mask = adata.var[gene_bool]
    
    # info about the gene
    gene_info = gene_mask
    if gene_col in gene_info:
        # just the column with the gene name
        gene_info = gene_info[gene_col] 
    
    # the index used in the dataframe, needed for extracting from `adata[:, gene_index]` or `adata.X[:, gene_index]`
    gene_index = mask.index.get_loc(ensembl_id)
    # assumes that the index is set ot ensembl_id
    ensembl_id = gene_info.index[0]
    # `gene_col` is the first value
    gene_symbol = gene_info.values[0]
    # note that `gene_index` will be an `int` but `ensembl_id` can be either `int` or `str`
    return gene_index, ensembl_id, gene_symbol

# %% ../nbs/02_utils.ipynb 21
def get_ndim(adata: ad.AnnData, basis: str | None = None, max_dim: int | None = 10):
    ndim = 2
    if basis in adata.obsm: 
        ndim = adata.obsm[basis].shape[1]
    if max_dim:
        ndim = min(ndim, max_dim)
    return ndim

# %% ../nbs/02_utils.ipynb 22
def configure_anndash_keys(
    adata: ad.AnnData,
    
    include_obs: set[str] | None = None,
    include_var: set[str] | None = None,
    include_obsm: set[str] | None = None,
    include_layers: set[str] | None = None,
    
    exclude_obs: set[str] | None = None,
    exclude_var: set[str] | None = None,
    exclude_obsm: set[str] | None = None,
    exclude_layers: set[str] | None = None,
) -> tuple[list[str], list[str], list[str], list[str]]:
    obs_keys = adata.obs_keys()
    var_keys = adata.var_keys()
    obsm_keys = adata.obsm_keys()
    layers_keys = list(adata.layers.keys())
    
    
    # keys to use for obs, var, obsm and layers
    # ensure keys are a set, or use all keys
    include_obs = set(include_obs or obs_keys)
    include_var = set(include_var or var_keys)
    include_obsm = set(include_obsm or obsm_keys)
    include_layers = set(include_layers or layers_keys)

    # only keys in intersection
    include_obs = (include_obs & obs_keys)
    include_var = (include_var & var_keys)
    include_obsm = (include_obsm & obsm_keys)
    include_layers = (include_layers & layers_keys)


    # keys to ignore for obs, var, obsm and layers    
    exclude_obs = set(exclude_obs or set())
    exclude_var = set(exclude_var or set())
    exclude_obsm = set(exclude_obsm or set())
    exclude_layers = set(exclude_layers or set())
    
    # only keys in difference
    obs_keys = list(include_obs - exclude_obs)
    var_keys = list(include_var - exclude_var)
    emb_keys = list(include_obsm - exclude_obsm)
    lay_keys = list(include_layers - exclude_layers)
    
    return obs_keys, var_keys, emb_keys, lay_keys

# %% ../nbs/02_utils.ipynb 24
def safe_colors(adata: ad.AnnData, colors: DataLike | None = None) -> DataLike:
    shape = getattr(colors, 'shape', ())
    n_obs = shape[0] if len(shape) else 0
    
    if colors is None or n_obs != getattr(adata, 'n_obs', 0):
        colors = np.zeros(adata.n_obs)
    
    if hasattr(colors, 'todense'): 
        colors = colors.todense()
            
    if not isinstance(colors, (np.ndarray, pd.Index, pd.Series)): 
        colors = np.asarray(colors)
        colors = colors.flatten()
        
    return colors

# %% ../nbs/02_utils.ipynb 25
def flatten_obs_array(
    array: np.ndarray | pd.Index | pd.Series | list | None = None,
    adata: ad.AnnData | None = None, 
    n_obs: int | None = 0
) -> np.ndarray | pd.Index | pd.Series | list:
    # shape of array e.g. (1000, 1)
    shape = getattr(array, 'shape', ())
    # number of elements in array
    n_els = shape[0] if len(shape) else 0
    # number of observation in anndata
    n_obs = getattr(adata, 'n_obs', n_obs)
    # default to zeros array if mismatch of elements and observations
    if not array or n_els != n_obs:
        array = np.zeros(n_obs)
    
    # if sparse matrix, make dense
    if hasattr(array, 'todense'): 
        array = array.todense()

    if not isinstance(array, (np.ndarray, pd.Index, pd.Series)): 
        array = np.asarray(array)
        array = array.flatten()
    
    return array

# %% ../nbs/02_utils.ipynb 26
def get_basis_dataframe(
    adata: ad.AnnData, 
    basis: str, 
    max_dims: int = 10,
    axis_name: str | None = None,
    include_basis: bool = False,
    remove_prefix: str | bool | None = True,
):
    if remove_prefix is None or remove_prefix is True: 
        remove_prefix = 'X_'
    
    name = basis.removeprefix(remove_prefix)
    if axis_name: 
        name = axis_name
    
    embed = adata.obsm[basis][:, :max_dims]
    n_obs, n_dim = embed.shape
    
    df_basis = pd.DataFrame(np.repeat(basis, n_obs), index=adata.obs.index, columns=['basis'])
    df_embed = pd.DataFrame(embed, index=adata.obs.index, columns=[f'{name}_{i}' for i in range(n_dim)])
    if include_basis: return df_basis.join(df_embed)
    return df_embed

def make_bases_dataframe(
    adata: ad.AnnData, 
    max_dims: int = 3,
):
    df_bases = [
        get_basis_dataframe(adata, basis, max_dims=max_dims) 
        for basis in adata.obsm_keys()
    ]
    return reduce(lambda left, right: left.join(right, how='right'), df_bases, pd.DataFrame())

def make_embed_dataframe(
    adata: ad.AnnData, 
    max_dims: int = 3,
    axis_name: str = 'ax',
    # include_basis: bool = True,
):
    df_bases = [
        get_basis_dataframe(adata, basis, max_dims=max_dims, axis_name=axis_name, include_basis=True) 
        for basis in adata.obsm_keys()
    ]
    return pd.concat(df_bases)
    
def make_color_dataframe(
    adata: ad.AnnData, 
    gene_col: str = 'gene_symbol'
):
    df_ann = pd.DataFrame(adata.X, columns=adata.var[gene_col], index=adata.obs.index)
    df_all = adata.obs.join(df_ann)
    return df_all
