"""
Contains functions to generate and combine a clustering ensemble.
"""
import numpy as np
import pandas as pd
from sklearn.metrics import pairwise_distances
from sklearn.metrics import adjusted_rand_score as ari
from sklearn.metrics import adjusted_mutual_info_score as ami
from sklearn.metrics import normalized_mutual_info_score as nmi
from tqdm import tqdm

from clustering.utils import reset_estimator, compare_arrays


def generate_ensemble(data, clusterers: dict, attributes: list, affinity_matrix=None):
    """
    It generates an ensemble from the data given a set of clusterers (a
    clusterer is an instance of a clustering algorithm with a fixed set of
    parameters).

    Args:
        data:
            A numpy array, pandas dataframe, or any other structure supported
            by the clusterers as data input.
        clusterers:
            A dictionary with clusterers specified in this format: { 'k-means
            #1': KMeans(n_clusters=2), ... }
        attributes:
            A list of attributes to save in the final dataframe; for example,
            including "n_clusters" will extract this attribute from the
            estimator and include it in the final dataframe returned.
        affinity_matrix:
            If the clustering algorithm is AgglomerativeClustering (from
            sklearn) and the linkage method is different than ward (which only
            support euclidean distance), the affinity_matrix is given as data
            input to the estimator instead of data.

    Returns:
        A pandas DataFrame with all the partitions generated by the clusterers.
        Columns include the clusterer name/id, the partition, the estimator
        parameters (obtained with the get_params() method) and any other
        attribute specified.
    """
    ensemble = []

    for clus_name, clus_obj in tqdm(clusterers.items(), total=len(clusterers)):
        # get partition
        #
        # for agglomerative clustering both data and affinity_matrix should be
        # given; for ward linkage, data is used, and for the other linkage
        # methods the affinity_matrix is used
        if (type(clus_obj).__name__ == "AgglomerativeClustering") and (
            clus_obj.linkage != "ward"
        ):
            partition = clus_obj.fit_predict(affinity_matrix).astype(float)
        else:
            partition = clus_obj.fit_predict(data).astype(float)

        # remove from partition noisy points (for example, if using DBSCAN)
        partition[partition < 0] = np.nan

        # get number of clusters
        partition_no_nan = partition[~np.isnan(partition)]
        n_clusters = np.unique(partition_no_nan).shape[0]

        # stop if n_clusters <= 1
        if n_clusters <= 1:
            reset_estimator(clus_obj)
            continue

        res = pd.Series(
            {
                "clusterer_id": clus_name,
                "clusterer_params": str(clus_obj.get_params()),
                "partition": partition,
            }
        )

        for attr in attributes:
            if attr == "n_clusters" and not hasattr(clus_obj, attr):
                res[attr] = n_clusters
            else:
                res[attr] = getattr(clus_obj, attr)

        ensemble.append(res)

        # for some estimators such as DBSCAN this is needed, because otherwise
        # the estimator saves references of huge data structures not needed in
        # this context
        reset_estimator(clus_obj)

    return pd.DataFrame(ensemble).set_index("clusterer_id")


def get_ensemble_distance_matrix(ensemble, n_jobs=1):
    """
    Given an ensemble, it computes the coassociation matrix (a distance matrix
    for all objects using the ensemble information). For each object pair, the
    coassociation matrix contains the percentage of times the pair of objects
    was clustered together in the ensemble.

    Args:
        ensemble:
            A numpy array representing a set of clustering solutions on the same
            data. Each row is a clustering solution (partition) and columns are
            objects.
        n_jobs:
            The number of jobs used by the pairwise_distance matrix from
            sklearn.

    Returns:
        A numpy array representing a square distance matrix for all objects
        (coassociation matrix).
    """

    def _compare(x, y):
        xy = np.array([x, y]).T
        xy = xy[~np.isnan(xy).any(axis=1)]
        return (xy[:, 0] != xy[:, 1]).sum() / xy.shape[0]

    return pairwise_distances(
        ensemble.T, metric=_compare, n_jobs=n_jobs, force_all_finite="allow-nan"
    )


