import logging

import numpy as np
import pandas as pd
import scanpy as sc
import scipy.sparse

from ..exceptions import OptionalDependencyNotInstalled, RLibraryNotFound
from ..utils import check_adata, check_batch
from .utils import NeighborsError, diffusion_conn, diffusion_nn


def kBET(
    adata,
    batch_key,
    label_key,
    scaled=True,
    embed="X_pca",
    type_=None,
    return_df=False,
    verbose=False,
):
    """kBET score

    Compute the average of k-nearest neighbour batch effect test (`kBET`_) score per label.

    .. _kBET: https://doi.org/10.1038/s41592-018-0254-1

    :param adata: anndata object to compute kBET on
    :param batch_key: name of batch column in adata.obs
    :param label_key: name of cell identity labels column in adata.obs
    :param scaled: whether to scale between 0 and 1
        with 0 meaning low batch mixing and 1 meaning optimal batch mixing
        if scaled=False, 0 means optimal batch mixing and 1 means low batch mixing
    :return:
        kBET score (average of kBET per label) based on observed rejection rate.
        If ``return_df=True``, also return a ``pd.DataFrame`` with kBET observed
        rejection rate per cluster
    """
    try:
        import rpy2.rinterface_lib.callbacks
        import rpy2.rinterface_lib.embedded
        import rpy2.robjects as ro

        rpy2.rinterface_lib.callbacks.logger.setLevel(logging.ERROR)
    except ModuleNotFoundError as e:
        raise OptionalDependencyNotInstalled(e)

    check_adata(adata)
    check_batch(batch_key, adata.obs)
    check_batch(label_key, adata.obs)

    try:
        ro.r("library(kBET)")
    except rpy2.rinterface_lib.embedded.RRuntimeError as e:
        raise OptionalDependencyNotInstalled(e, "kBET")

    # compute connectivities for non-knn type data integrations
    # and increase neighborhoods for knn type data integrations
    if type_ != "knn":
        adata_tmp = sc.pp.neighbors(adata, n_neighbors=50, use_rep=embed, copy=True)
    else:
        # check if pre-computed neighbours are stored in input file
        adata_tmp = adata.copy()
        if "diffusion_connectivities" not in adata.uns["neighbors"]:
            if verbose:
                print("Compute diffusion neighbours")
            adata_tmp = diffusion_conn(adata, min_k=50, copy=True)
        adata_tmp.obsp["connectivities"] = adata_tmp.uns["neighbors"][
            "diffusion_connectivities"
        ]

    if verbose:
        print(f"batch: {batch_key}")

    # set upper bound for k0
    size_max = 2**31 - 1

    # prepare call of kBET per cluster
    kBET_scores = {"cluster": [], "kBET": []}
    for clus in adata_tmp.obs[label_key].unique():

        # subset by label
        adata_sub = adata_tmp[adata_tmp.obs[label_key] == clus, :].copy()

        # check if neighborhood size too small or only one batch in subset
        if np.logical_or(
            adata_sub.n_obs < 10, len(adata_sub.obs[batch_key].cat.categories) == 1
        ):
            print(f"{clus} consists of a single batch or is too small. Skip.")
            score = np.nan
        else:
            quarter_mean = np.floor(
                np.mean(adata_sub.obs[batch_key].value_counts()) / 4
            ).astype("int")
            k0 = np.min([70, np.max([10, quarter_mean])])
            # check k0 for reasonability
            if k0 * adata_sub.n_obs >= size_max:
                k0 = np.floor(size_max / adata_sub.n_obs).astype("int")

            matrix = np.zeros(shape=(adata_sub.n_obs, k0 + 1))

            if verbose:
                print(f"Use {k0} nearest neighbors.")
            n_comp, labs = scipy.sparse.csgraph.connected_components(
                adata_sub.obsp["connectivities"], connection="strong"
            )

            if n_comp == 1:  # a single component to compute kBET on
                try:
                    nn_index_tmp = diffusion_nn(adata_sub, k=k0).astype("float")
                    # call kBET
                    score = kBET_single(
                        matrix=matrix,
                        batch=adata_sub.obs[batch_key],
                        knn=nn_index_tmp
                        + 1,  # nn_index in python is 0-based and 1-based in R
                        verbose=verbose,
                        k0=k0,
                    )
                except NeighborsError:
                    print("Not enough neighbours")
                    score = 1  # i.e. 100% rejection

            else:
                # check the number of components where kBET can be computed upon
                comp_size = pd.value_counts(labs)
                # check which components are small
                comp_size_thresh = 3 * k0
                idx_nonan = np.flatnonzero(
                    np.in1d(labs, comp_size[comp_size >= comp_size_thresh].index)
                )

                # check if 75% of all cells can be used for kBET run
                if len(idx_nonan) / len(labs) >= 0.75:
                    # create another subset of components, assume they are not visited in a diffusion process
                    adata_sub_sub = adata_sub[idx_nonan, :].copy()
                    nn_index_tmp = np.empty(shape=(adata_sub.n_obs, k0))
                    nn_index_tmp[:] = np.nan

                    try:
                        nn_index_tmp[idx_nonan] = diffusion_nn(
                            adata_sub_sub, k=k0
                        ).astype("float")
                        # call kBET
                        score = kBET_single(
                            matrix=matrix,
                            batch=adata_sub.obs[batch_key],
                            knn=nn_index_tmp
                            + 1,  # nn_index in python is 0-based and 1-based in R
                            verbose=verbose,
                            k0=k0,
                        )
                    except NeighborsError:
                        print("Not enough neighbours")
                        score = 1  # i.e. 100% rejection
                else:  # if there are too many too small connected components, set kBET score to 1
                    score = 1  # i.e. 100% rejection

        kBET_scores["cluster"].append(clus)
        kBET_scores["kBET"].append(score)

    kBET_scores = pd.DataFrame.from_dict(kBET_scores)
    kBET_scores = kBET_scores.reset_index(drop=True)

    if return_df:
        return kBET_scores

    final_score = np.nanmean(kBET_scores["kBET"])
    return 1 - final_score if scaled else final_score


