# Path: scib/preprocessing.py
# def score_cell_cycle(adata, organism='mouse'):
#     """
#     Tirosh et al. cell cycle marker genes downloaded from
#     https://raw.githubusercontent.com/theislab/scanpy_usage/master/180209_cell_cycle/data/regev_lab_cell_cycle_genes.txt
#     return: (s_genes, g2m_genes)
#         s_genes: S-phase genes
#         g2m_genes: G2- and M-phase genes
#     """
#     import pathlib
#     root = pathlib.Path(__file__).parent
# 
#     cc_files = {'mouse': [root / 'resources/s_genes_tirosh.txt',
#                           root / 'resources/g2m_genes_tirosh.txt'],
#                 'human': [root / 'resources/s_genes_tirosh_hm.txt',
#                           root / 'resources/g2m_genes_tirosh_hm.txt']}
# 
#     with open(cc_files[organism][0], "r") as f:
#         s_genes = [x.strip() for x in f.readlines() if x.strip() in adata.var.index]
#     with open(cc_files[organism][1], "r") as f:
#         g2m_genes = [x.strip() for x in f.readlines() if x.strip() in adata.var.index]
# 
#     if (len(s_genes) == 0) or (len(g2m_genes) == 0):
#         rand_choice = np.random.randint(1, adata.n_vars, 10)
#         rand_genes = adata.var_names[rand_choice].tolist()
#         raise ValueError(f"cell cycle genes not in adata\n organism: {organism}\n varnames: {rand_genes}")
# 
#     sc.tl.score_genes_cell_cycle(adata, s_genes, g2m_genes)
# 
# Path: scib/utils.py
# def check_adata(adata):
#     if type(adata) is not anndata.AnnData:
#         raise TypeError('Input is not a valid AnnData object')
# 
# Path: scib/metrics/pcr.py
# def pc_regression(
#         data,
#         covariate,
#         pca_var=None,
#         n_comps=50,
#         svd_solver='arpack',
#         verbose=False
# ):
#     """
#     :params data: Expression or PC matrix. Assumed to be PC, if pca_sd is given.
#     :param covariate: series or list of batch assignments
#     :param n_comps: number of PCA components for computing PCA, only when pca_sd is not given.
#         If no pca_sd is not defined and n_comps=None, compute PCA and don't reduce data
#     :param pca_var: Iterable of variances for `n_comps` components.
#         If `pca_sd` is not `None`, it is assumed that the matrix contains PC,
#         otherwise PCA is computed on `data`.
#     :param svd_solver:
#     :param verbose:
#     :return:
#         R2Var of regression
#     """
# 
#     if isinstance(data, (np.ndarray, sparse.csr_matrix, sparse.csc_matrix)):
#         matrix = data
#     else:
#         raise TypeError(f'invalid type: {data.__class__} is not a numpy array or sparse matrix')
# 
#     # perform PCA if no variance contributions are given
#     if pca_var is None:
# 
#         if n_comps is None or n_comps > min(matrix.shape):
#             n_comps = min(matrix.shape)
# 
#         if n_comps == min(matrix.shape):
#             svd_solver = 'full'
# 
#         if verbose:
#             print("compute PCA")
#         pca = sc.tl.pca(matrix, n_comps=n_comps, use_highly_variable=False,
#                         return_info=True, svd_solver=svd_solver, copy=True)
#         X_pca = pca[0].copy()
#         pca_var = pca[3].copy()
#         del pca
#     else:
#         X_pca = matrix
#         n_comps = matrix.shape[1]
# 
#     ## PC Regression
#     if verbose:
#         print("fit regression on PCs")
# 
#     # handle categorical values
#     if pd.api.types.is_numeric_dtype(covariate):
#         covariate = np.array(covariate).reshape(-1, 1)
#     else:
#         if verbose:
#             print("one-hot encode categorical values")
#         covariate = pd.get_dummies(covariate)
# 
#     # fit linear model for n_comps PCs
#     r2 = []
#     for i in range(n_comps):
#         pc = X_pca[:, [i]]
#         lm = LinearRegression()
#         lm.fit(covariate, pc)
#         r2_score = np.maximum(0, lm.score(covariate, pc))
#         r2.append(r2_score)
# 
#     Var = pca_var / sum(pca_var) * 100
#     R2Var = sum(r2 * Var) / 100
# 
#     return R2Var


import numpy as np
import pandas as pd
from ..preprocessing import score_cell_cycle
from ..utils import check_adata
from .pcr import pc_regression
    # sanity check: subsets have same number of rows?
    if raw_sub.shape[0] != int_sub.shape[0]:
        raise ValueError(
            f'batch "{batch}" of batch_key "{batch_key}" has unequal number of '
            f'entries before and after integration.\n'
            f'before: {raw_sub.shape[0]} after: {int_sub.shape[0]}'
        )

    # compute cc scores if necessary
    if recompute_cc:
        if verbose:
            print("Score cell cycle...")
        score_cell_cycle(raw_sub, organism=organism)

    # regression variable
    covariate = raw_sub.obs[['S_score', 'G2M_score']]

    # PCR on adata before integration
    if recompute_pcr:
        before = pc_regression(
            raw_sub.X,
            covariate,
            pca_var=None,
            n_comps=n_comps,
            verbose=verbose
        )
    else:
        before = pd.Series(raw_sub.uns[pcr_key])

    # PCR on adata after integration

