"""Core module for `anndash`"""

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

# %% auto 0
__all__ = ['DataLike', 'fullpath', 'default_env', 'env_dependencies', 'matched_source', 'map_matched_source', 'first_matched',
           'ColorDType', 'DTypePattern', 'color_dtype', 'guess_categorical', 'check_columns', 'check_values',
           'extract_colors', 'AnnDashboard']

# %% ../nbs/00_core.ipynb 4
import os, re, io, yaml
from importlib_resources import files

from typing import Any, Literal, LiteralString, Sequence, TypeAlias
from enum import StrEnum, auto
from pathlib import Path

# widgets
from IPython.display import display
import ipywidgets as widgets

# data analysis packages
import pandas as pd, numpy as np
from sklearn.preprocessing import MinMaxScaler

# single cell packages
import anndata as ad

# plotting packages
import matplotlib as mpl, matplotlib.pyplot as plt, seaborn as sns
import plotly.graph_objects as go, plotly.express as px

from dash import Dash, dcc, html, Input, Output, State
import dash.development.base_component as ddbc
import dash_bootstrap_components as dbc

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

# %% ../nbs/00_core.ipynb 7
def default_env(stem: str = 'env', resource_path: str = '.') -> Path:
    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]:
    env_file = file or default_env()
    env_data = yaml.safe_load(io.open(env_file))
    return env_data.get('dependencies', [])

# %% ../nbs/00_core.ipynb 10
def matched_source(pattern: re.Pattern, string: str, flags: re.RegexFlag = re.RegexFlag.NOFLAG):
    if re.match(pattern, string, flags): return string
    
def map_matched_source(pattern: re.Pattern, strings: Sequence[str], flags: re.RegexFlag = re.RegexFlag.NOFLAG) -> list[str]:
    match_with_flags = lambda s: matched_source(pattern, s, flags)
    return list(filter(None, map(match_with_flags, strings)))

def first_matched(pattern: re.Pattern, strings: Sequence[str], flags: re.RegexFlag = re.RegexFlag.NOFLAG) -> str | None:
    return next(iter(map_matched_source(pattern, strings, flags)), None)

# %% ../nbs/00_core.ipynb 12
DataLike: TypeAlias = Sequence | np.ndarray | pd.Series
'''Type alias for data like variables e.g. NumPy's `np.ndarray`, Pandas's `pd.Series`, etc''';

# %% ../nbs/00_core.ipynb 13
class ColorDType(StrEnum):
    '''Color datatypes'''
    OBJECT = auto()
    STRING = auto()
    NUMBER = auto()
    BOOLEAN = auto()
    CATEGORY = auto()
    DATETIME = auto()
    UNKNOWN = auto()

# %% ../nbs/00_core.ipynb 14
class DTypePattern(StrEnum):
    '''RegEx patterns for Pandas / Numpy'''
    
    OBJECT = r'^(object).*'
    '''object regex pattern'''
    
    STRING = r'^(u?)(str|utf|byte).*'
    '''string regex pattern'''
    
    NUMBER = r'^(u?)(int|float|complex|number|long|short|double).*'
    '''number regex pattern'''
    
    BOOLEAN = r'^(bool|boolean).*'
    '''boolean regex pattern'''
    
    CATEGORY = r'^(category|categorical).*'
    '''categorical regex pattern'''
    
    DATETIME = r'^(date|time).*'
    '''date / time regex pattern''';
    
    UNKNOWN = r'.*'
    
    def pattern(self) -> re.Pattern:
        return re.Pattern(self.value)
    
    def match(self, string: str, flags: re.RegexFlag = re.RegexFlag.IGNORECASE) -> re.Match[str] | None:
        return re.match(self.pattern(), string, flags)