def supraconsensus(ensemble, k, methods, selection_criterion, n_jobs=1, use_tqdm=False):
    """
    It combines a clustering ensemble using a set of methods that the user can
    specify. Each of these methods combines the ensemble and returns a single
    partition. This function returns the combined partition that maximizes the
    selection criterion.

    Args:
        ensemble:
            a clustering ensemble (rows are partitions, columns are objects).
        k:
            the final number of clusters for the combined partition.
        methods:
            a list of methods to apply on the ensemble; each returns a combined
            partition.
        selection_criterion:
            a function that represents the selection criterion; this function
            has to accept an ensemble as the first argument, and a partition as
            the second one.
        n_jobs:
            number of jobs.
        use_tqdm:
            ensembles/disables the use of tqdm to show a progress bar.

    Returns:
        Returns a tuple: (partition, best method name, best criterion value)
    """
    from concurrent.futures import ProcessPoolExecutor, as_completed

    methods_results = {}

    with ProcessPoolExecutor(max_workers=n_jobs) as executor:
        tasks = {executor.submit(m, ensemble, k): m.__name__ for m in methods}

        for future in tqdm(
            as_completed(tasks),
            total=len(tasks),
            disable=(not use_tqdm),
            ncols=100,
        ):
            method_name = tasks[future]
            part = future.result()
            criterion_value = selection_criterion(ensemble, part)

            methods_results[method_name] = {
                "partition": part,
                "criterion_value": criterion_value,
            }

    # select the best performing method according to the selection criterion
    best_method = max(
        methods_results, key=lambda x: methods_results[x]["criterion_value"]
    )
    best_method_results = methods_results[best_method]

    return (
        best_method_results["partition"],
        best_method,
        best_method_results["criterion_value"],
    )


def run_method_and_compute_agreement(method_func, ensemble_data, ensemble, k, **kwargs):
    """
    Runs a consensus clustering method on the ensemble data, obtains the
    consolidated partition with the desired number of clusters, and computes
    a series of performance measures.

    Args:
        method_func:
            A consensus function (first argument is either the ensemble or
            the coassociation matrix derived from the ensemble).
        ensemble_data:
            A numpy array with the ensemble data that will be given to the
            specified method. For evidence accumulation methods, this is the
            coassociation matrix (a square matrix with the distance between
            object pairs derived from the ensemble).
        ensemble:
            A numpy array representing the ensemble (partitions in rows, objects
            in columns).
        k:
            The number of clusters to obtain from the ensemble data using the
            specified method.
        kwargs:
            Other parameters passed to `method_func`.

    Returns:
        It returns a tuple with the data partition derived from the ensemble
        data using the specified method, and some performance measures of this
        partition.
    """
    part = method_func(ensemble_data, k, **kwargs)

    nmi_values = np.array(
        [
            compare_arrays(ensemble_member, part, nmi, use_weighting=True)
            for ensemble_member in ensemble
        ]
    )

    ami_values = np.array(
        [
            compare_arrays(ensemble_member, part, ami, use_weighting=True)
            for ensemble_member in ensemble
        ]
    )

    ari_values = np.array(
        [
            compare_arrays(ensemble_member, part, ari, use_weighting=True)
            for ensemble_member in ensemble
        ]
    )

    performance_values = {
        "ari_mean": np.mean(ari_values),
        "ari_median": np.median(ari_values),
        "ari_std": np.std(ari_values),
        "ami_mean": np.mean(ami_values),
        "ami_median": np.median(ami_values),
        "ami_std": np.std(ami_values),
        "nmi_mean": np.mean(nmi_values),
        "nmi_median": np.median(nmi_values),
        "nmi_std": np.std(nmi_values),
    }

    return part, performance_values

