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

# %% auto 0
__all__ = ['make_dummy_genes', 'make_dummy_cells', 'make_mock_genes_x_tbins', 'make_mock_trajectories', 'df_trj_to_cells_x_tbins',
           'df_trj_to_genes_x_tbins']

# %% ../nbs/04_mock.ipynb 3
import pandas as pd, numpy as np
from typing import List, Optional, Union

# %% ../nbs/04_mock.ipynb 4
from atyp import SeriesLike, DataFrame, NPArray

# %% ../nbs/04_mock.ipynb 5
from .static import (MOCK_GENES)

# %% ../nbs/04_mock.ipynb 6
def make_dummy_genes(n:int):
    return [f'gene_{i}' for i in range(1, n+1)]

def make_dummy_cells(n:int):
    return [f'cell{i}' for i in range(1, n+1)]

# %% ../nbs/04_mock.ipynb 7
def make_mock_genes_x_tbins(
    genes: Optional[Union[int, List[str]]] = MOCK_GENES,
    tbins: int = 100,
) -> DataFrame:
    
    if isinstance(genes, int):
        genes = make_dummy_genes(genes)        
    n_genes = len(genes)

    forward = lambda : np.sort(np.random.randint(0, n_genes, (tbins)))
    reverse = lambda : forward()[::-1]

    timeseries = [forward() if i % 2 == 0 else reverse() for i in range(n_genes)]
    df_trj = pd.DataFrame(timeseries, index=genes)
    return df_trj

# %% ../nbs/04_mock.ipynb 9
def make_mock_trajectories(
    tbins: int = 100,
    cells: int = 50,
    genes: int = 10,
) -> NPArray:
    df = pd.DataFrame()
    for i in range(cells):
        df_gxt = make_mock_genes_x_tbins(genes, tbins)
        df = pd.concat([df, df_gxt])
    res = df.values.reshape(tbins, cells, genes)
    return res

# %% ../nbs/04_mock.ipynb 11
def df_trj_to_cells_x_tbins(trajectories:NPArray, agg_fn=np.mean) -> NPArray:
    '''
    Transpose and aggregate trajectories matrix (timebinse, cells, gene)
    to produce (cells, timebins)
    ''' 
    return agg_fn(np.transpose(trajectories, (1, 0, 2)), axis=2)

def df_trj_to_genes_x_tbins(trajectories:NPArray, agg_fn=np.mean) -> NPArray:
    '''
    Transpose and aggregate trajectories matrix (timebinse, cells, gene)
    to produce (genes, timebins)
    '''
    return agg_fn(np.transpose(trajectories, (2, 0, 1)), axis=2)
