# Path: scib/utils.py
# def split_batches(adata, batch, hvg=None, return_categories=False):
#     split = []
#     batch_categories = adata.obs[batch].unique()
#     if hvg is not None:
#         adata = adata[:, hvg]
#     for i in batch_categories:
#         split.append(adata[adata.obs[batch] == i].copy())
#     if return_categories:
#         return split, batch_categories
#     return split


import numpy as np
import scanpy as sc
from ..utils import split_batches



def precompute_hvg_batch(adata, batch, features, n_hvg=500, save_hvg=False):
    adata_list = split_batches(adata, batch, hvg=features)
    hvg_dir = {}
    for i in adata_list:
        sc.pp.filter_genes(i, min_cells=1)
        n_hvg_tmp = np.minimum(n_hvg, int(0.5 * i.n_vars))
        if n_hvg_tmp < n_hvg:
            print(i.obs[batch][0] + ' has less than the specified number of genes')
            print('Number of genes: ' + str(i.n_vars))
        hvg = sc.pp.highly_variable_genes(i, flavor='cell_ranger', n_top_genes=n_hvg_tmp, inplace=False)
        hvg_dir[i.obs[batch][0]] = i.var.index[hvg['highly_variable']]

    if save_hvg:
        adata.uns['hvg_before'] = hvg_dir
    else:
        return hvg_dir


def hvg_overlap(adata_pre, adata_post, batch, n_hvg=500, verbose=False):
    hvg_post = adata_post.var_names


# Path: scib/utils.py
# def split_batches(adata, batch, hvg=None, return_categories=False):
#     split = []
#     batch_categories = adata.obs[batch].unique()
#     if hvg is not None:
#         adata = adata[:, hvg]
#     for i in batch_categories:
#         split.append(adata[adata.obs[batch] == i].copy())
#     if return_categories:
#         return split, batch_categories
#     return split


# Path: scib/metrics/highly_variable_genes.py
import numpy as np
import scanpy as sc
from ..utils import split_batches




def precompute_hvg_batch(adata, batch, features, n_hvg=500, save_hvg=False):
    adata_list = split_batches(adata, batch, hvg=features)
    hvg_dir = {}
    for i in adata_list:
        sc.pp.filter_genes(i, min_cells=1)
        n_hvg_tmp = np.minimum(n_hvg, int(0.5 * i.n_vars))
        if n_hvg_tmp < n_hvg:
            print(i.obs[batch][0] + ' has less than the specified number of genes')
            print('Number of genes: ' + str(i.n_vars))
        hvg = sc.pp.highly_variable_genes(i, flavor='cell_ranger', n_top_genes=n_hvg_tmp, inplace=False)
        hvg_dir[i.obs[batch][0]] = i.var.index[hvg['highly_variable']]

    if save_hvg:
        adata.uns['hvg_before'] = hvg_dir
    else:
        return hvg_dir


def hvg_overlap(adata_pre, adata_post, batch, n_hvg=500, verbose=False):
    hvg_post = adata_post.var_names


    adata_post_list = split_batches(adata_post, batch)