# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_data.ipynb.

# %% auto 0
__all__ = ['clipsd2', 'clipsd3', 'merge_dfs', 'calculate_contribution', 'calculate_contributions', 'BoundsSpec', 'ClipSpec',
           'FiltSpec', 'zscore_genes', 'zscore_genes_sd2', 'zscore_genes_sd3', 'score_genes', 'CutoffWarning',
           'zscore_cells', 'Basis', 'IRootSpec', 'hvgset', 'embed_hvg']

# %% ../nbs/04_data.ipynb 6
import warnings, string, itertools
from itertools import chain
from pathlib import Path

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

    from matplotlib.colors import LinearSegmentedColormap
    from matplotlib.patches import Patch

    from scipy.stats import zscore


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

# %% ../nbs/04_data.ipynb 8
from typing import Any, NamedTuple, Union, Optional, Tuple, Unpack, Self, Type

# %% ../nbs/04_data.ipynb 10
#| export


# %% ../nbs/04_data.ipynb 12
from atyp import AnnData, CallQ, NPArray, IntQ, BoolIndex, BoolQ
from chck import isad, isarr, islist, isany, isnone, iscall, isstr

from nchr import NIL, DOT
from nlit import GENE_SYMBOL, X_, X_PCA_HVG, PHATE_HVG, X_PHATE_HVG, EXT_H5AD, HIGHLY_VARIABLE

from asto import asarr

# %% ../nbs/04_data.ipynb 14
from .attr import getkey
from .attr import findgenes, subset_genes, genes2str, labelaxes, labelaxis
from .plot import emphplot

# %% ../nbs/04_data.ipynb 17
def merge_dfs(df_dict, on, how='outer', **kwargs):
    df_fin = pd.DataFrame()
    for df_val in df_dict.values():
        if df_fin.empty: 
            df_fin = df_val
            continue
        df_fin = pd.merge(df_fin, df_val, on=on, how=how, **kwargs)
    return df_fin

# %% ../nbs/04_data.ipynb 18
def calculate_contribution(
    df, category_key: str, cluster_key = 'cluster', total_key = 'total_cells'):    
    count_key = f'count_{category_key}'    
    percentage_key = f'percent_{category_key}'
    df_tmp = df.groupby([cluster_key, category_key]).size().reset_index(name=count_key)
    cnt_df = df_tmp.groupby(cluster_key)[count_key]
    counts = cnt_df.sum().reset_index(name=total_key)
    df_tmp = df_tmp.merge(counts, on=cluster_key, how='inner')    
    df_tmp.loc[:, percentage_key] = 100 * df_tmp[count_key] / df_tmp[total_key]
    # df_tmp.loc[:, percentage_key] = df_tmp.apply(lambda x: 100 * x[count_key] / x[total_key], 1)
    return df_tmp

# %% ../nbs/04_data.ipynb 19
def calculate_contributions(
    adata,
    condition_key = 'con',
    replicate_key = 'rep',
    timepoint_key = 'day',
    cluster_key = 'cluster',
    sample_key = 'sample',
    index_key = 'barcodes',
    total_key = 'total_cells',
    return_cats = True,
):
    
    column_keys = [cluster_key, condition_key, replicate_key, timepoint_key]
    contrib_keys = [condition_key, replicate_key, timepoint_key, sample_key]

    df = adata.obs.loc[:, column_keys].copy().reset_index(allow_duplicates=False)
    df.columns = [index_key, *column_keys]
    df.loc[:, sample_key] = df.apply(lambda x : f'{x[condition_key]}_r{x[replicate_key]}_t{x[timepoint_key]}', 1, result_type='expand')

    # Initialize an empty DataFrame for results
    df_res = pd.DataFrame()
    df_cats = {}
    for category_key in contrib_keys:
        df_tmp = calculate_contribution(df, category_key, cluster_key, total_key)
        df_cats[category_key] = (df_tmp)

        # Add to result DataFrame
        if df_res.empty: df_res = df_tmp
        else: df_res = pd.merge(df_res, df_tmp, on=[cluster_key, total_key], how='outer')

    if return_cats: return df_res, df_cats
    return df_res