def kBET_single(matrix, batch, k0=10, knn=None, verbose=False):
    """Single kBET run

    Compute k-nearest neighbour batch effect test (kBET) score as described in
    https://doi.org/10.1038/s41592-018-0254-1

    :param matrix: expression matrix (at the moment: a PCA matrix, so ``do.pca`` is set to ``FALSE``)
    :param batch: series or list of batch assignments
    :returns: kBET observed rejection rate
    """
    try:
        import anndata2ri
        import rpy2.rinterface_lib.callbacks
        import rpy2.rinterface_lib.embedded
        import rpy2.robjects as ro

        rpy2.rinterface_lib.callbacks.logger.setLevel(logging.ERROR)
    except ModuleNotFoundError as e:
        raise OptionalDependencyNotInstalled(e)

    try:
        ro.r("library(kBET)")
    except rpy2.rinterface_lib.embedded.RRuntimeError as ex:
        RLibraryNotFound(ex)

    anndata2ri.activate()

    if verbose:
        print("importing expression matrix")
    ro.globalenv["data_mtrx"] = matrix
    ro.globalenv["batch"] = batch

    if verbose:
        print("kBET estimation")

    ro.globalenv["knn_graph"] = knn
    ro.globalenv["k0"] = k0
    ro.r(
        "batch.estimate <- kBET("
        "  data_mtrx,"
        "  batch,"
        "  knn=knn_graph,"
        "  k0=k0,"
        "  plot=FALSE,"
        "  do.pca=FALSE,"
        "  heuristic=FALSE,"
        "  adapt=FALSE,"
        f"  verbose={str(verbose).upper()}"
        ")"
    )

    try:
        score = ro.r("batch.estimate$summary$kBET.observed")[0]
    except rpy2.rinterface_lib.embedded.RRuntimeError as ex:
        print(f"Error computing kBET: {ex}\nSetting value to np.nan")
        score = np.nan

    anndata2ri.deactivate()

    return score
