"""
Tests for the birch clustering algorithm taken from
  sklearn/cluster/tests/test_birch.py

patched with scikit-learn/scikit-learn#8808

Ensures that the patched version of Birch included in
FreeDiscovery version is still compatible with upstream.
"""

from scipy import sparse
import numpy as np

from sklearn.cluster.tests.common import generate_clustered_data
from freediscovery.cluster.birch import Birch
from sklearn.cluster.hierarchical import AgglomerativeClustering
from sklearn.datasets import make_blobs
from sklearn.linear_model import ElasticNet
from sklearn.metrics import pairwise_distances_argmin, v_measure_score

from sklearn.utils.testing import assert_greater_equal
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_greater
from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_warns


def test_n_samples_leaves_roots():
    # Sanity check for the number of samples in leaves and roots
    X, y = make_blobs(n_samples=10)
    brc = Birch()
    brc.fit(X)
    n_samples_root = sum([sc.n_samples_ for sc in brc.root_.subclusters_])
    n_samples_leaves = sum([sc.n_samples_ for leaf in brc._get_leaves()
                            for sc in leaf.subclusters_])
    assert_equal(n_samples_leaves, X.shape[0])
    assert_equal(n_samples_root, X.shape[0])


def test_birch_predict():
    # Test the predict method predicts the nearest centroid.
    rng = np.random.RandomState(0)
    X = generate_clustered_data(n_clusters=3, n_features=3,
                                n_samples_per_cluster=10)

    # n_samples * n_samples_per_cluster
    shuffle_indices = np.arange(30)
    rng.shuffle(shuffle_indices)
    X_shuffle = X[shuffle_indices, :]
    brc = Birch(n_clusters=4, threshold=1.)
    brc.fit(X_shuffle)
    centroids = brc.subcluster_centers_
    assert_array_equal(brc.labels_, brc.predict(X_shuffle))
    nearest_centroid = pairwise_distances_argmin(X_shuffle, centroids)
    assert_almost_equal(v_measure_score(nearest_centroid, brc.labels_), 1.0)


def test_n_clusters():
    # Test that n_clusters param works properly
    X, y = make_blobs(n_samples=100, centers=10)
    brc1 = Birch(n_clusters=10)
    brc1.fit(X)
    assert_greater(len(brc1.subcluster_centers_), 10)
    assert_equal(len(np.unique(brc1.labels_)), 10)

    # Test that n_clusters = Agglomerative Clustering gives
    # the same results.
    gc = AgglomerativeClustering(n_clusters=10)
    brc2 = Birch(n_clusters=gc)
    brc2.fit(X)
    assert_array_equal(brc1.subcluster_labels_, brc2.subcluster_labels_)
    assert_array_equal(brc1.labels_, brc2.labels_)

    # Test that the wrong global clustering step raises an Error.
    clf = ElasticNet()
    brc3 = Birch(n_clusters=clf)
    assert_raises(ValueError, brc3.fit, X)

    # Test that a small number of clusters raises a warning.
    brc4 = Birch(threshold=10000.)
    assert_warns(UserWarning, brc4.fit, X)


def test_sparse_X():
    # Test that sparse and dense data give same results
    X, y = make_blobs(n_samples=100, centers=10)
    brc = Birch(n_clusters=10)
    brc.fit(X)

    csr = sparse.csr_matrix(X)
    brc_sparse = Birch(n_clusters=10)
    brc_sparse.fit(csr)

    assert_array_equal(brc.labels_, brc_sparse.labels_)
    assert_array_almost_equal(brc.subcluster_centers_,
                              brc_sparse.subcluster_centers_)


def check_branching_factor(node, branching_factor):
    subclusters = node.subclusters_
    assert_greater_equal(branching_factor, len(subclusters))
    for cluster in subclusters:
        if cluster.child_:
            check_branching_factor(cluster.child_, branching_factor)


def test_branching_factor():
    # Test that nodes have at max branching_factor number of subclusters
    X, y = make_blobs()
    branching_factor = 9

    # Purposefully set a low threshold to maximize the subclusters.
    brc = Birch(n_clusters=None, branching_factor=branching_factor,
                threshold=0.01)
    brc.fit(X)
    check_branching_factor(brc.root_, branching_factor)
    brc = Birch(n_clusters=3, branching_factor=branching_factor,
                threshold=0.01)
    brc.fit(X)
    check_branching_factor(brc.root_, branching_factor)

    # Raises error when branching_factor is set to one.
    brc = Birch(n_clusters=None, branching_factor=1, threshold=0.01)
    assert_raises(ValueError, brc.fit, X)


def check_threshold(birch_instance, threshold):
    """Use the leaf linked list for traversal"""
    current_leaf = birch_instance.dummy_leaf_.next_leaf_
    while current_leaf:
        subclusters = current_leaf.subclusters_
        for sc in subclusters:
            assert_greater_equal(threshold, sc.radius)
        current_leaf = current_leaf.next_leaf_


def test_threshold():
    # Test that the leaf subclusters have a threshold lesser than radius
    X, y = make_blobs(n_samples=80, centers=4)
    brc = Birch(threshold=0.5, n_clusters=None)
    brc.fit(X)
    check_threshold(brc, 0.5)

    brc = Birch(threshold=5.0, n_clusters=None)
    brc.fit(X)
    check_threshold(brc, 5.)


def test_birch_hierarchy():
    X, y = make_blobs(random_state=40)
    brc = Birch(n_clusters=None, branching_factor=5,
                compute_sample_indices=True)
    brc.fit(X)

    # make sure that leave nodes contain all the samples
    n_leaves = 1
    sample_id = []
    current_leaf = brc.dummy_leaf_.next_leaf_
    while current_leaf:
        subclusters = current_leaf.subclusters_
        for sc in subclusters:
            assert sc.n_samples_ == len(sc.samples_id_)
            sample_id += sc.samples_id_
        current_leaf = current_leaf.next_leaf_
        n_leaves += 1
    assert_array_equal(np.sort(sample_id), np.arange(X.shape[0]))

    # Verify that the resulting hierarchical tree is deeper than 1 level
    # (i.e. subclusters of the root node are nor tree leaves )
    assert len(brc.root_.subclusters_) < n_leaves

    # Make sure that subclusters of the root_ node contain all the samples
    sample_id = []
    for sc in brc.root_.subclusters_:
        sample_id += sc.samples_id_
        assert sc.n_samples_ == len(sc.samples_id_)
    assert_array_equal(np.sort(sample_id), np.arange(X.shape[0]))

    # Pick a sample at random and make sure that reported samples_id_
    # matches with the subcluster the sample is closest to
    document_id = 45
    document_in_subcluster = []
    distance_to_centroid = []
    for sc in brc.root_.subclusters_:
        centroid = X[sc.samples_id_, :].mean(axis=0)
        distance_to_centroid.append(((X[[document_id]] - centroid)**2).sum())
        document_in_subcluster.append(document_id in sc.samples_id_)

    assert np.argmin(distance_to_centroid) == \
        np.nonzero(document_in_subcluster)[0][0]

    # Make sure that we can recompute labels from tree leaves
    labels2 = np.zeros(X.shape[0], dtype=int)
    cluster_id = 0
    for current_leaf in brc._get_leaves():
        subclusters = current_leaf.subclusters_
        for sc in subclusters:
            labels2[list(sc.samples_id_)] = cluster_id
            cluster_id += 1

    assert np.unique(brc.labels_).shape == np.unique(labels2).shape
    # The two methods yield approximately equal labels
    assert v_measure_score(brc.labels_, labels2) > 0.95
