"""Color module for `anndash`"""

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

# %% auto 0
__all__ = ['rescale_points', 'blank_colors', 'color_categories', 'color_dtype', 'guess_categorical', 'extract_colors',
           'obs_color', 'var_color', 'get_gene_colors', 'get_colors', 'get_obs_colors', 'get_var_colors',
           'color_series']

# %% ../nbs/04_color.ipynb 4
import os, re, io, yaml
from pathlib import Path
from inspect import isfunction, ismethod
from importlib_resources import files
from enum import StrEnum, auto
from warnings import warn
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/04_color.ipynb 5
from anndash.hints import (
    DataLike, Columns, AnnOption
)
from anndash.utils import (
    subset_columns, check_columns, var_index,
    get_gene_info, safe_colors, flatten_obs_array
)
from .enums import DTypePattern, ColorDType


# %% ../nbs/04_color.ipynb 7
def rescale_points(points: np.ndarray, scaler: MinMaxScaler | None = None):
    scaler = scaler or MinMaxScaler()
    scaler.fit(points)
    return scaler.transform(points)

def blank_colors(data: np.ndarray, colors: DataLike | None = None):
    return colors or np.zeros(data.shape[0])

def color_categories(colors: DataLike):
    color_map, categories = pd.factorize(colors, sort=True)
    return color_map, categories

# %% ../nbs/04_color.ipynb 9
def color_dtype(colors: DataLike) -> ColorDType:
    '''Given an sequence of data determine its datatype.'''
    # try to get dtype property of object
    dtype = getattr(colors, 'dtype', float)
    
    # see if data-type has a name otherwise ensure string
    dname = getattr(dtype, 'name', str(dtype)).lower()
    
    match dname:
        case dname if DTypePattern.OBJECT(dname): return ColorDType.OBJECT
        case dname if DTypePattern.STRING(dname): return ColorDType.STRING
        case dname if DTypePattern.NUMBER(dname): return ColorDType.NUMBER
        case dname if DTypePattern.BOOLEAN(dname): return ColorDType.BOOLEAN
        case dname if DTypePattern.CATEGORY(dname): return ColorDType.CATEGORY
        case dname if DTypePattern.DATETIME(dname): return ColorDType.DATETIME
        case _: return ColorDType.UNKNOWN

# %% ../nbs/04_color.ipynb 10
def guess_categorical(data: DataLike, threshold: float = 0.05, force: bool = False) -> bool:
    '''Given a sequence of data guess if it is categorical data based on a threshold (percentage) of
    unique values compared to total number of inputs, with the option to force the guess to be categorical.
    '''
    unique = np.unique(data)
    heuristic = (len(unique) / len(data) < threshold) or force
    dtype = color_dtype(data)
    is_cat = DTypePattern.CATEGORY(dtype) or heuristic
    return is_cat

# %% ../nbs/04_color.ipynb 13
def extract_colors(
    adata: ad.AnnData,
    value: Any,
    layer: str | None = None,
    annots: AnnOption | Sequence[AnnOption] | None = None,
    subset: Columns = None,
    insensitive: bool = True,
) -> DataLike | None:
    # the parts of the AnnData object to check
    parts = [annots] if isinstance(annots, str) else annots or ['obs', 'var', 'gene']
    
    # default color values
    colors = None
    
    # attempt to get the matrix of the specified layer, falling back to `anndata.X`
    mat = adata.layers.get(layer, adata.X)
    
    # Check adata.obs and adata.var
    for part in parts:
        # e.g. `anndata.obs`, `anndata.var`
        # NOTE: that means 'gene' doesn't get past here.
        df = getattr(adata, part, None)
        if df is None: continue
        
        # Check if value is a column
        col = check_columns(df, value, subset, insensitive)
        if col is not None and len(col) == adata.n_obs:
            # most likely occurs when value is apart of `anndata.obs` e.g. `anndata.obs['condition']`
            colors = col
            break
        
        # Check if value is an entry in any column of the provided subset.
        idx = check_values(df, value, subset, insensitive)
        if idx is None: continue
        # if we do not have an index, we can continue

        match part:
            case 'obs': 
                # if index is a string e.g. `'ensembl_id'` or `'gene_symbol'` get the integer position
                if isinstance(idx, str):
                    idx = adata.obs.index.get_loc(idx)
                
                # and then extract that row
                colors = mat[idx, :]
                break
            
            case 'var':
                # if index is a string e.g. `'ensembl_id'` or `'gene_symbol'` get the integer position
                if isinstance(idx, str):
                    idx = adata.var.index.get_loc(idx)
                
                # and then extract that column
                colors = mat[:, idx]
                break
            
            case 'gene': 
                # if index is a string e.g. `'ensembl_id'` or `'gene_symbol'` get the integer position
                if isinstance(idx, str):
                    idx = adata.var.index.get_loc(idx)
                
                # and then extract that column
                colors = mat[:, idx]
                break
            case _:
                ...
            
    # failed to find colors that match the number of observations, create zeros
    if colors is None or colors.shape[0] != adata.n_obs:
        colors = np.zeros(adata.n_obs)
    
    # colors is not a safe format for plotly, e.g. sparse array
    if not isinstance(colors, (np.ndarray, pd.Series)): 
        if hasattr(colors, 'todense'):
            colors = colors.todense()
        colors = np.asarray(colors)
        colors = colors.flatten()
        
    return colors