# %% ../nbs/04_data.ipynb 21
class BoundsSpec(NamedTuple):
    '''Boolean Specification class'''
    lower: float = float('-inf')
    upper: float = float('inf')  
    name:  str = None   
    desc:  str = None
    replace: dict = {}
    replace_kws: dict = dict(inplace=False)

    def __call__(self, *args, **kwargs): return self.apply(*args, **kwargs)
    def __iter__(self):  return iter((self.lower, self.upper))
    def apply(self, data, key = None, *args, **kwargs):
        match type(data):
            case ad.AnnData:
                try: arr = getkey(data, key)
                except ValueError as err:
                    if kwargs.get('copy', True) == False:
                        raise err
                    arr = getkey(data.copy(), key)

            case pd.DataFrame: arr = getkey(data, key)
            case np.ndarray:   arr = pd.Series(data)
            case _: arr = pd.Series(data)

        res = (self.lower < arr) & (arr < self.upper)
        srs = pd.Series(res, index=arr.index, name=key)

        rep = kwargs.get('replace', self.replace)
        kws = kwargs.get('replace_kws', self.replace_kws)
        inp = kws.get('inplace', False)
        if isad(data) and inp:
            srs.replace(rep, **kws)
        else:
            if kwargs.get('copy', True): 
                srs = srs.copy()
            srs.replace(rep, **kws)
        return srs
    
    
class ClipSpec(NamedTuple):
    '''Clipping Specification class'''
    lower: float = float('-inf')    
    upper: float = float('inf')    
    vmin:  float = None  
    vmax:  float = None  
    name:  str = None   
    desc:  str = None

    def __call__(self, *args, **kwargs): return self.apply(*args, **kwargs)
    def __iter__(self): return iter((self.lower, self.upper, self.vmin or self.lower, self.vmax or self.upper))

    def apply(self, data, *args, **kwargs):
        if isarr(data) or islist(data): data = pd.Series(data)
        data[data < self.lower] = self.vmin if isany(self.vmin) else self.lower
        data[data > self.upper] = self.vmax if isany(self.vmax) else self.upper
        return data
    

class FiltSpec(NamedTuple):
    key: str
    lower: float = float('-inf')
    upper: float = float('inf') 

    def apply(self, data, verb: bool = False, inplace: bool = False, copy: bool = True):
        bidx = BoundsSpec(self.lower, self.upper)(data, key = self.key)
        if isad(data) and inplace:
            data._inplace_subset_obs(bidx)
        else:
            data = data[bidx]
            if copy: data = data.copy()
        if isad(data) and verb:
            print(data.obs[self.key].value_counts())
        return data

# %% ../nbs/04_data.ipynb 22
# Alternative score - the number of markers above 2 SD in smoothed expression
clipsd2 = ClipSpec(
    lower=2, upper=2, vmin=0, vmax=1, name='sd2', 
    desc='Alternative score - the number of markers above 2 SD in smoothed expression'
)

# Remove influence of markers less than 1 SD above the mean and stop individual markers dominating
clipsd3 = ClipSpec(
    lower=1, upper=3, vmin=0, vmax=3, name='sd3', 
    desc='Remove influence of markers less than 1 SD above the mean and stop individual markers dominating'
)

# %% ../nbs/04_data.ipynb 26
def zscore_genes(
    adata, genes, layer = None, genecol = GENE_SYMBOL,
    clip = False, lower = 1, upper = 3, sdmin = 0, sdmax = 3
):
    genes = findgenes(adata, genes)
    # Create dataframe of smoothed MAGIC counts for current markers
    df = subset_genes(adata, genes, layer, genecol)
    # Scale expression per gene
    df = df.apply(zscore)   
    if clip: df = ClipSpec(lower, upper, sdmin, sdmax)(df)
    return df

