# Path: scib/utils.py
# def check_batch(batch, obs, verbose=False):
#     if batch not in obs:
#         raise ValueError(f'column {batch} is not in obs')
#     elif verbose:
#         print(f'Object contains {obs[batch].nunique()} batches.')
# 
# Path: scib/metrics/utils.py
# class RootCellError(Exception):
#     def __init__(self, message):
#         self.message = message


import numpy as np
import pandas as pd
import scanpy as sc
from scipy.sparse.csgraph import connected_components
from ..utils import check_batch
from .utils import RootCellError
    :param label_key: column in `adata_pre.obs` of the groups used to precompute the trajectory
    :param pseudotime_key: column in `adata_pre.obs` in which the pseudotime is saved in.
        Column can contain empty entries, the dataset will be subset to the cells with scores.
    :param batch_key: set to batch key if if you want to compute the trajectory metric by batch
    """
    # subset to cells for which pseudotime has been computed
    cell_subset = adata_pre.obs.index[adata_pre.obs[pseudotime_key].notnull()]
    adata_pre_ti = adata_pre[cell_subset]
    adata_post_ti = adata_post[cell_subset]
    try:
        iroot, adata_post_ti2 = get_root(adata_pre_ti, adata_post_ti, label_key, pseudotime_key)
    except RootCellError:
        print('No root cell found, setting trajectory conservation metric to 0.')
        return 0  # failure to find root cell means no TI conservation

    adata_post_ti2.uns['iroot'] = iroot

    sc.tl.dpt(adata_post_ti2)  # stored in 'dpt_pseudotime'
    adata_post_ti2.obs.loc[adata_post_ti2.obs['dpt_pseudotime'] > 1, 'dpt_pseudotime'] = 0
    adata_post_ti.obs['dpt_pseudotime'] = 0
    adata_post_ti.obs['dpt_pseudotime'] = adata_post_ti2.obs['dpt_pseudotime']
    adata_post_ti.obs['dpt_pseudotime'].fillna(0, inplace=True)

    if batch_key == None:
        pseudotime_before = adata_pre_ti.obs[pseudotime_key]
        pseudotime_after = adata_post_ti.obs['dpt_pseudotime']
        correlation = pseudotime_before.corr(pseudotime_after, 'spearman')
        return (correlation + 1) / 2  # scaled
    else:
        check_batch(batch_key, adata_pre.obs)

# Path: scib/utils.py
# def check_batch(batch, obs, verbose=False):
#     if batch not in obs:
#         raise ValueError(f'column {batch} is not in obs')
#     elif verbose:
#         print(f'Object contains {obs[batch].nunique()} batches.')
# 
# Path: scib/metrics/utils.py
# class RootCellError(Exception):
#     def __init__(self, message):
#         self.message = message


# Path: scib/metrics/trajectory.py
import numpy as np
import pandas as pd
import scanpy as sc
from scipy.sparse.csgraph import connected_components
from ..utils import check_batch
from .utils import RootCellError

    :param label_key: column in `adata_pre.obs` of the groups used to precompute the trajectory
    :param pseudotime_key: column in `adata_pre.obs` in which the pseudotime is saved in.
        Column can contain empty entries, the dataset will be subset to the cells with scores.
    :param batch_key: set to batch key if if you want to compute the trajectory metric by batch
    """
    # subset to cells for which pseudotime has been computed
    cell_subset = adata_pre.obs.index[adata_pre.obs[pseudotime_key].notnull()]
    adata_pre_ti = adata_pre[cell_subset]
    adata_post_ti = adata_post[cell_subset]
    try:
        iroot, adata_post_ti2 = get_root(adata_pre_ti, adata_post_ti, label_key, pseudotime_key)
    except RootCellError:
        print('No root cell found, setting trajectory conservation metric to 0.')
        return 0  # failure to find root cell means no TI conservation

    adata_post_ti2.uns['iroot'] = iroot

    sc.tl.dpt(adata_post_ti2)  # stored in 'dpt_pseudotime'
    adata_post_ti2.obs.loc[adata_post_ti2.obs['dpt_pseudotime'] > 1, 'dpt_pseudotime'] = 0
    adata_post_ti.obs['dpt_pseudotime'] = 0
    adata_post_ti.obs['dpt_pseudotime'] = adata_post_ti2.obs['dpt_pseudotime']
    adata_post_ti.obs['dpt_pseudotime'].fillna(0, inplace=True)

    if batch_key == None:
        pseudotime_before = adata_pre_ti.obs[pseudotime_key]
        pseudotime_after = adata_post_ti.obs['dpt_pseudotime']
        correlation = pseudotime_before.corr(pseudotime_after, 'spearman')
        return (correlation + 1) / 2  # scaled
    else:
        check_batch(batch_key, adata_pre.obs)

        check_batch(batch_key, adata_post.obs)