"""Plots module for `anndash`"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_plots.ipynb.

# %% auto 0
__all__ = ['base_dropdown', 'obs_dropdown', 'var_dropdown', 'axis_dropdown', 'axis_label', 'basis_labels', 'marker_color_button',
           'dropdown_marker_color_button', 'column_dropdown', 'make_button', 'axis_index_button']

# %% ../nbs/05_plots.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

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/05_plots.ipynb 5
from anndash.hints import (
    DataLike, Columns, AnnOption, DropDownKeywords
)
from anndash.utils import (
    subset_columns, check_columns, var_index,
    get_gene_info, get_basis_keys, get_ndim
)
from .enums import DTypePattern, ColorDType
from .color import get_gene_colors, obs_color, var_color, color_dtype, extract_colors, color_series

# %% ../nbs/05_plots.ipynb 7
def base_dropdown(
    *,
    xpos: float = 0., 
    ypos: float = 0., 
    visible: bool = True, 
    xanchor: str = 'left', 
    yanchor: str = 'top', 
    direction: str = 'down',
    **kwargs: DropDownKeywords
):
    spec = dict(
        x=xpos, xanchor=xanchor, y=ypos, yanchor=yanchor,
        direction=direction, showactive=True, visible=visible
    )
    spec.update(**kwargs)
    return spec

# %% ../nbs/05_plots.ipynb 8
def base_dropdown(
    *,
    xpos: float = 0., ypos: float = 0, visible: bool = True, 
    xanchor: str = 'left', yanchor: str = 'top', direction: str = 'down',
    **kwargs
):
    spec = dict(
        x=xpos, xanchor=xanchor, y=ypos, yanchor=yanchor,
        direction=direction, showactive=True, visible=visible
    )
    spec.update(**kwargs)
    return spec

# %% ../nbs/05_plots.ipynb 9
def obs_dropdown(
    adata: ad.AnnData,
    keys: set[str] | None = None,
    label: str | None = None, 
    **kwargs
):
    label = label or f'obs'
    keys = keys or adata.obs_keys()
    spec = base_dropdown(**kwargs)
    spec.update(buttons = [
        dict(
            args=[{
                'marker.color': [
                    obs_color(adata, col)
                ],
            }],
            method='restyle',
            label=label,
        ) for col in keys
    ])
    return spec

def var_dropdown(
    adata: ad.AnnData,
    col: str | None = None,
    layer: str | None = None,
    label: str | None = None, 
    **kwargs
):
    label = label or f'var'
    keys = adata.var[col].values if col else adata.var.index
    spec = base_dropdown(**kwargs)
    spec.update(buttons = [
        dict(
            args=[{
                'marker.color': [
                    var_color(adata, col, key, layer=layer)
                ],
            }],
            method='restyle',
            label=label,
        ) for key in keys
    ])
    return spec

def axis_dropdown(
    data: np.ndarray,
    ax: str = 'x', 
    ndim: int | None = None,
    max_dim: int = 10,
    label: str | None = None, 
    **kwargs
):
    ndim = min(ndim or data.shape[-1], max_dim)
    label = label or f'{ax.title()}-axis'
    spec = base_dropdown(**kwargs)
    spec.update(buttons = [
        dict(
            args=[{ax: [data[:, i]]}],
            method='restyle',
            label=label,
        ) for i in range(ndim)
    ])
    return spec

def axis_label(
    index: int | None = 0, axis: str | None = None, name: str = 'axis', *, 
    x: str | None = None, y: str | None = None, z: str | None = None,    
) -> dict:
    if index is None: return dict()    
    axis = str(axis or x or y or z or 'x')
    return {f'{axis}': f'{name}_{index}'}

def basis_labels(
    basis: str, x: int = 0, y: int = 1, z: int | None = None,
    key_suffix: str | None = 'axis_title'
) -> dict:
    key_suffix = key_suffix or ''
    name = basis.removeprefix('X_')
    spec = axis_label(x, axis='x', name=name) | \
           axis_label(y, axis='y', name=name) | \
           axis_label(z, axis='z', name=name)
    return {f'{k}{key_suffix}': v for k, v in spec.items()}



# %% ../nbs/05_plots.ipynb 12
def marker_color_button(colors, label: str = None, method: str = 'restyle'):
    args = [{'marker.color': [colors],}]
    spec = dict(args=args, label=label, method=method)
    return spec

# %% ../nbs/05_plots.ipynb 13
def dropdown_marker_color_button(
    adata,
    column: str,
    target = None,
    categorical: bool | None = None, 
    boolean: bool | None = None,
    layer : str | None = None,
    annots: list | None = None,
    subset: list | None = None, 
    insensitive: bool = True,
):
    colors = extract_colors(
        adata, 
        column,
        layer = layer,
        annots = annots, 
        subset = subset,
        insensitive = insensitive
    )
    colors = color_series(
        colors, target, 
        categorical = categorical if categorical is not None else color_dtype(colors) == 'category', 
        boolean = boolean if boolean is not None else color_dtype(colors) == bool
    )
    return marker_color_button(colors, column)

def column_dropdown(adata, keys: list, dropdown_kws: dict = {}, **kwargs):  
    spec = base_dropdown(**dropdown_kws)
    spec.update(buttons = [
        dropdown_marker_color_button(adata, col, **kwargs) for col in keys
    ])
    return spec

# %% ../nbs/05_plots.ipynb 14
def make_button(
    args: list[dict] = None,
    label: str = '', 
    method: str = 'restyle', 
    argkws: dict | None = None,
    
    marker_color: list | None = None,
    
    x: list | None = None,
    y: list | None = None,
    z: list | None = None,
    
    i: int | None = None,
    j: int | None = None,
    k: int | None = None,
    
    obs: str | None = None,
    var: str | None = None,
    basis: str | None = None,
    adata: ad.AnnData | None = None,
    layer: str | None = None,
    gene_col: str = 'gene_symbol',
    max_dim: int | None = 10,
    **kwargs
) -> dict:
    if args is None: 
        args = []
    
    if argkws is None: 
        argkws = dict()
            
    if isinstance(args, list) and len(args) and isinstance(args[0], dict):
        args[0].update(argkws)
        
        if adata is not None:
            # print(obs, marker_color)
            if obs is not None and obs in adata.obs:
                marker_color = adata.obs[obs]
            print(obs, marker_color)
            
            if var is not None and var in adata.var[gene_col]:
                marker_color = get_gene_colors(adata, var, layer, gene_col)
                
            if basis is None:
                basis = get_basis_keys(adata)[0]

            ndims = get_ndim(adata, basis, max_dim)
                        
            if i is not None and i <= ndims:
                x = adata.obsm[basis][:, i]
            
            if j is not None and j <= ndims:
                y = adata.obsm[basis][:, j]
            
            if k is not None and k <= ndims:
                z = adata.obsm[basis][:, k]
            
        if x is not None: 
            args[0].update(x=x)
        
        if y is not None: 
            args[0].update(y=y)
        
        if z is not None: 
            args[0].update(z=z)
            
        if marker_color is not None:
            args[0].update({'marker.color': marker_color})
    
    return dict(args=args, label=label, method=method, **kwargs)


# %% ../nbs/05_plots.ipynb 15
def axis_index_button(
    data: np.ndarray, 
    axis: str = 'x', 
    index: int = 0,
    label: str | None = None, 
    method: str = 'restyle'
):
    label = label or f'{axis.title()}-axis'
    args = [{axis: [data[:, index]]}]
    spec = dict(args=args, label=label, method=method)
    return spec

def axis_dropdown(
    data: np.ndarray,
    ax: str = 'x', 
    ndim: int | None = None,
    max_dim: int = 10,
    label: str | None = None, 
    dropdown_kws: dict = {},
    **kwargs
):
    ndim = min(ndim or data.shape[-1], max_dim)
    spec = base_dropdown(**dropdown_kws)
    spec.update(buttons = [
        axis_index_button(data, ax, i, label)
        for i in range(ndim)
    ])
    return spec