def zscore_genes_sd2(adata, genes, layer = None, genecol = GENE_SYMBOL):
    return clipsd2(zscore_genes(adata, genes, layer, genecol, clip=False))

def zscore_genes_sd3(adata, genes, layer = None, genecol = GENE_SYMBOL): 
    return clipsd3(zscore_genes(adata, genes, layer, genecol, clip=False))


def score_genes(
    adata, genes, layer = None, genecol = GENE_SYMBOL,
    clipby = clipsd3, obskey = None
):
    df = zscore_genes(adata, genes, layer, genecol, clip=False)    
    
    # Clip by Standard Deviation
    if clipby is None: clipby = clipsd3
    elif isstr(clipby):
        if clipby == 'default' or clipby.endswith('sd3'): clipby = clipsd3
        elif 'alter' in clipby or clipby.endswith('sd2'): clipby = clipsd2
        elif clipby.startswith('clip_'): clipby = globals().get(clipby, None)
        
    if iscall(clipby): df = clipby(df)

    # Sum up scores
    scores = df.sum(axis=1)

    # Store scores in adata.obs
    obskey = obskey or genes2str(genes, 'score', modfix=False)
    adata.obs[obskey] = scores
    return scores

class CutoffWarning(Warning):    
    @classmethod
    def msg(cls, cutoff, scores): return f'Cutoff {cutoff} greater than max score {np.max(scores)}'

def zscore_cells(
    adata, genes, layer = None, genecol = GENE_SYMBOL,
    clipby = clipsd3, cutoff = 3, plot = False,
    obskey = None, subkey = None, zsckey = None,
):
    zsckey = obskey or genes2str(genes, 'zscore', modfix=False)
    subkey = obskey or genes2str(genes, 'subset', modfix=False)
    obskey = obskey or genes2str(genes, modfix=False)
    if zsckey == obskey: zsckey = f'{obskey}_zscore'
    if subkey == obskey: zsckey = f'{obskey}_subset'

    scores = score_genes(adata, genes, layer, genecol, clipby, obskey=zsckey)
    # Histogram of tissue scores to find min threshold across all cells
    adata.obs[subkey] = (scores >= cutoff).replace({True: obskey, False: pd.NA,}) 
    adata.obs[obskey] = adata.obs[subkey] == obskey
    if plot: scprep.plot.histogram(scores, cutoff=cutoff, title=obskey)
    # if cutoff > scores.max(): warnings.warn(CutoffWarning.msg(cutoff, scores), CutoffWarning)
    return adata

# %% ../nbs/04_data.ipynb 28
class Basis:
    def __init__(self, name: str = PHATE_HVG, prepend_X: BoolQ = None) -> None:
        if isnone(prepend_X): prepend_X = name[:2] != X_         
        self.name = name
        self.prex = prepend_X

    @property
    def basis(self) -> str: return (X_ if self.prex else NIL) + self.name
    @property
    def label(self) -> str: return self.basis.removeprefix(X_)

    def __str__(self) -> str: return self.basis
    def __repr__(self) -> str: return f'{type(self).__name__}({self.label})'

    
    def __call__(self, adata: AnnData, *args, **kwargs) -> NPArray:
        kwargs.setdefault('locs', ['obsm'])
        return getkey(adata, self.basis, *args, **kwargs)
    
    def axis(self, dim: int = None, modify = str.upper, spaces: bool = True) -> str:        
        return labelaxis(self.label, dim, modify, spaces)
   
    def axes(self, dims: int, oneidx: bool = False, useint: bool = True) -> list[str]:
        return labelaxes(self.label, dims, oneidx, useint)

    def xyz(self, *dims) -> tuple[int, int, int]:
        '''Extracts x, y, z dims from *dims defaulting to None if missing'''
        x, y, z = (0, 1, 2)
        for i in range(len(dims)):
            match i:
                case 0: x = dims[i]
                case 1: y = dims[i]
                case 2: z = dims[i]
                case _: break
        return (x, y, z)
    
    def dim(self, adata: AnnData, dim: int) -> NPArray: return self(adata)[:, dim]
    def dims(self, adata: AnnData, *dims) -> tuple: return tuple(self.dim(adata, dim) for dim in dims)
    def shape(self, adata: AnnData) -> int: return self(adata).shape    
    def nrows(self, adata: AnnData) -> int: return self.shape(adata)[0]
    def ncols(self, adata: AnnData) -> int: return self.shape(adata)[1]

