# Path: scib/metrics/nmi.py
# def nmi(adata, group1, group2, method="arithmetic", nmi_dir=None):
#     """
#     Wrapper for normalized mutual information NMI between two different cluster assignments
# 
#     :param adata: Anndata object
#     :param group1: column name of `adata.obs`
#     :param group2: column name of `adata.obs`
#     :param method: NMI implementation
#         'max': scikit method with `average_method='max'`
#         'min': scikit method with `average_method='min'`
#         'geometric': scikit method with `average_method='geometric'`
#         'arithmetic': scikit method with `average_method='arithmetic'`
#         'Lancichinetti': implementation by A. Lancichinetti 2009 et al. https://sites.google.com/site/andrealancichinetti/mutual
#         'ONMI': implementation by Aaron F. McDaid et al. https://github.com/aaronmcdaid/Overlapping-NMI
#     :param nmi_dir: directory of compiled C code if 'Lancichinetti' or 'ONMI' are specified as `method`.
#         These packages need to be compiled as specified in the corresponding READMEs.
#     :return:
#         Normalized mutual information NMI value
#     """
# 
#     check_adata(adata)
#     check_batch(group1, adata.obs)
#     check_batch(group2, adata.obs)
# 
#     group1 = adata.obs[group1].tolist()
#     group2 = adata.obs[group2].tolist()
# 
#     if len(group1) != len(group2):
#         raise ValueError(
#             f'different lengths in group1 ({len(group1)}) and group2 ({len(group2)})'
#         )
# 
#     # choose method
#     if method in ['max', 'min', 'geometric', 'arithmetic']:
#         nmi_value = normalized_mutual_info_score(group1, group2, average_method=method)
#     elif method == "Lancichinetti":
#         nmi_value = nmi_Lanc(group1, group2, nmi_dir=nmi_dir)
#     elif method == "ONMI":
#         nmi_value = onmi(group1, group2, nmi_dir=nmi_dir)
#     else:
#         raise ValueError(f"Method {method} not valid")
# 
#     return nmi_value


import matplotlib.pyplot as plt
import pandas as pd
import scanpy as sc
import seaborn as sns
from .nmi import nmi



def opt_louvain(adata, label_key, cluster_key, function=None, resolutions=None,
                use_rep=None,
                inplace=True, plot=False, force=True, verbose=True, **kwargs):
    """
    params:
        label_key: name of column in adata.obs containing biological labels to be
            optimised against
        cluster_key: name of column to be added to adata.obs during clustering. 
            Will be overwritten if exists and `force=True`
        function: function that computes the cost to be optimised over. Must take as
            arguments (adata, group1, group2, **kwargs) and returns a number for maximising
        resolutions: list if resolutions to be optimised over. If `resolutions=None`,
            default resolutions of 20 values ranging between 0.1 and 2 will be used
        use_rep: key of embedding to use only if adata.uns['neighbors'] is not defined,
            otherwise will be ignored
    returns:
        res_max: resolution of maximum score
        score_max: maximum score
        score_all: `pd.DataFrame` containing all scores at resolutions. Can be used to plot the score profile.
        clustering: only if `inplace=False`, return cluster assignment as `pd.Series`
        plot: if `plot=True` plot the score profile over resolution
    """

    if verbose:
        print('Clustering...')

    if function is None:

# Path: scib/metrics/nmi.py
# def nmi(adata, group1, group2, method="arithmetic", nmi_dir=None):
#     """
#     Wrapper for normalized mutual information NMI between two different cluster assignments
# 
#     :param adata: Anndata object
#     :param group1: column name of `adata.obs`
#     :param group2: column name of `adata.obs`
#     :param method: NMI implementation
#         'max': scikit method with `average_method='max'`
#         'min': scikit method with `average_method='min'`
#         'geometric': scikit method with `average_method='geometric'`
#         'arithmetic': scikit method with `average_method='arithmetic'`
#         'Lancichinetti': implementation by A. Lancichinetti 2009 et al. https://sites.google.com/site/andrealancichinetti/mutual
#         'ONMI': implementation by Aaron F. McDaid et al. https://github.com/aaronmcdaid/Overlapping-NMI
#     :param nmi_dir: directory of compiled C code if 'Lancichinetti' or 'ONMI' are specified as `method`.
#         These packages need to be compiled as specified in the corresponding READMEs.
#     :return:
#         Normalized mutual information NMI value
#     """
# 
#     check_adata(adata)
#     check_batch(group1, adata.obs)
#     check_batch(group2, adata.obs)
# 
#     group1 = adata.obs[group1].tolist()
#     group2 = adata.obs[group2].tolist()
# 
#     if len(group1) != len(group2):
#         raise ValueError(
#             f'different lengths in group1 ({len(group1)}) and group2 ({len(group2)})'
#         )
# 
#     # choose method
#     if method in ['max', 'min', 'geometric', 'arithmetic']:
#         nmi_value = normalized_mutual_info_score(group1, group2, average_method=method)
#     elif method == "Lancichinetti":
#         nmi_value = nmi_Lanc(group1, group2, nmi_dir=nmi_dir)
#     elif method == "ONMI":
#         nmi_value = onmi(group1, group2, nmi_dir=nmi_dir)
#     else:
#         raise ValueError(f"Method {method} not valid")
# 
#     return nmi_value


# Path: scib/metrics/clustering.py
import matplotlib.pyplot as plt
import pandas as pd
import scanpy as sc
import seaborn as sns
from .nmi import nmi




def opt_louvain(adata, label_key, cluster_key, function=None, resolutions=None,
                use_rep=None,
                inplace=True, plot=False, force=True, verbose=True, **kwargs):
    """
    params:
        label_key: name of column in adata.obs containing biological labels to be
            optimised against
        cluster_key: name of column to be added to adata.obs during clustering. 
            Will be overwritten if exists and `force=True`
        function: function that computes the cost to be optimised over. Must take as
            arguments (adata, group1, group2, **kwargs) and returns a number for maximising
        resolutions: list if resolutions to be optimised over. If `resolutions=None`,
            default resolutions of 20 values ranging between 0.1 and 2 will be used
        use_rep: key of embedding to use only if adata.uns['neighbors'] is not defined,
            otherwise will be ignored
    returns:
        res_max: resolution of maximum score
        score_max: maximum score
        score_all: `pd.DataFrame` containing all scores at resolutions. Can be used to plot the score profile.
        clustering: only if `inplace=False`, return cluster assignment as `pd.Series`
        plot: if `plot=True` plot the score profile over resolution
    """

    if verbose:
        print('Clustering...')

    if function is None:

        function = nmi