[((121, 11, 123, 5), 'sklearn.metrics.pairwise_distances', 'pairwise_distances', (), '', False, 'from sklearn.metrics import pairwise_distances\n'), ((90, 8, 90, 33), 'clustering.utils.reset_estimator', 'reset_estimator', ({(90, 24, 90, 32): 'clus_obj'}, {}), '(clus_obj)', False, 'from clustering.utils import reset_estimator, compare_arrays\n'), ((157, 9, 157, 48), 'concurrent.futures.ProcessPoolExecutor', 'ProcessPoolExecutor', (), '', False, 'from concurrent.futures import ProcessPoolExecutor, as_completed\n'), ((241, 20, 241, 39), 'numpy.mean', 'np.mean', ({(241, 28, 241, 38): 'ari_values'}, {}), '(ari_values)', True, 'import numpy as np\n'), ((242, 22, 242, 43), 'numpy.median', 'np.median', ({(242, 32, 242, 42): 'ari_values'}, {}), '(ari_values)', True, 'import numpy as np\n'), ((243, 19, 243, 37), 'numpy.std', 'np.std', ({(243, 26, 243, 36): 'ari_values'}, {}), '(ari_values)', True, 'import numpy as np\n'), ((244, 20, 244, 39), 'numpy.mean', 'np.mean', ({(244, 28, 244, 38): 'ami_values'}, {}), '(ami_values)', True, 'import numpy as np\n'), ((245, 22, 245, 43), 'numpy.median', 'np.median', ({(245, 32, 245, 42): 'ami_values'}, {}), '(ami_values)', True, 'import numpy as np\n'), ((246, 19, 246, 37), 'numpy.std', 'np.std', ({(246, 26, 246, 36): 'ami_values'}, {}), '(ami_values)', True, 'import numpy as np\n'), ((247, 20, 247, 39), 'numpy.mean', 'np.mean', ({(247, 28, 247, 38): 'nmi_values'}, {}), '(nmi_values)', True, 'import numpy as np\n'), ((248, 22, 248, 43), 'numpy.median', 'np.median', ({(248, 32, 248, 42): 'nmi_values'}, {}), '(nmi_values)', True, 'import numpy as np\n'), ((249, 19, 249, 37), 'numpy.std', 'np.std', ({(249, 26, 249, 36): 'nmi_values'}, {}), '(nmi_values)', True, 'import numpy as np\n'), ((68, 12, 68, 37), 'clustering.utils.reset_estimator', 'reset_estimator', ({(68, 28, 68, 36): 'clus_obj'}, {}), '(clus_obj)', False, 'from clustering.utils import reset_estimator, compare_arrays\n'), ((92, 11, 92, 33), 'pandas.DataFrame', 'pd.DataFrame', ({(92, 24, 92, 32): 'ensemble'}, {}), '(ensemble)', True, 'import pandas as pd\n'), ((117, 13, 117, 29), 'numpy.array', 'np.array', ({(117, 22, 117, 28): '[x, y]'}, {}), '([x, y])', True, 'import numpy as np\n'), ((161, 12, 161, 31), 'concurrent.futures.as_completed', 'as_completed', ({(161, 25, 161, 30): 'tasks'}, {}), '(tasks)', False, 'from concurrent.futures import ProcessPoolExecutor, as_completed\n'), ((221, 12, 221, 74), 'clustering.utils.compare_arrays', 'compare_arrays', (), '', False, 'from clustering.utils import reset_estimator, compare_arrays\n'), ((228, 12, 228, 74), 'clustering.utils.compare_arrays', 'compare_arrays', (), '', False, 'from clustering.utils import reset_estimator, compare_arrays\n'), ((235, 12, 235, 74), 'clustering.utils.compare_arrays', 'compare_arrays', (), '', False, 'from clustering.utils import reset_estimator, compare_arrays\n'), ((64, 21, 64, 48), 'numpy.unique', 'np.unique', ({(64, 31, 64, 47): 'partition_no_nan'}, {}), '(partition_no_nan)', True, 'import numpy as np\n'), ((63, 38, 63, 57), 'numpy.isnan', 'np.isnan', ({(63, 47, 63, 56): 'partition'}, {}), '(partition)', True, 'import numpy as np\n'), ((118, 17, 118, 29), 'numpy.isnan', 'np.isnan', ({(118, 26, 118, 28): 'xy'}, {}), '(xy)', True, 'import numpy as np\n')]