# %% ../nbs/04_data.ipynb 30
class IRootSpec:    
    def __init__(self, dim: int = 0, val: str | CallQ = 'min'):
        self.dim = dim
        self.get = getattr(np, val, val if iscall(val) else lambda x: val)

    def __call__(self, adata: AnnData, basis: Basis, dim: IntQ = None) -> BoolIndex:
        return self.find(adata, basis, dim)

    def find(self, adata: AnnData, basis: str | Basis, dim: IntQ = None) -> BoolIndex:
        '''Find the indicies matchingthe root cell in the specified basis'''
        basis = Basis(str(basis))
        return (emb := basis.dim(adata, dim or self.dim)) == self.get(emb)

    def dims(self, adata: AnnData, basis: str | Basis, *dims) -> tuple[BoolIndex, ...]:
        '''Apply the instance to multiple dims'''
        return tuple(self(adata, basis, dim) for dim in dims)
    
    @classmethod
    def map(cls: Type[Self], adata: AnnData, basis: str | Basis, *specs: Tuple[Type[Self], ...]) -> tuple[BoolIndex, ...]:
        '''Map multiple specs to the same basis'''
        return tuple(spec(adata, basis) for spec in specs)
    
    @classmethod
    def iroot(
        cls: Type[Self], adata: AnnData, basis: str | Basis = Basis('phate'), 
        specs: Optional[Tuple[Type[Self], ...]] = None
    ) -> BoolIndex:
        if isnone(specs): specs = (cls(0, 'min'), )        
        found = cls.map(adata, basis, *specs)
        bools = np.all(found, axis=0)
        iroot = np.where(bools)[0][0]
        return iroot
    
    @classmethod
    def select(
        cls: Type[Self], adata: AnnData, basis: str | Basis = Basis('phate'), 
        specs: Optional[Tuple[Type[Self], ...]] = None, plot: bool = False, **kwargs
    ):
        basis = Basis(str(basis))
        if isnone(specs): specs = (cls(0, 'min'), )
        iroot = cls.iroot(adata, basis, specs)
        if not plot: return iroot

        dims = kwargs.pop('dims', dict(x=0, y=1, z=2))
        for i, c in enumerate('xyz'): dims.setdefault(c, i)
        if basis.ncols(adata) < 3 or kwargs.pop('is2d', False): dims.pop('z', None)
        fig, ax = emphplot(adata, *(iroot, ), basis=str(basis), **dims, **kwargs)
        return iroot

# %% ../nbs/04_data.ipynb 35
def hvgset(adata):
    return adata[:, adata.var[HIGHLY_VARIABLE]]

def embed_hvg(
    subset, 
    pca_kws = dict(), g_kws = dict(), phate_kws = dict(),
    save: bool = True, filename: str = None, savepath: Path = None,
    varname: str = None, toglobals: bool = False, 
):
    subhvg = hvgset(subset)    
    pcs, svs = scprep.reduce.pca(asarr(subhvg), **pca_kws)    
    subset.obsm[X_PCA_HVG] = pcs
    g_hvg = gt.Graph(pcs, **g_kws)
    
    pht_op = phate.PHATE(**phate_kws)
    embpht = pht_op.fit_transform(g_hvg)
    subset.obsm[X_PHATE_HVG] = embpht

    if varname is None and not save: return subset
    if filename is None: filename = f'{varname}{EXT_H5AD}'
    if toglobals: globals()[varname] = subset
    if save and (file := ((savepath or Path(DOT)) / filename) and not file.is_file()): 
        subset.write_h5ad(file)