# Path: scib/preprocessing.py
# def score_cell_cycle(adata, organism='mouse'):
#     """
#     Tirosh et al. cell cycle marker genes downloaded from
#     https://raw.githubusercontent.com/theislab/scanpy_usage/master/180209_cell_cycle/data/regev_lab_cell_cycle_genes.txt
#     return: (s_genes, g2m_genes)
#         s_genes: S-phase genes
#         g2m_genes: G2- and M-phase genes
#     """
#     import pathlib
#     root = pathlib.Path(__file__).parent
# 
#     cc_files = {'mouse': [root / 'resources/s_genes_tirosh.txt',
#                           root / 'resources/g2m_genes_tirosh.txt'],
#                 'human': [root / 'resources/s_genes_tirosh_hm.txt',
#                           root / 'resources/g2m_genes_tirosh_hm.txt']}
# 
#     with open(cc_files[organism][0], "r") as f:
#         s_genes = [x.strip() for x in f.readlines() if x.strip() in adata.var.index]
#     with open(cc_files[organism][1], "r") as f:
#         g2m_genes = [x.strip() for x in f.readlines() if x.strip() in adata.var.index]
# 
#     if (len(s_genes) == 0) or (len(g2m_genes) == 0):
#         rand_choice = np.random.randint(1, adata.n_vars, 10)
#         rand_genes = adata.var_names[rand_choice].tolist()
#         raise ValueError(f"cell cycle genes not in adata\n organism: {organism}\n varnames: {rand_genes}")
# 
#     sc.tl.score_genes_cell_cycle(adata, s_genes, g2m_genes)
# 
# Path: scib/utils.py
# def check_adata(adata):
#     if type(adata) is not anndata.AnnData:
#         raise TypeError('Input is not a valid AnnData object')
# 
# Path: scib/metrics/pcr.py
# def pc_regression(
#         data,
#         covariate,
#         pca_var=None,
#         n_comps=50,
#         svd_solver='arpack',
#         verbose=False
# ):
#     """
#     :params data: Expression or PC matrix. Assumed to be PC, if pca_sd is given.
#     :param covariate: series or list of batch assignments
#     :param n_comps: number of PCA components for computing PCA, only when pca_sd is not given.
#         If no pca_sd is not defined and n_comps=None, compute PCA and don't reduce data
#     :param pca_var: Iterable of variances for `n_comps` components.
#         If `pca_sd` is not `None`, it is assumed that the matrix contains PC,
#         otherwise PCA is computed on `data`.
#     :param svd_solver:
#     :param verbose:
#     :return:
#         R2Var of regression
#     """
# 
#     if isinstance(data, (np.ndarray, sparse.csr_matrix, sparse.csc_matrix)):
#         matrix = data
#     else:
#         raise TypeError(f'invalid type: {data.__class__} is not a numpy array or sparse matrix')
# 
#     # perform PCA if no variance contributions are given
#     if pca_var is None:
# 
#         if n_comps is None or n_comps > min(matrix.shape):
#             n_comps = min(matrix.shape)
# 
#         if n_comps == min(matrix.shape):
#             svd_solver = 'full'
# 
#         if verbose:
#             print("compute PCA")
#         pca = sc.tl.pca(matrix, n_comps=n_comps, use_highly_variable=False,
#                         return_info=True, svd_solver=svd_solver, copy=True)
#         X_pca = pca[0].copy()
#         pca_var = pca[3].copy()
#         del pca
#     else:
#         X_pca = matrix
#         n_comps = matrix.shape[1]
# 
#     ## PC Regression
#     if verbose:
#         print("fit regression on PCs")
# 
#     # handle categorical values
#     if pd.api.types.is_numeric_dtype(covariate):
#         covariate = np.array(covariate).reshape(-1, 1)
#     else:
#         if verbose:
#             print("one-hot encode categorical values")
#         covariate = pd.get_dummies(covariate)
# 
#     # fit linear model for n_comps PCs
#     r2 = []
#     for i in range(n_comps):
#         pc = X_pca[:, [i]]
#         lm = LinearRegression()
#         lm.fit(covariate, pc)
#         r2_score = np.maximum(0, lm.score(covariate, pc))
#         r2.append(r2_score)
# 
#     Var = pca_var / sum(pca_var) * 100
#     R2Var = sum(r2 * Var) / 100
# 
#     return R2Var


# Path: scib/metrics/cell_cycle.py
import numpy as np
import pandas as pd
from ..preprocessing import score_cell_cycle
from ..utils import check_adata
from .pcr import pc_regression

    # sanity check: subsets have same number of rows?
    if raw_sub.shape[0] != int_sub.shape[0]:
        raise ValueError(
            f'batch "{batch}" of batch_key "{batch_key}" has unequal number of '
            f'entries before and after integration.\n'
            f'before: {raw_sub.shape[0]} after: {int_sub.shape[0]}'
        )

    # compute cc scores if necessary
    if recompute_cc:
        if verbose:
            print("Score cell cycle...")
        score_cell_cycle(raw_sub, organism=organism)

    # regression variable
    covariate = raw_sub.obs[['S_score', 'G2M_score']]

    # PCR on adata before integration
    if recompute_pcr:
        before = pc_regression(
            raw_sub.X,
            covariate,
            pca_var=None,
            n_comps=n_comps,
            verbose=verbose
        )
    else:
        before = pd.Series(raw_sub.uns[pcr_key])

    # PCR on adata after integration

    after = pc_regression(