# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/100_test.ipynb.

# %% auto 0
__all__ = ['make_adata']

# %% ../nbs/100_test.ipynb 6
import os, re, io, inspect, string, itertools, requests, tarfile, gzip, shutil, pickle, warnings
from pathlib import Path
from contextlib import redirect_stderr

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 sklearn.datasets import make_blobs
    from sklearn.cluster import KMeans
    from sklearn.preprocessing import StandardScaler


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

# %% ../nbs/100_test.ipynb 8
from typing import Any, NamedTuple, Optional

# %% ../nbs/100_test.ipynb 10
#| export


# %% ../nbs/100_test.ipynb 12
from atyp import AnnData, StrQ, CallQ, IterStr, IndexQ, Strs, NPArray, DataFrame
from psrc import dirdirs

from nchr import U1
from nlit import (GENE_SYMBOL, ENSEMBL_ID, X_PHATE_HVG, X_PHATE)

from chck import isarr, islist, isany, isnone, isad, notnone, iscall, isstr
from asto import asarr

# %% ../nbs/100_test.ipynb 14
#| export


# %% ../nbs/100_test.ipynb 16
def make_adata(
    cons: tuple = ('wt', 'ko', ),
    days: tuple = (1, 5, 10, ),
    reps: tuple = (1, 2),    
    cells: int = 100,
    feats: int = 10,

    path: Path = '../data',
    filename: str = 'adata.h5ad',
    cluster_key: str = 'cluster',
    resolution: float = 0.5,
    seed: int = 3,
):
    
    np.random.seed(seed)

    genes = list(map(lambda i: f'gene_{i}', range(1, feats + 1)))
    total_conds = len(days) * len(cons)
    total_samps = total_conds * len(reps)
    total_cells = total_samps * cells

    barcodes = list(map(lambda i: f'cell_{i}', range(1, total_cells + 1)))

    resdir = Path(path)
    resdir.mkdir(exist_ok=True, parents=True)
    datafile = (resdir / filename)

    n_pcs = (feats / 2)
    
    skip = datafile.exists()

    params = dict(
        cons=cons, days=days, reps=reps, cells=cells, feats=feats, path=path, filename=filename,
        clstkey=cluster_key, resolution=resolution, seed=seed,
        genes=genes, total_cells=total_cells, barcodes=barcodes,
        datafile=datafile, skip=skip
    )

    if skip: 
        adata = ad.read_h5ad(datafile)
        return adata, params
    
    X, y = make_blobs(
        n_samples = total_cells,
        n_features = feats,
        centers = np.random.randn(total_samps, feats),
        random_state=seed
    )

    samp_dict = dict(list(enumerate(itertools.product(cons, days, reps))))
    cons_dict = {k: v[0] for k, v in samp_dict.items()}
    days_dict = {k: v[1] for k, v in samp_dict.items()}
    reps_dict = {k: v[2] for k, v in samp_dict.items()}
    cxrs_dict = {k: ''.join((v[0], str(v[2]))) for k, v in samp_dict.items()}
    cxts_dict = {k: ''.join((v[0], '_d', str(v[1]))) for k, v in samp_dict.items()}
    sample_idx = pd.Series(y, name='sample_idx').astype('category')

    with redirect_stderr(f_err := io.StringIO()):
        pop = phate.PHATE(3)
        phy = pop.fit_transform(X)

    df = pd.DataFrame(X, columns=genes, index=barcodes).assign(
        barcodes = barcodes,
        con=sample_idx.replace(cons_dict).values,
        day=sample_idx.replace(days_dict).values,
        rep=sample_idx.replace(reps_dict).values,
        cxr=sample_idx.replace(cxrs_dict).values,
        cxt=sample_idx.replace(cxts_dict).values,
        sample_idx=sample_idx.values,
    )

    adata = ad.AnnData(
        df.loc[:, df.columns.isin(genes)],
        obs=df.loc[:, ~df.columns.isin(genes)],
        var=pd.Series(df.columns[df.columns.isin(genes)], index=genes, name=GENE_SYMBOL).to_frame(),
        obsm={X_PHATE: phy},
    )
    adata.var_names_make_unique()

    sc.pp.neighbors(adata, n_neighbors=10)
    sc.tl.dendrogram(adata, groupby='cxr', use_rep=X_PHATE)

    sc.tl.leiden(adata, key_added=cluster_key, resolution=resolution, random_state=seed)

    adata.write(datafile)    
    return adata, params