# %% ../nbs/00_core.ipynb 16
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/00_core.ipynb 17
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/00_core.ipynb 19
def check_columns(
    df: pd.DataFrame, 
    value: Any, 
    subset: list[str] | 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_matched(str(value), cols, flags)
    if first:
        return df[first]
    return None

def check_values(
    df: pd.DataFrame, 
    value: Any, 
    subset: list[str] | None = None, 
    insensitive: bool = True
) -> Any | None:
    '''
    Search dataframe `df`'s values in all columns (or subset thereof) for an
    entry equal to `value` case `insensitive` as specified by the user.
    '''
    el = str(value).lower() if insensitive else value
    sf = df if subset is None else df[subset]
    
    if insensitive and ColorDType.OBJECT == sf.columns.dtype:
        mask = sf.apply(lambda col: col.str.lower() == el if ColorDType.OBJECT else False)
    else:
        mask = sf == el
    
    idx = mask.any(axis=1)
    if idx.any():
        return df.index[idx][0]
    
    # Check index
    if insensitive:
        idx = df.index.map(str).str.lower()
        if el in idx:
            return df.index[idx == el][0]
    else:
        if value in df.index:
            return value
        
    return None

# %% ../nbs/00_core.ipynb 20
def extract_colors(
    adata: ad.AnnData,
    value: Any,
    layer: str | None = None,
    annots: list | None = None,
    subset: list | None = None,
    insensitive: bool = True,
) -> DataLike | None:
    parts = [annots] if isinstance(annots, str) else annots or ['obs', 'var', 'gene']
    
    colors = None
    
    mat = adata.layers.get(layer, adata.X)
    
    # Check adata.obs and adata.var
    for part in parts:
        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:
            colors = col
            break
        
        # Check if value is a value in any column
        idx = check_values(df, value, subset, insensitive)
        if idx is None: continue

        match part:
            case 'obs': 
                if isinstance(idx, str):
                    idx = adata.obs.index.get_loc(idx)
                colors = mat[idx, :]
                break
            case 'var': 
                if isinstance(idx, str):
                    idx = adata.var.index.get_loc(idx)
                colors = mat[:, idx]
                break
            case 'gene': 
                if isinstance(idx, str):
                    idx = adata.var.index.get_loc(idx)
                colors = mat[:, idx]
                break
            case _:
                ...
            
    if colors is None or colors.shape[0] != adata.n_obs:
        colors = np.zeros(adata.n_obs)
    
    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/00_core.ipynb 22
class AnnDashboard:
    def __init__(self, adata: ad.AnnData, *args, **kwargs):
        # Initialize Dash app
        self._app: Dash = Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
        
        gene_col = kwargs.pop('gene_col', 'gene_symbol')
        if gene_col not in adata.var_keys(): 
            gene_col = None
            
        self.layer_keys: list[str] = list(adata.layers.keys())
        self.basis_keys: list[str] = self._filter_keys(adata.obsm_keys(), r'^X_')
        
        self.gene_items = self._make_gene_items(adata, gene_col)
        
        self.obs_options = adata.obs_keys()
        self.var_options = adata.var_keys()
        self.dim_default = [dict(label=str(i), value=i) for i in range(3)]
        self.dim_options = self.dim_default
        
        self._basis = None
        self._fig = None
        self._make_layout()
        self._make_callbacks(adata) 
            
    
    def _make_gene_items(self, adata: ad.AnnData, gene_col: str | None = None) -> list[dict[str,str]]:
        genes = []
        
        for index in adata.var.index:
            label = getattr(index, 'value', index)
            item = dict(label=label, value=index)
            
            if gene_col in adata.var_keys():
                value = adata.var.loc[index, gene_col]
                label = getattr(value, 'value', value)
                item.update(label=label)
            
            genes.append(item)
            
        return genes
        
    def _filter_keys(self, keys: list[str], pattern: re.Pattern | str | None = None) -> list[str]:
        pattern = pattern or r'.*'
        filt_fn = lambda s: re.match(pattern, s)
        return list(filter(filt_fn, keys))
        
    def _make_select(
        self,
        id: str, 
        keys: list[str], 
        extra: list[dict] | None = None,                 
        select: ddbc.Component | str = 'Dropdown',
        default = None, 
        use_key = False,
        suffix: str | None = '_dropdown',
        **kwargs
    ):
        options = [dict(label=key, value=key) for key in keys]
        options += ([extra] if isinstance(extra, dict) else extra if (extra and len(extra)) else [])
        
        html_id = f'{id}{suffix}' if suffix else id
        
        if not use_key or use_key and not keys: default = (default or None)
        
        if use_key and isinstance(keys, list) and len(keys):
            default = keys[0] 
            
        if isinstance(default, dict):
            default = default.get('value', default)
        
        if isinstance(select, str): 
            select = dcc.Dropdown
        
        return select(id=html_id, options=options, value=default, **kwargs)
        
    def _make_radios(
        self, id: str, keys: list[str], extra: list[dict] | None = None,         
        default = None, use_key = False, suffix: str | None = '_radios', **kwargs
    ) -> dcc.RadioItems:
        return self._make_select(id, keys, extra, dcc.RadioItems, default, use_key, suffix, **kwargs)    
    
    def _make_drowdown(
        self, 
        id: str, keys: list[str], extra: list[dict] | None = None,         
        default = None, use_key = False, suffix: str | None = '_dropdown', **kwargs
    ) -> dcc.Dropdown:
        return self._make_select(id, keys, extra, dcc.Dropdown, default, use_key, suffix, **kwargs)
        
    def _make_layout(self):
        # Layout
        self._app.layout = dbc.Container([
            dbc.Row([
                dbc.Col([
                    html.H3("Embedding Basis"),
                    self._make_drowdown('basis', self.basis_keys, use_key=True),
                ], width=3),
                dbc.Col([
                    html.H3("Layer"),
                    self._make_drowdown(
                        'layer', self.layer_keys, extra=dict(label='X (default)', value='X'), default='X',
                        placeholder='Select a layer if other than X',
                    ),
                ], width=3),
                dbc.Col([
                    html.H3("Color By"),
                    self._make_radios('color_by', ['obs', 'gene'], default='gene', use_key=False, labelStyle={'display': 'block'}),
                ], width=3),
                dbc.Col([
                    html.H3("Plot Dims"),
                    self._make_radios('plot_dims', ['2D', '3D'], default='3D', use_key=False, labelStyle={'display': 'block'}),
                ], width=3),
            ]),
            dbc.Row([
                dbc.Col([
                    html.H3("Color Variable"),
                    dcc.Dropdown(
                        id='color_var_dropdown',
                        options=self.obs_options,
                        value=None,
                        multi=False
                    ),
                ], width=6),
                dbc.Col([
                    html.H3("Gene Selection"),
                    dcc.Dropdown(
                        id='gene_dropdown', 
                        options=self.gene_items, 
                        value=self.gene_items,
                        multi=False, 
                        placeholder='Select a gene',
                    ),
                ], width=6),
            ]),
            dbc.Row([
                dbc.Col([
                    html.H3("X Axis"),
                    dcc.Dropdown(id='x_axis_dropdown', options=self.dim_options, value=0)
                ], width=3),
                dbc.Col([
                    html.H3("Y Axis"),
                    dcc.Dropdown(id='y_axis_dropdown', options=self.dim_options, value=1)
                ], width=3),
                dbc.Col([
                    html.H3("Z Axis"),
                    dcc.Dropdown(id='z_axis_dropdown', options=self.dim_options, value=2, disabled=True)
                ], width=3),
            ]),
            dbc.Row([
                dbc.Col([
                    dcc.Graph(id='embedding_plot')
                ])
            ])
        ], fluid=True)
        
        return self
    
    def _make_colors(self, adata: ad.AnnData, color_by, color_var, gene: str | None = None, layer: str | None = None):
        # Handle coloring
        match color_by:
            case 'obs': 
                # color = adata.obs[color_var]
                color = extract_colors(adata, color_var, annots=[color_by], layer=layer)
            case 'var': 
                color = extract_colors(adata, color_var, annots=[color_by], layer=layer)
            case 'gene':
                color = extract_colors(adata, gene, annots=['var'], layer=layer)
            case _: 
                if gene is not None:
                    color = extract_colors(adata, gene, layer=layer)
                else:
                    color = extract_colors(adata, color_var, layer=layer)
        return color
    
    def _make_plot(
        self, adata: ad.AnnData, basis, layer, color_by, color_var, 
        gene, plot_type, x_axis, y_axis, z_axis
    ) -> go.Figure | dict:
        if not basis or basis not in adata.obsm: return {}
        if layer == 'X': layer = None
        
        data = adata.obsm[basis]
        if data.shape[1] < 2: return {}
        
        colors = self._make_colors(adata, color_by, color_var, gene, layer)
        
        df = pd.DataFrame(
            dict(
                x=data[:, x_axis], 
                y=data[:, y_axis],
                z=data[:, z_axis] if plot_type == '3D' else None
            ), index=adata.obs_names
        )
        match plot_type:
            case '2D': 
                fig = px.scatter(df, x='x', y='y', color=colors, hover_data=[adata.obs_names])
            case '3D': 
                fig = px.scatter_3d(df, x='x', y='y', z='z', color=colors, hover_data=[adata.obs_names])
            case _: 
                fig = {}
        return fig
    
    def _make_callbacks(self, adata: ad.AnnData):
        @self._app.callback(
            Output('x_axis_dropdown', 'options'),
            Output('y_axis_dropdown', 'options'),
            Output('z_axis_dropdown', 'options'),
            Output('z_axis_dropdown', 'disabled'),
            Input('basis_dropdown', 'value'),
            Input('plot_dims_radios', 'value')
        )
        def update_axes_options(basis, plot_type):
            x_opts, y_opts, z_opts, z_off = [], [], [], True
            if not basis or basis not in adata.obsm:
                return x_opts, y_opts, z_opts, z_off
            
            
            dims = adata.obsm[basis].shape[1]
            name = basis.removeprefix('X_')
            
            options = [dict(label=f'{name}_{i}', value=i) for i in range(dims)]
            
            self._basis = basis
            # self._axis_options = options
            
            x_opts = options
            y_opts = options
            
            if dims >= 3 and plot_type == '3D': 
                z_opts, z_off = options, False
            
            return x_opts, y_opts, z_opts, z_off
    
    
        # what goes in the color_by column
        @self._app.callback(
            Output('color_var_dropdown', 'options'),
            Output('color_var_dropdown', 'value'),
            Output('gene_dropdown', 'disabled'),
            Input('color_by_radios', 'value')
        )
        def update_color_variable_options(color_by: str | None = None):
            gene_off = True
            match color_by:
                case 'gene':
                    options = [] # self.gene_items
                    value = options[0] if len(options) else None
                    gene_off = False
                case 'obs':
                    options = [{'label': col, 'value': col} for col in adata.obs.columns]
                    value = adata.obs.columns[0] if len(adata.obs.columns) > 0 else None
                case 'var':
                    options = [{'label': col, 'value': col} for col in adata.var.columns]
                    value = adata.var.columns[0] if len(adata.var.columns) > 0 else None
                    gene_off = False
                case _:
                    options = self.gene_items
                    value = None
            if isinstance(value, dict):
                value = value.get('value', value)
            return options, value, gene_off
     
     
        
            
    
        @self._app.callback(
            Output('z_axis_dropdown', 'value'),
            Input('z_axis_dropdown', 'options')
        )
        def set_z_default(options):
            if options: 
                return 2 if len(options) > 2 else options[-1]['value']
            return None
    
        @self._app.callback(
            Output('plot_dims_radios', 'options'),
            Output('plot_dims_radios', 'value'),
            # Output('z_axis_dropdown', 'disabled'),
            Input('basis_dropdown', 'value')
        )
        def update_plot_type_options(basis):
            opts, value = [], '3D'
            if (not basis or basis not in adata.obsm):
                return opts, value
            
            item2d = dict(label='2D', value='2D')
            item3d = dict(label='3D', value='3D')
            
            dims = adata.obsm[basis].shape[1]
            opts = [item2d]
            if dims >= 3: 
                opts += [item3d]
            return opts, value#
    
        @self._app.callback(
            Output('embedding_plot', 'figure'),
            Input('basis_dropdown', 'value'),
            Input('layer_dropdown', 'value'),
            Input('color_by_radios', 'value'),
            Input('color_var_dropdown', 'value'),
            Input('gene_dropdown', 'value'),
            Input('plot_dims_radios', 'value'),
            Input('x_axis_dropdown', 'value'),
            Input('y_axis_dropdown', 'value'),
            Input('z_axis_dropdown', 'value'),
        )
        def update_plot(basis, layer, color_by, color_var, gene, plot_type, x_axis, y_axis, z_axis):
            if layer == 'X': layer = None
            fig = self._make_plot(adata, basis, layer, color_by, color_var, gene, plot_type, x_axis, y_axis, z_axis)
            self._fig = fig
            return fig
            
    def run(self, debug: bool = False, **kwargs):
        self._app.run_server(debug=debug)