# %% ../nbs/04_color.ipynb 15
def obs_color(
    adata: ad.AnnData,
    col: str, 
    val: str | None = None,
    categorical: bool = False,
    boolean: bool = False
):
    obs = adata.obs[col]
    if categorical: 
        obs = obs.astype('category')
    if val and boolean:
        obs = obs.eq(val).rename({True: val, False: 'other'})
    return obs

def var_color(
    adata: ad.AnnData, col: str, val: str | None = None, 
    layer: str | None = None
):
    idx = var_index(adata, col, val)
    return adata[:, idx].layers.get(layer, adata[:, idx].X).flatten()

# %% ../nbs/04_color.ipynb 16
def get_gene_colors(adata: ad.AnnData, gene: str, layer: str | None = None, gene_col: str = 'gene_symbol'):
    '''DeprecationWarning'''
    warn('Use `anndash.color.extract_colors` instead', DeprecationWarning)
    gidx, _, _ = get_gene_info(adata, gene, gene_col)
    gcol = adata[:, gidx]
    colors = gcol.X
    if layer in adata.layers:
        colors = gcol.layers[layer]
    return safe_colors(adata, colors)

def get_colors(
    adata: ad.AnnData,
    obs: str | None = None,
    var: str | None = None,
    layer: str | None = None,
    gene_col: str | None = 'gene_symbol'
):
    warn('Use `anndash.color.extract_colors` instead', DeprecationWarning)
    colors = []
    if adata is None: return colors
        
    
    if obs is not None and obs in adata.obs:
        colors = adata.obs[obs]
        
    if var is not None and var in adata.var[gene_col]:
        colors = get_gene_colors(adata, var, layer, gene_col)
    
    return safe_colors(adata, colors)

# %% ../nbs/04_color.ipynb 17
def get_obs_colors(adata: ad.AnnData, obs: str | None = None):
    if obs in adata.obs: return adata.obs[obs]
    return None

def get_var_colors(adata: ad.AnnData, var: str | None = None):
    if var in adata.var: return adata.var[var]
    return None

def get_gene_colors(adata: ad.AnnData, gene: str, gene_col: str | None = 'gene_symbol', layer: str | None = None):
    gene_index, ensembl_id, gene_symbol = get_gene_info(adata, gene, gene_col)
    gene_adata = adata[:, gene_index]
    gene_colors = gene_adata.X
    if layer in adata.layers:
        gene_colors = gene_adata.layers[layer]
    return flatten_obs_array(gene_colors, adata, adata.n_obs)

# %% ../nbs/04_color.ipynb 18
def color_series(
    colors: np.ndarray | pd.Series, 
    target: Any | None = None,
    categorical: bool = False, 
    boolean: bool = False,
):
    warn('Use `anndash.color.extract_colors` instead', DeprecationWarning)
    if not isinstance(colors, pd.Series):
        colors = pd.Series(colors)
    
    ctype = color_dtype(colors)
    is_cat = guess_categorical(colors, force=categorical)
    
    if is_cat: 
        colors = colors.astype('category')
        
    if (target and is_cat) or (target and ctype == bool) or boolean:
        colors = colors.eq(target).rename({True: target, False: 'other'})
    return colors
