"""Testing for K-means"""
import sys

import numpy as np
from scipy import sparse as sp

from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import SkipTest
from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_raises_regexp
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_greater
from sklearn.utils.testing import assert_less
from sklearn.utils.testing import assert_warns
from sklearn.utils.testing import if_safe_multiprocessing_with_blas

from sklearn.utils.validation import DataConversionWarning
from sklearn.utils.extmath import row_norms
from sklearn.metrics.cluster import v_measure_score
from sklearn.cluster import KMeans, k_means
from sklearn.cluster import MiniBatchKMeans
from sklearn.cluster.k_means_ import _labels_inertia
from sklearn.cluster.k_means_ import _mini_batch_step
from sklearn.datasets.samples_generator import make_blobs
from sklearn.externals.six.moves import cStringIO as StringIO

# non centered, sparse centers to check the
centers = np.array([
    [0.0, 5.0, 0.0, 0.0, 0.0],
    [1.0, 1.0, 4.0, 0.0, 0.0],
    [1.0, 0.0, 0.0, 5.0, 1.0],
])
n_samples = 100
n_clusters, n_features = centers.shape
X, true_labels = make_blobs(n_samples=n_samples, centers=centers,
                            cluster_std=1., random_state=42)
X_csr = sp.csr_matrix(X)

def test_kmeans_dtype():
    rnd = np.random.RandomState(0)
    X = rnd.normal(size=(40, 2))
    X = (X * 10).astype(np.uint8)
    km = KMeans(n_init=1).fit(X)
    pred_x = assert_warns(DataConversionWarning, km.predict, X)
    assert_array_equal(km.labels_, pred_x)

def test_labels_assignment_and_inertia():
    # pure numpy implementation as easily auditable reference gold
    # implementation
    rng = np.random.RandomState(42)
    noisy_centers = centers + rng.normal(size=centers.shape)
    labels_gold = - np.ones(n_samples, dtype=np.int)
    mindist = np.empty(n_samples)
    mindist.fill(np.infty)
    for center_id in range(n_clusters):
        dist = np.sum((X - noisy_centers[center_id]) ** 2, axis=1)
        labels_gold[dist < mindist] = center_id
        mindist = np.minimum(dist, mindist)
    inertia_gold = mindist.sum()
    assert_true((mindist >= 0.0).all())
    assert_true((labels_gold != -1).all())

    # perform label assignment using the dense array input
    x_squared_norms = (X ** 2).sum(axis=1)
    labels_array, inertia_array = _labels_inertia(
        X, x_squared_norms, noisy_centers)
    assert_array_almost_equal(inertia_array, inertia_gold)
    assert_array_equal(labels_array, labels_gold)

    # perform label assignment using the sparse CSR input
    x_squared_norms_from_csr = row_norms(X_csr, squared=True)
    labels_csr, inertia_csr = _labels_inertia(
        X_csr, x_squared_norms_from_csr, noisy_centers)
    assert_array_almost_equal(inertia_csr, inertia_gold)
    assert_array_equal(labels_csr, labels_gold)

def test_minibatch_update_consistency():
    # Check that dense and sparse minibatch update give the same results
    rng = np.random.RandomState(42)
    old_centers = centers + rng.normal(size=centers.shape)

    new_centers = old_centers.copy()
    new_centers_csr = old_centers.copy()

    counts = np.zeros(new_centers.shape[0], dtype=np.int32)
    counts_csr = np.zeros(new_centers.shape[0], dtype=np.int32)

    x_squared_norms = (X ** 2).sum(axis=1)
    x_squared_norms_csr = row_norms(X_csr, squared=True)

    buffer = np.zeros(centers.shape[1], dtype=np.double)
    buffer_csr = np.zeros(centers.shape[1], dtype=np.double)

    # extract a small minibatch
    X_mb = X[:10]
    X_mb_csr = X_csr[:10]
    x_mb_squared_norms = x_squared_norms[:10]
    x_mb_squared_norms_csr = x_squared_norms_csr[:10]

    # step 1: compute the dense minibatch update
    old_inertia, incremental_diff = _mini_batch_step(
        X_mb, x_mb_squared_norms, new_centers, counts,
        buffer, 1, None, random_reassign=False)
    assert_greater(old_inertia, 0.0)

    # compute the new inertia on the same batch to check that it decreased
    labels, new_inertia = _labels_inertia(
        X_mb, x_mb_squared_norms, new_centers)
    assert_greater(new_inertia, 0.0)
    assert_less(new_inertia, old_inertia)

    # check that the incremental difference computation is matching the
    # final observed value
    effective_diff = np.sum((new_centers - old_centers) ** 2)
    assert_almost_equal(incremental_diff, effective_diff)

    # step 2: compute the sparse minibatch update
    old_inertia_csr, incremental_diff_csr = _mini_batch_step(
        X_mb_csr, x_mb_squared_norms_csr, new_centers_csr, counts_csr,
        buffer_csr, 1, None, random_reassign=False)
    assert_greater(old_inertia_csr, 0.0)

    # compute the new inertia on the same batch to check that it decreased
    labels_csr, new_inertia_csr = _labels_inertia(
        X_mb_csr, x_mb_squared_norms_csr, new_centers_csr)
    assert_greater(new_inertia_csr, 0.0)
    assert_less(new_inertia_csr, old_inertia_csr)

    # check that the incremental difference computation is matching the
    # final observed value
    effective_diff = np.sum((new_centers_csr - old_centers) ** 2)
    assert_almost_equal(incremental_diff_csr, effective_diff)

    # step 3: check that sparse and dense updates lead to the same results
    assert_array_equal(labels, labels_csr)
    assert_array_almost_equal(new_centers, new_centers_csr)
    assert_almost_equal(incremental_diff, incremental_diff_csr)
    assert_almost_equal(old_inertia, old_inertia_csr)
    assert_almost_equal(new_inertia, new_inertia_csr)

def _check_fitted_model(km):
    # check that the number of clusters centers and distinct labels match
    # the expectation
    centers = km.cluster_centers_
    assert_equal(centers.shape, (n_clusters, n_features))

    labels = km.labels_
    assert_equal(np.unique(labels).shape[0], n_clusters)

    # check that the labels assignment are perfect (up to a permutation)
    assert_equal(v_measure_score(true_labels, labels), 1.0)
    assert_greater(km.inertia_, 0.0)

    # check error on dataset being too small
    assert_raises(ValueError, km.fit, [[0., 1.]])

def test_k_means_plus_plus_init():
    km = KMeans(init="k-means++", n_clusters=n_clusters,
                random_state=42).fit(X)
    _check_fitted_model(km)

def test_k_means_new_centers():
    # Explore the part of the code where a new center is reassigned
    X = np.array([[0, 0, 1, 1],
                  [0, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 0, 0],
                  [0, 0, 0, 0],
                  [0, 1, 0, 0]])
    labels = [0, 1, 2, 1, 1, 2]
    bad_centers = np.array([[+0, 1, 0, 0],
                            [.2, 0, .2, .2],
                            [+0, 0, 0, 0]])

    km = KMeans(n_clusters=3, init=bad_centers, n_init=1, max_iter=10,
                random_state=1)
    for this_X in (X, sp.coo_matrix(X)):
        km.fit(this_X)
        this_labels = km.labels_
        # Reorder the labels so that the first instance is in cluster 0,
        # the second in cluster 1, ...
        this_labels = np.unique(this_labels, return_index=True)[1][this_labels]
        np.testing.assert_array_equal(this_labels, labels)

@if_safe_multiprocessing_with_blas
def test_k_means_plus_plus_init_2_jobs():
    if sys.version_info[:2] < (3, 4):
        raise SkipTest(
            "Possible multi-process bug with some BLAS under Python < 3.4")

    km = KMeans(init="k-means++", n_clusters=n_clusters, n_jobs=2,
                random_state=42).fit(X)
    _check_fitted_model(km)

def test_k_means_precompute_distances_flag():
    # check that a warning is raised if the precompute_distances flag is not
    # supported
    km = KMeans(precompute_distances="wrong")
    assert_raises(ValueError, km.fit, X)

def test_k_means_plus_plus_init_sparse():
    km = KMeans(init="k-means++", n_clusters=n_clusters, random_state=42)
    km.fit(X_csr)
    _check_fitted_model(km)

def test_k_means_random_init():
    km = KMeans(init="random", n_clusters=n_clusters, random_state=42)
    km.fit(X)
    _check_fitted_model(km)

def test_k_means_random_init_sparse():
    km = KMeans(init="random", n_clusters=n_clusters, random_state=42)
    km.fit(X_csr)
    _check_fitted_model(km)

def test_k_means_plus_plus_init_not_precomputed():
    km = KMeans(init="k-means++", n_clusters=n_clusters, random_state=42,
                precompute_distances=False).fit(X)
    _check_fitted_model(km)

def test_k_means_random_init_not_precomputed():
    km = KMeans(init="random", n_clusters=n_clusters, random_state=42,
                precompute_distances=False).fit(X)
    _check_fitted_model(km)

def test_k_means_perfect_init():
    km = KMeans(init=centers.copy(), n_clusters=n_clusters, random_state=42,
                n_init=1)
    km.fit(X)
    _check_fitted_model(km)

def test_k_means_n_init():
    rnd = np.random.RandomState(0)
    X = rnd.normal(size=(40, 2))

    # two regression tests on bad n_init argument
    # previous bug: n_init <= 0 threw non-informative TypeError (#3858)
    assert_raises_regexp(ValueError, "n_init", KMeans(n_init=0).fit, X)
    assert_raises_regexp(ValueError, "n_init", KMeans(n_init=-1).fit, X)

def test_k_means_fortran_aligned_data():
    # Check the KMeans will work well, even if X is a fortran-aligned data.
    X = np.asfortranarray([[0, 0], [0, 1], [0, 1]])
    centers = np.array([[0, 0], [0, 1]])
    labels = np.array([0, 1, 1])
    km = KMeans(n_init=1, init=centers, precompute_distances=False,
                random_state=42)
    km.fit(X)
    assert_array_equal(km.cluster_centers_, centers)
    assert_array_equal(km.labels_, labels)

def test_mb_k_means_plus_plus_init_dense_array():
    mb_k_means = MiniBatchKMeans(init="k-means++", n_clusters=n_clusters,
                                 random_state=42)
    mb_k_means.fit(X)
    _check_fitted_model(mb_k_means)

def test_mb_kmeans_verbose():
    mb_k_means = MiniBatchKMeans(init="k-means++", n_clusters=n_clusters,
                                 random_state=42, verbose=1)
    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        mb_k_means.fit(X)
    finally:
        sys.stdout = old_stdout

def test_mb_k_means_plus_plus_init_sparse_matrix():
    mb_k_means = MiniBatchKMeans(init="k-means++", n_clusters=n_clusters,
                                 random_state=42)
    mb_k_means.fit(X_csr)
    _check_fitted_model(mb_k_means)

def test_minibatch_init_with_large_k():
    mb_k_means = MiniBatchKMeans(init='k-means++', init_size=10, n_clusters=20)
    # Check that a warning is raised, as the number clusters is larger
    # than the init_size
    assert_warns(RuntimeWarning, mb_k_means.fit, X)

def test_minibatch_k_means_random_init_dense_array():
    # increase n_init to make random init stable enough
    mb_k_means = MiniBatchKMeans(init="random", n_clusters=n_clusters,
                                 random_state=42, n_init=10).fit(X)
    _check_fitted_model(mb_k_means)

def test_minibatch_k_means_random_init_sparse_csr():
    # increase n_init to make random init stable enough
    mb_k_means = MiniBatchKMeans(init="random", n_clusters=n_clusters,
                                 random_state=42, n_init=10).fit(X_csr)
    _check_fitted_model(mb_k_means)

def test_minibatch_k_means_perfect_init_dense_array():
    mb_k_means = MiniBatchKMeans(init=centers.copy(), n_clusters=n_clusters,
                                 random_state=42, n_init=1).fit(X)
    _check_fitted_model(mb_k_means)

def test_minibatch_k_means_init_multiple_runs_with_explicit_centers():
    mb_k_means = MiniBatchKMeans(init=centers.copy(), n_clusters=n_clusters,
                                 random_state=42, n_init=10)
    assert_warns(RuntimeWarning, mb_k_means.fit, X)

def test_minibatch_k_means_perfect_init_sparse_csr():
    mb_k_means = MiniBatchKMeans(init=centers.copy(), n_clusters=n_clusters,
                                 random_state=42, n_init=1).fit(X_csr)
    _check_fitted_model(mb_k_means)

def test_minibatch_sensible_reassign_fit():
    # check if identical initial clusters are reassigned
    # also a regression test for when there are more desired reassignments than
    # samples.
    zeroed_X, true_labels = make_blobs(n_samples=100, centers=5,
                                       cluster_std=1., random_state=42)
    zeroed_X[::2, :] = 0
    mb_k_means = MiniBatchKMeans(n_clusters=20, batch_size=10, random_state=42,
                                 init="random")
    mb_k_means.fit(zeroed_X)
    # there should not be too many exact zero cluster centers
    assert_greater(mb_k_means.cluster_centers_.any(axis=1).sum(), 10)

    # do the same with batch-size > X.shape[0] (regression test)
    mb_k_means = MiniBatchKMeans(n_clusters=20, batch_size=201,
                                 random_state=42, init="random")
    mb_k_means.fit(zeroed_X)
    # there should not be too many exact zero cluster centers
    assert_greater(mb_k_means.cluster_centers_.any(axis=1).sum(), 10)

def test_minibatch_sensible_reassign_partial_fit():
    zeroed_X, true_labels = make_blobs(n_samples=n_samples, centers=5,
                                       cluster_std=1., random_state=42)
    zeroed_X[::2, :] = 0
    mb_k_means = MiniBatchKMeans(n_clusters=20, random_state=42, init="random")
    for i in range(100):
        mb_k_means.partial_fit(zeroed_X)
    # there should not be too many exact zero cluster centers
    assert_greater(mb_k_means.cluster_centers_.any(axis=1).sum(), 10)

def test_minibatch_reassign():
    # Give a perfect initialization, but a large reassignment_ratio,
    # as a result all the centers should be reassigned and the model
    # should not longer be good
    for this_X in (X, X_csr):
        mb_k_means = MiniBatchKMeans(n_clusters=n_clusters, batch_size=100,
                                     random_state=42)
        mb_k_means.fit(this_X)

        score_before = mb_k_means.score(this_X)
        try:
            old_stdout = sys.stdout
            sys.stdout = StringIO()
            # Turn on verbosity to smoke test the display code
            _mini_batch_step(this_X, (X ** 2).sum(axis=1),
                             mb_k_means.cluster_centers_,
                             mb_k_means.counts_,
                             np.zeros(X.shape[1], np.double),
                             False, distances=np.zeros(X.shape[0]),
                             random_reassign=True, random_state=42,
                             reassignment_ratio=1, verbose=True)
        finally:
            sys.stdout = old_stdout
        assert_greater(score_before, mb_k_means.score(this_X))

    # Give a perfect initialization, with a small reassignment_ratio,
    # no center should be reassigned
    for this_X in (X, X_csr):
        mb_k_means = MiniBatchKMeans(n_clusters=n_clusters, batch_size=100,
                                     init=centers.copy(),
                                     random_state=42, n_init=1)
        mb_k_means.fit(this_X)
        clusters_before = mb_k_means.cluster_centers_
        # Turn on verbosity to smoke test the display code
        _mini_batch_step(this_X, (X ** 2).sum(axis=1),
                         mb_k_means.cluster_centers_,
                         mb_k_means.counts_,
                         np.zeros(X.shape[1], np.double),
                         False, distances=np.zeros(X.shape[0]),
                         random_reassign=True, random_state=42,
                         reassignment_ratio=1e-15)
        assert_array_almost_equal(clusters_before, mb_k_means.cluster_centers_)

def test_minibatch_with_many_reassignments():
    # Test for the case that the number of clusters to reassign is bigger
    # than the batch_size
    n_samples = 550
    rnd = np.random.RandomState(42)
    X = rnd.uniform(size=(n_samples, 10))
    # Check that the fit works if n_clusters is bigger than the batch_size.
    # Run the test with 550 clusters and 550 samples, because it turned out
    # that this values ensure that the number of clusters to reassign
    # is always bigger than the batch_size
    n_clusters = 550
    MiniBatchKMeans(n_clusters=n_clusters,
                    batch_size=100,
                    init_size=n_samples,
                    random_state=42).fit(X)

def test_sparse_mb_k_means_callable_init():

    def test_init(X, k, random_state):
        return centers

    # Small test to check that giving the wrong number of centers
    # raises a meaningful error
    assert_raises(ValueError,
                  MiniBatchKMeans(init=test_init, random_state=42).fit, X_csr)

    # Now check that the fit actually works
    mb_k_means = MiniBatchKMeans(n_clusters=3, init=test_init,
                                 random_state=42).fit(X_csr)
    _check_fitted_model(mb_k_means)

def test_mini_batch_k_means_random_init_partial_fit():
    km = MiniBatchKMeans(n_clusters=n_clusters, init="random", random_state=42)

    # use the partial_fit API for online learning
    for X_minibatch in np.array_split(X, 10):
        km.partial_fit(X_minibatch)

    # compute the labeling on the complete dataset
    labels = km.predict(X)
    assert_equal(v_measure_score(true_labels, labels), 1.0)

def test_minibatch_default_init_size():
    mb_k_means = MiniBatchKMeans(init=centers.copy(), n_clusters=n_clusters,
                                 batch_size=10, random_state=42,
                                 n_init=1).fit(X)
    assert_equal(mb_k_means.init_size_, 3 * mb_k_means.batch_size)
    _check_fitted_model(mb_k_means)

def test_minibatch_tol():
    mb_k_means = MiniBatchKMeans(n_clusters=n_clusters, batch_size=10,
                                 random_state=42, tol=.01).fit(X)
    _check_fitted_model(mb_k_means)

def test_minibatch_set_init_size():
    mb_k_means = MiniBatchKMeans(init=centers.copy(), n_clusters=n_clusters,
                                 init_size=666, random_state=42,
                                 n_init=1).fit(X)
    assert_equal(mb_k_means.init_size, 666)
    assert_equal(mb_k_means.init_size_, n_samples)
    _check_fitted_model(mb_k_means)

def test_k_means_invalid_init():
    km = KMeans(init="invalid", n_init=1, n_clusters=n_clusters)
    assert_raises(ValueError, km.fit, X)

def test_mini_match_k_means_invalid_init():
    km = MiniBatchKMeans(init="invalid", n_init=1, n_clusters=n_clusters)
    assert_raises(ValueError, km.fit, X)

def test_k_means_copyx():
    # Check if copy_x=False returns nearly equal X after de-centering.
    my_X = X.copy()
    km = KMeans(copy_x=False, n_clusters=n_clusters, random_state=42)
    km.fit(my_X)
    _check_fitted_model(km)

    # check if my_X is centered
    assert_array_almost_equal(my_X, X)

def test_k_means_non_collapsed():
    # Check k_means with a bad initialization does not yield a singleton
    # Starting with bad centers that are quickly ignored should not
    # result in a repositioning of the centers to the center of mass that
    # would lead to collapsed centers which in turns make the clustering
    # dependent of the numerical unstabilities.
    my_X = np.array([[1.1, 1.1], [0.9, 1.1], [1.1, 0.9], [0.9, 1.1]])
    array_init = np.array([[1.0, 1.0], [5.0, 5.0], [-5.0, -5.0]])
    km = KMeans(init=array_init, n_clusters=3, random_state=42, n_init=1)
    km.fit(my_X)

    # centers must not been collapsed
    assert_equal(len(np.unique(km.labels_)), 3)

    centers = km.cluster_centers_
    assert_true(np.linalg.norm(centers[0] - centers[1]) >= 0.1)
    assert_true(np.linalg.norm(centers[0] - centers[2]) >= 0.1)
    assert_true(np.linalg.norm(centers[1] - centers[2]) >= 0.1)

def test_predict():
    km = KMeans(n_clusters=n_clusters, random_state=42)

    km.fit(X)

    # sanity check: predict centroid labels
    pred = km.predict(km.cluster_centers_)
    assert_array_equal(pred, np.arange(n_clusters))

    # sanity check: re-predict labeling for training set samples
    pred = km.predict(X)
    assert_array_equal(pred, km.labels_)

    # re-predict labels for training set using fit_predict
    pred = km.fit_predict(X)
    assert_array_equal(pred, km.labels_)

def test_score():
    km1 = KMeans(n_clusters=n_clusters, max_iter=1, random_state=42)
    s1 = km1.fit(X).score(X)
    km2 = KMeans(n_clusters=n_clusters, max_iter=10, random_state=42)
    s2 = km2.fit(X).score(X)
    assert_greater(s2, s1)

def test_predict_minibatch_dense_input():
    mb_k_means = MiniBatchKMeans(n_clusters=n_clusters, random_state=40).fit(X)

    # sanity check: predict centroid labels
    pred = mb_k_means.predict(mb_k_means.cluster_centers_)
    assert_array_equal(pred, np.arange(n_clusters))

    # sanity check: re-predict labeling for training set samples
    pred = mb_k_means.predict(X)
    assert_array_equal(mb_k_means.predict(X), mb_k_means.labels_)

def test_predict_minibatch_kmeanspp_init_sparse_input():
    mb_k_means = MiniBatchKMeans(n_clusters=n_clusters, init='k-means++',
                                 n_init=10).fit(X_csr)

    # sanity check: re-predict labeling for training set samples
    assert_array_equal(mb_k_means.predict(X_csr), mb_k_means.labels_)

    # sanity check: predict centroid labels
    pred = mb_k_means.predict(mb_k_means.cluster_centers_)
    assert_array_equal(pred, np.arange(n_clusters))

    # check that models trained on sparse input also works for dense input at
    # predict time
    assert_array_equal(mb_k_means.predict(X), mb_k_means.labels_)

def test_predict_minibatch_random_init_sparse_input():
    mb_k_means = MiniBatchKMeans(n_clusters=n_clusters, init='random',
                                 n_init=10).fit(X_csr)

    # sanity check: re-predict labeling for training set samples
    assert_array_equal(mb_k_means.predict(X_csr), mb_k_means.labels_)

    # sanity check: predict centroid labels
    pred = mb_k_means.predict(mb_k_means.cluster_centers_)
    assert_array_equal(pred, np.arange(n_clusters))

    # check that models trained on sparse input also works for dense input at
    # predict time
    assert_array_equal(mb_k_means.predict(X), mb_k_means.labels_)

def test_input_dtypes():
    X_list = [[0, 0], [10, 10], [12, 9], [-1, 1], [2, 0], [8, 10]]
    X_int = np.array(X_list, dtype=np.int32)
    X_int_csr = sp.csr_matrix(X_int)
    init_int = X_int[:2]

    fitted_models = [
        KMeans(n_clusters=2).fit(X_list),
        KMeans(n_clusters=2).fit(X_int),
        KMeans(n_clusters=2, init=init_int, n_init=1).fit(X_list),
        KMeans(n_clusters=2, init=init_int, n_init=1).fit(X_int),
        # mini batch kmeans is very unstable on such a small dataset hence
        # we use many inits
        MiniBatchKMeans(n_clusters=2, n_init=10, batch_size=2).fit(X_list),
        MiniBatchKMeans(n_clusters=2, n_init=10, batch_size=2).fit(X_int),
        MiniBatchKMeans(n_clusters=2, n_init=10, batch_size=2).fit(X_int_csr),
        MiniBatchKMeans(n_clusters=2, batch_size=2,
                        init=init_int, n_init=1).fit(X_list),
        MiniBatchKMeans(n_clusters=2, batch_size=2,
                        init=init_int, n_init=1).fit(X_int),
        MiniBatchKMeans(n_clusters=2, batch_size=2,
                        init=init_int, n_init=1).fit(X_int_csr),
    ]
    expected_labels = [0, 1, 1, 0, 0, 1]
    scores = np.array([v_measure_score(expected_labels, km.labels_)
                       for km in fitted_models])
    assert_array_equal(scores, np.ones(scores.shape[0]))

def test_transform():
    km = KMeans(n_clusters=n_clusters)
    km.fit(X)
    X_new = km.transform(km.cluster_centers_)

    for c in range(n_clusters):
        assert_equal(X_new[c, c], 0)
        for c2 in range(n_clusters):
            if c != c2:
                assert_greater(X_new[c, c2], 0)

def test_fit_transform():
    X1 = KMeans(n_clusters=3, random_state=51).fit(X).transform(X)
    X2 = KMeans(n_clusters=3, random_state=51).fit_transform(X)
    assert_array_equal(X1, X2)

def test_n_init():
    # Check that increasing the number of init increases the quality
    n_runs = 5
    n_init_range = [1, 5, 10]
    inertia = np.zeros((len(n_init_range), n_runs))
    for i, n_init in enumerate(n_init_range):
        for j in range(n_runs):
            km = KMeans(n_clusters=n_clusters, init="random", n_init=n_init,
                        random_state=j).fit(X)
            inertia[i, j] = km.inertia_

    inertia = inertia.mean(axis=1)
    failure_msg = ("Inertia %r should be decreasing"
                   " when n_init is increasing.") % list(inertia)
    for i in range(len(n_init_range) - 1):
        assert_true(inertia[i] >= inertia[i + 1], failure_msg)

def test_k_means_function():
    # test calling the k_means function directly
    # catch output
    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        cluster_centers, labels, inertia = k_means(X, n_clusters=n_clusters,
                                                   verbose=True)
    finally:
        sys.stdout = old_stdout
    centers = cluster_centers
    assert_equal(centers.shape, (n_clusters, n_features))

    labels = labels
    assert_equal(np.unique(labels).shape[0], n_clusters)

    # check that the labels assignment are perfect (up to a permutation)
    assert_equal(v_measure_score(true_labels, labels), 1.0)
    assert_greater(inertia, 0.0)

    # check warning when centers are passed
    assert_warns(RuntimeWarning, k_means, X, n_clusters=n_clusters,
                 init=centers)

    # to many clusters desired
    assert_raises(ValueError, k_means, X, n_clusters=X.shape[0] + 1)

def test_x_squared_norms_init_centroids():
    """Test that x_squared_norms can be None in _init_centroids"""
    from sklearn.cluster.k_means_ import _init_centroids

    X_norms = np.sum(X**2, axis=1)
    precompute = _init_centroids(
        X, 3, "k-means++", random_state=0, x_squared_norms=X_norms)
    assert_array_equal(
        precompute,
        _init_centroids(X, 3, "k-means++", random_state=0))


import os

# ------------------------------------------------------------------------------
#
def get_experiment_frames(experiments, datadir=None):
    """
    read profiles for all sessions in the given 'experiments' dict.  That dict
    is expected to be like this:

    { 'test 1' : [ [ 'rp.session.thinkie.merzky.016609.0007',         'stampede popen sleep 1/1/1/1 (?)'] ],
      'test 2' : [ [ 'rp.session.ip-10-184-31-85.merzky.016610.0112', 'stampede shell sleep 16/8/8/4'   ] ],
      'test 3' : [ [ 'rp.session.ip-10-184-31-85.merzky.016611.0013', 'stampede shell mdrun 16/8/8/4'   ] ],
      'test 4' : [ [ 'rp.session.titan-ext4.marksant1.016607.0005',   'titan    shell sleep 1/1/1/1 a'  ] ],
      'test 5' : [ [ 'rp.session.titan-ext4.marksant1.016607.0006',   'titan    shell sleep 1/1/1/1 b'  ] ],
      'test 6' : [ [ 'rp.session.ip-10-184-31-85.merzky.016611.0013', 'stampede - isolated',            ],
                   [ 'rp.session.ip-10-184-31-85.merzky.016612.0012', 'stampede - integrated',          ],
                   [ 'rp.session.titan-ext4.marksant1.016607.0006',   'blue waters - integrated'        ] ]
    }  name in

    ie. iname in t is a list of experiment names, and each label has a list of
    session/label pairs, where the label will be later used to label (duh) plots.

    we return a similar dict where the session IDs are data frames
    """
    import pandas as pd

    exp_frames  = dict()

    if not datadir:
        datadir = os.getcwd()

    print 'reading profiles in %s' % datadir

    for exp in experiments:
        print " - %s" % exp
        exp_frames[exp] = list()

        for sid, label in experiments[exp]:
            print "   - %s" % sid

            import glob
            for prof in glob.glob ("%s/%s-pilot.*.prof" % (datadir, sid)):
                print "     - %s" % prof
                frame = get_profile_frame (prof)
                exp_frames[exp].append ([frame, label])

    return exp_frames

# ------------------------------------------------------------------------------
#
def get_profile_frame (prof):
    import pandas as pd
    return pd.read_csv(prof)

# ------------------------------------------------------------------------------
#
tmp = None
def add_concurrency (frame, tgt, spec):
    """
    add a column 'tgt' which is a cumulative sum of conditionals of enother row.

    The purpose is the following: if a unit enters a component, the tgt row counter is
    increased by 1, if the unit leaves the component, the counter is decreases by 1.
    For any time, the resulting row contains the number of units which is in the
    component.  Or state.  Or whatever.

    The arguments are:
        'tgt'  : name of the new column
        'spec' : a set of filters to determine if a unit enters or leaves

    'spec' is expected to be a dict of the following format:

        spec = { 'in'  : [{'col1' : 'pat1',
                           'col2' : 'pat2'},
                          ...],
                 'out' : [{'col3' : 'pat3',
                           'col4' : 'pat4'},
                          ...]
               }

    where:
        'in'    : filter set to determine the unit entering
        'out'   : filter set to determine the unit leaving
        'col'   : name of column for which filter is defined
        'event' : event which correlates to entering/leaving
        'msg'   : qualifier on the event, if event is not unique

    Example:
        spec = {'in'  : [{'state' :'Executing'}],
                'out' : [{'state' :'Done'},
                         {'state' :'Failed'},
                         {'state' :'Cancelled'}]
               }
        get_concurrency (df, 'concurrently_running', spec)
    """

    import numpy

    # create a temporary row over which we can do the commulative sum
    # --------------------------------------------------------------------------
    def _conc (row, spec):

        # row must match any filter dict in 'spec[in/out]'
        # for any filter dict it must match all col/pat pairs

        # for each in filter
        for f in spec['in']:
            match = 1
            # for each col/val in that filter
            for col, pat in f.iteritems():
                if row[col] != pat:
                    match = 0
                    break
            if match:
                # one filter matched!
              # print " + : %-20s : %.2f : %-20s : %s " % (row['uid'], row['time'], row['event'], row['message'])
                return 1

        # for each out filter
        for f in spec['out']:
            match = 1
            # for each col/val in that filter
            for col, pat in f.iteritems():
                if row[col] != pat:
                    match = 0
                    break
            if match:
                # one filter matched!
              # print " - : %-20s : %.2f : %-20s : %s " % (row['uid'], row['time'], row['event'], row['message'])
                return -1

        # no filter matched
      # print "   : %-20s : %.2f : %-20s : %s " % (row['uid'], row['time'], row['event'], row['message'])
        return  0
    # --------------------------------------------------------------------------

    # we only want to later look at changes of the concurrency -- leading or trailing
    # idle times are to be ignored.  We thus set repeating values of the cumsum to NaN,
    # so that they can be filtered out when ploting: df.dropna().plot(...).
    # That specifically will limit the plotted time range to the area of activity.
    # The full time range can still be plotted when ommitting the dropna() call.
    # --------------------------------------------------------------------------
    def _time (x):
        global tmp
        if     x != tmp: tmp = x
        else           : x   = numpy.NaN
        return x

    # --------------------------------------------------------------------------
    # sanitize concurrency: negative values indicate incorrect event ordering,
    # so we set the repesctive values to 0
    # --------------------------------------------------------------------------
    def _abs (x):
        if x < 0:
            return numpy.NaN
        return x
    # --------------------------------------------------------------------------

    frame[tgt] = frame.apply(lambda row: _conc(row, spec), axis=1).cumsum()
    frame[tgt] = frame.apply(lambda row: _abs (row[tgt]),  axis=1)
    frame[tgt] = frame.apply(lambda row: _time(row[tgt]),  axis=1)
  # print frame[[tgt, 'time']]

# ------------------------------------------------------------------------------
#
t0 = None
def calibrate_frame(frame, spec):
    """
    move the time axis of a profiling frame so that t_0 is at the first event
    matching the given 'spec'.  'spec' has the same format as described in
    'add_concurrency' (list of dicts with col:pat filters)
    """

    # --------------------------------------------------------------------------
    def _find_t0 (row, spec):

        # row must match any filter dict in 'spec[in/out]'
        # for any filter dict it must match all col/pat pairs
        global t0
        if t0 is not None:
            # already found t0
            return

        # for each col/val in that filter
        for f in spec:
            match = 1
            for col, pat in f.iteritems():
                if row[col] != pat:
                    match = 0
                    break
            if match:
                # one filter matched!
                t0 = row['time']
                return
    # --------------------------------------------------------------------------

    # --------------------------------------------------------------------------
    def _calibrate (row, t0):

        if t0 is None:
            # no t0...
            return

        return row['time'] - t0
    # --------------------------------------------------------------------------

    # we need to iterate twice over the frame: first to find t0, then to
    # calibrate the time axis
    global t0
    t0 = None # no t0
    frame.apply(lambda row: _find_t0  (row, spec), axis=1)

    if t0 == None:
        print "Can't recalibrate, no matching timestamp found"
        return
    frame['time'] = frame.apply(lambda row: _calibrate(row, t0  ), axis=1)

# ------------------------------------------------------------------------------
#
def create_plot():
    """
    create a plot object and tune its layout to our liking.
    """

    import matplotlib.pyplot as plt

    fig, plot = plt.subplots(figsize=(12,6))

    plot.xaxis.set_tick_params(width=1, length=7)
    plot.yaxis.set_tick_params(width=1, length=7)

    plot.spines['right' ].set_position(('outward', 10))
    plot.spines['top'   ].set_position(('outward', 10))
    plot.spines['bottom'].set_position(('outward', 10))
    plot.spines['left'  ].set_position(('outward', 10))

    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)

    fig.tight_layout()

    return fig, plot

# ------------------------------------------------------------------------------
#
def frame_plot (frames, axis, title=None, logx=False, logy=False,
                legend=True, figdir=None):
    """
    plot the given axis from the give data frame.  We create a plot, and plot
    all frames given in the list.  The list is expected to contain [frame,label]
    pairs

    frames: list of tuples of dataframes and labels
    frames  = [[stampede_df_1, 'stampede - popen'],
               [stampede_df_2, 'stampede - shell'],
               [stampede_df_3, 'stampede - ORTE' ]]

    axis:   tuple of data frame column index and axis label
    axis    = ['time', 'time (s)']
    """

    # create figure and layout
    fig, plot = create_plot()

    # set plot title
    if title:
        plot.set_title(title, y=1.05, fontsize=18)

    # plot the data frames
    # NOTE: we need to set labels separately, because of
    #       https://github.com/pydata/pandas/issues/9542
    labels = list()
    for frame, label in frames:
        try:
            frame.dropna().plot(ax=plot, logx=logx, logy=logy,
                    x=axis[0][0], y=axis[1][0],
                    drawstyle='steps',
                    label=label, legend=False)
        except Exception as e:
            print "skipping frame '%s': '%s'" % (label, e)

    if legend:
        plot.legend(labels=labels, loc='upper right', fontsize=14, frameon=True)

    # set axis labels
    plot.set_xlabel(axis[0][1], fontsize=14)
    plot.set_ylabel(axis[1][1], fontsize=14)
    plot.set_frame_on(True)

    # save as png and pdf.  Use the title as base for names
    if title: base = title
    else    : base = "%s_%s" % (axis[0][1], axis[1][1])

    # clean up base name -- only keep alphanum and such
    import re
    base = re.sub('[^a-zA-Z0-9\.\-]', '_', base)
    base = re.sub('_+',               '_', base)

    if not figdir:
        figdir = os.getcwd()

    print 'saving %s/%s.png' % (figdir, base)
    fig.savefig('%s/%s.png' % (figdir, base), bbox_inches='tight')

    print 'saving %s/%s.pdf' % (figdir, base)
    fig.savefig('%s/%s.pdf' % (figdir, base), bbox_inches='tight')

    return fig, plot

# ------------------------------------------------------------------------------
#
def create_analytical_frame (idx, kind, args, limits, step):
    """
    create an artificial data frame, ie. a data frame which does not contain
    data gathered from an experiment, but data representing an analytical
    construct of some 'kind'.

    idx:    data frame column index to fill (a time column is always created)
    kind:   construct to use (only 'rate' is supporte right now)
    args:   construct specific parameters
    limits: time range for which data are to be created
    step:   time steps for which data are to be created
    """

    import pandas as pd

    # --------------------------------------------------------------------------
    def _frange(start, stop, step):
        while start <= stop:
            yield start
            start += step
    # --------------------------------------------------------------------------

    if kind == 'rate' :
        t_0  = args.get ('t_0',  0.0)
        rate = args.get ('rate', 1.0)
        data = list()
        for t in _frange(limits[0], limits[1], step):
            data.append ({'time': t+t_0, idx: t*rate})
        return pd.DataFrame (data)

    else:
        raise ValueError ("No such frame kind '%s'" % kind)

# ------------------------------------------------------------------------------


# Author: Denis A. Engemann <d.engemann@gmail.com>
#
# License: BSD (3-clause)

import os.path as op
import numpy as np
from numpy.testing import assert_array_almost_equal
from nose.tools import assert_equal, assert_raises

from mne import io, Epochs, read_events, pick_types
from mne.utils import requires_sklearn, check_version
from mne.decoding import compute_ems, EMS

data_dir = op.join(op.dirname(__file__), '..', '..', 'io', 'tests', 'data')
curdir = op.join(op.dirname(__file__))

raw_fname = op.join(data_dir, 'test_raw.fif')
event_name = op.join(data_dir, 'test-eve.fif')

tmin, tmax = -0.2, 0.5
event_id = dict(aud_l=1, vis_l=3)

@requires_sklearn
def test_ems():
    """Test event-matched spatial filters"""
    raw = io.read_raw_fif(raw_fname, preload=False)

    # create unequal number of events
    events = read_events(event_name)
    events[-2, 2] = 3
    picks = pick_types(raw.info, meg=True, stim=False, ecg=False,
                       eog=False, exclude='bads')
    picks = picks[1:13:3]
    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
                    baseline=(None, 0), preload=True)
    assert_raises(ValueError, compute_ems, epochs, ['aud_l', 'vis_l'])
    epochs = epochs.equalize_event_counts(epochs.event_id, copy=False)[0]

    assert_raises(KeyError, compute_ems, epochs, ['blah', 'hahah'])
    surrogates, filters, conditions = compute_ems(epochs)
    assert_equal(list(set(conditions)), [1, 3])

    events = read_events(event_name)
    event_id2 = dict(aud_l=1, aud_r=2, vis_l=3)
    epochs = Epochs(raw, events, event_id2, tmin, tmax, picks=picks,
                    baseline=(None, 0), preload=True)
    epochs = epochs.equalize_event_counts(epochs.event_id, copy=False)[0]

    n_expected = sum([len(epochs[k]) for k in ['aud_l', 'vis_l']])

    assert_raises(ValueError, compute_ems, epochs)
    surrogates, filters, conditions = compute_ems(epochs, ['aud_r', 'vis_l'])
    assert_equal(n_expected, len(surrogates))
    assert_equal(n_expected, len(conditions))
    assert_equal(list(set(conditions)), [2, 3])

    # test compute_ems cv
    epochs = epochs['aud_r', 'vis_l']
    epochs.equalize_event_counts(epochs.event_id)
    if check_version('sklearn', '0.18'):
        from sklearn.model_selection import StratifiedKFold
        cv = StratifiedKFold()
    else:
        from sklearn.cross_validation import StratifiedKFold
        cv = StratifiedKFold(epochs.events[:, 2])
    compute_ems(epochs, cv=cv)
    compute_ems(epochs, cv=2)
    assert_raises(ValueError, compute_ems, epochs, cv='foo')
    assert_raises(ValueError, compute_ems, epochs, cv=len(epochs) + 1)
    raw.close()

    # EMS transformer, check that identical to compute_ems
    X = epochs.get_data()
    y = epochs.events[:, 2]
    X = X / np.std(X)  # X scaled outside cv in compute_ems
    Xt, coefs = list(), list()
    ems = EMS()
    assert_equal(ems.__repr__(), '<EMS: not fitted.>')
    # manual leave-one-out to avoid sklearn version problem
    for test in range(len(y)):
        train = np.setdiff1d(range(len(y)), test)
        ems.fit(X[train], y[train])
        coefs.append(ems.filters_)
        Xt.append(ems.transform(X[[test]]))
    assert_equal(ems.__repr__(), '<EMS: fitted with 4 filters on 2 classes.>')
    assert_array_almost_equal(filters, np.mean(coefs, axis=0))
    assert_array_almost_equal(surrogates, np.vstack(Xt))


""" This is the Bokeh charts testing interface.

"""
#-----------------------------------------------------------------------------
# Copyright (c) 2012 - 2014, Continuum Analytics, Inc. All rights reserved.
#
# Powered by the Bokeh Development Team.
#
# The full license is in the file LICENSE.txt, distributed with this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
from __future__ import absolute_import

from collections import OrderedDict
import unittest

import numpy as np
from numpy.testing import assert_array_equal
import pandas as pd

from bokeh.charts import Dot
from bokeh.util.testing import create_chart

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

class TestDot(unittest.TestCase):
    def test_supported_input(self):
        xyvalues = OrderedDict()
        xyvalues['python']=[2, 5]
        xyvalues['pypy']=[12, 40]
        xyvalues['jython']=[22, 30]

        xyvaluesdf = pd.DataFrame(xyvalues, index=['lists', 'loops'])

        cat = ['lists', 'loops']
        catjython = ['lists:0.75', 'loops:0.75']
        catpypy = ['lists:0.5', 'loops:0.5']
        catpython = ['lists:0.25', 'loops:0.25']
        python = seg_top_python = [2, 5]
        pypy = seg_top_pypy = [12, 40]
        jython = seg_top_jython = [22, 30]
        zero = [0, 0]

        for i, _xy in enumerate([xyvalues, xyvaluesdf]):
            hm = create_chart(Dot, _xy, cat=cat)
            builder = hm._builders[0]
            self.assertEqual(sorted(builder._groups), sorted(list(xyvalues.keys())))
            assert_array_equal(builder._data['cat'], cat)
            assert_array_equal(builder._data['catjython'], catjython)
            assert_array_equal(builder._data['catpython'], catpython)
            assert_array_equal(builder._data['catpypy'], catpypy)

            assert_array_equal(builder._data['python'], python)
            assert_array_equal(builder._data['jython'], jython)
            assert_array_equal(builder._data['pypy'], pypy)

            assert_array_equal(builder._data['seg_top_python'], seg_top_python)
            assert_array_equal(builder._data['seg_top_jython'], seg_top_jython)
            assert_array_equal(builder._data['seg_top_pypy'], seg_top_pypy)

            assert_array_equal(builder._data['z_python'], zero)
            assert_array_equal(builder._data['z_pypy'], zero)
            assert_array_equal(builder._data['z_jython'], zero)
            assert_array_equal(builder._data['zero'], zero)

        lvalues = [[2, 5], [12, 40], [22, 30]]
        for _xy in [lvalues, np.array(lvalues)]:
            hm = create_chart(Dot, _xy, cat=cat)
            builder = hm._builders[0]

            self.assertEqual(builder._groups, ['0', '1', '2'])
            assert_array_equal(builder._data['cat'], cat)
            assert_array_equal(builder._data['cat0'], catpython)
            assert_array_equal(builder._data['cat1'], catpypy)
            assert_array_equal(builder._data['cat2'], catjython)
            assert_array_equal(builder._data['0'], python)
            assert_array_equal(builder._data['1'], pypy)
            assert_array_equal(builder._data['2'], jython)

            assert_array_equal(builder._data['seg_top_0'], seg_top_python)
            assert_array_equal(builder._data['seg_top_1'], seg_top_pypy)
            assert_array_equal(builder._data['seg_top_2'], seg_top_jython)

            assert_array_equal(builder._data['z_0'], zero)
            assert_array_equal(builder._data['z_1'], zero)
            assert_array_equal(builder._data['z_2'], zero)
            assert_array_equal(builder._data['zero'], zero)


import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
import aplpy
from astropy.wcs import WCS
import sys, os
from getIQU import IQU
from astropy import coordinates as coord
from astropy.coordinates import SkyCoord
from astropy import units as u
from scipy.interpolate import griddata
plt.ion()

root_dir = '/home/wizwit/miscellaneous_projects/carina/carinaData'
blast250_file = os.path.join(root_dir, 'smooth/3.0_arcmin/carinaneb_250_smoothed_3.0_rl.fits')

beta = 1.27

def getPsi(path_to_file):
    I, Q, U, __, wcs = IQU(path_to_file)
    Pvals = np.sqrt(Q**2 + U**2)
    pvals = Pvals/I
    # pvals /= pol_eff[band_idx]
    psi = 0.5*np.arctan2(U,Q)
    return I, Q, U, wcs, psi

I, __, __, wcs_250, __, = getPsi(blast250_file)

#tau_d = (nu/nu0)**beta
# See Walker pg. 71
# nu0 = frequency at which dust emission becomes optically thin
#nu0 = 0.103 * Td # 0.103 (THz/K) * Td
#Inu_dust = Bnu(Td)*(1.0 - np.exp(1.0 - e**(-1.0*tau_d))

# See Walker pg. 69
# Av = 1.086*tau_d
# N_H = 1.79e21 * Av # (atoms/cm**2 mag)

# 1) Solve tau_d for temperature
# 2) Plug into Inu_dust equation


import numpy as np

from pandas import Index, Timedelta, timedelta_range
import pandas._testing as tm

class TestJoin:
    def test_append_join_nondatetimeindex(self):
        rng = timedelta_range("1 days", periods=10)
        idx = Index(["a", "b", "c", "d"])

        result = rng.append(idx)
        assert isinstance(result[0], Timedelta)

        # it works
        rng.join(idx, how="outer")

    def test_join_self(self, join_type):
        index = timedelta_range("1 day", periods=10)
        joined = index.join(index, how=join_type)
        tm.assert_index_equal(index, joined)

    def test_does_not_convert_mixed_integer(self):
        df = tm.makeCustomDataframe(
            10,
            10,
            data_gen_f=lambda *args, **kwargs: np.random.randn(),
            r_idx_type="i",
            c_idx_type="td",
        )
        str(df)

        cols = df.columns.join(df.index, how="outer")
        joined = cols.join(df.columns)
        assert cols.dtype == np.dtype("O")
        assert cols.dtype == joined.dtype
        tm.assert_index_equal(cols, joined)

    def test_join_preserves_freq(self):
        # GH#32157
        tdi = timedelta_range("1 day", periods=10)
        result = tdi[:5].join(tdi[5:], how="outer")
        assert result.freq == tdi.freq
        tm.assert_index_equal(result, tdi)

        result = tdi[:5].join(tdi[6:], how="outer")
        assert result.freq is None
        expected = tdi.delete(5)
        tm.assert_index_equal(result, expected)


import pytest
from os import path

DB_PATH = '~/.demolog/meta.db'
DATA_DIR = path.join(path.dirname(path.realpath(__file__)),'../../example_data/')

def test_implant_angle_filter():
	from labbookdb.report.selection import animal_id, animal_treatments, animal_operations
	import numpy as np

	db_path=DB_PATH
	df = animal_operations(db_path)
	#validate target by code
	df = df[~df['OrthogonalStereotacticTarget_code'].isnull()]
	df = df[df['OrthogonalStereotacticTarget_code'].str.contains('dr')]
	#check pitch
	df = df[~df['OrthogonalStereotacticTarget_pitch'].isin([0,np.NaN])]
	animals = df['Animal_id'].tolist()
	animals_eth = [animal_id(db_path,'ETH/AIC',i,reverse=True) for i in animals]

	assert animals_eth == ['5684']

def test_animal_cage_treatment_control_in_report():
	"""Check if animal which died before the cagetreatment was applied to its last home cage is indeed not showing a cage treatment, but still showing the animal treatment."""
	from labbookdb.report.tracking import animals_info

	df = animals_info(DB_PATH,
		save_as=None,
		functional_scan_responders=True,
		treatments=True,
		)
	assert df[df['ETH/AIC']=='6255']['cage_treatment'].values[0] == ""
	assert df[df['ETH/AIC']=='6255']['animal_treatment'].values[0] == 'aFluIV_'

def test_animal_id():
	"""Check if LabbookDB animal ID is correctly reported based on external database identifier."""
	from labbookdb.report.selection import animal_id

	my_id = animal_id(DB_PATH,
		database='ETH/AIC',
		identifier='6255'
		)
	assert my_id == 41

def test_bids_eventsfile():
	"""Check if correct BIDS events file can be sourced."""
	from labbookdb.report.tracking import bids_eventsfile
	import pandas as pd

	df = bids_eventsfile(DB_PATH,'chr_longSOA')
	bids_eventsfile = path.join(DATA_DIR,'bids_eventsfile.csv')
	df_ = pd.read_csv(bids_eventsfile, index_col=0)

	assert df[['onset','duration']].equals(df_[['onset','duration']])

def test_drinking_by_cage_treatment(
	treatment_relative_date=True,
	rounding='D',
	):
	from labbookdb.report.tracking import treatment_group, append_external_identifiers, qualitative_dates, cage_consumption
	from labbookdb.report.selection import cage_periods, cage_drinking_measurements

	known_cage_ids = [25, 38, 41]
	known_consumption_values = [2.35, 2.51, 2.94, 2.95, 3.16, 3.17, 3.22, 3.23, 3.24, 3.25, 3.49, 3.63, 3.72, 4.04, 4.09, 4.58, 4.98, 5.15, 5.31, 5.39, 5.54, 5.97, 6.73, 6.78]

	df = cage_drinking_measurements(DB_PATH,['cFluDW'])
	df = cage_consumption(DB_PATH,df)

	fuzzy_matching = {
		"ofM":[-14,-15,-13,-7,-8,-6],
		"ofMaF":[0,-1],
		"ofMcF1":[14,13,15],
		"ofMcF2":[28,27,29],
		"ofMpF":[45,44,46,43,47],
	}
	df = qualitative_dates(df,
		iterator_column='Cage_id',
		date_column='relative_end_date',
		label='qualitative_date',
		fuzzy_matching=fuzzy_matching,
		)

	cage_ids = sorted(df['Cage_id'].unique())
	assert cage_ids == known_cage_ids

	consumption_values = df['day_animal_consumption'].values
	consumption_values = [round(i, 2) for i in consumption_values]
	consumption_values = sorted(list(set(consumption_values)))
	assert consumption_values == known_consumption_values

def test_groups():
	"""Create a `pandas.DataFrame` containing treatment and genotype group assignments"""
	from labbookdb.report.tracking import treatment_group, append_external_identifiers

	known_sorted_ids = [
		'5667',
		'5668',
		'5673',
		'5674',
		'5675',
		'5689',
		'5690',
		'5691',
		'5692',
		'5694',
		'5699',
		'5700',
		'5704',
		'5705',
		'5706',
		'6254',
		'6256',
		'6262',
		]

	df = treatment_group(DB_PATH, ['cFluDW','cFluDW_'], level='cage')
	df = append_external_identifiers(DB_PATH, df, ['Genotype_code'])
	sorted_ids = sorted(df['ETH/AIC'].tolist())

	assert sorted_ids == known_sorted_ids


"""
====================
Active Contour Model
====================

The active contour model is a method to fit open or closed splines to lines or
edges in an image. It works by minimising an energy that is in part defined by
the image and part by the spline's shape: length and smoothness. The
minimization is done implicitly in the shape energy and explicitly in the
image energy.

In the following two examples the active contour model is used (1) to segment
the face of a person from the rest of an image by fitting a closed curve
to the edges of the face and (2) to find the darkest curve between two fixed
points while obeying smoothness considerations. Typically it is a good idea to
smooth images a bit before analyzing, as done in the following examples.

.. [1] *Snakes: Active contour models*. Kass, M.; Witkin, A.; Terzopoulos, D.
       International Journal of Computer Vision 1 (4): 321 (1988).

We initialize a circle around the astronaut's face and use the default boundary
condition ``bc='periodic'`` to fit a closed curve. The default parameters
``w_line=0, w_edge=1`` will make the curve search towards edges, such as the
boundaries of the face.
"""

import numpy as np
import matplotlib.pyplot as plt
from skimage.color import rgb2gray
from skimage import data
from skimage.filters import gaussian_filter
from skimage.segmentation import active_contour

# Test scipy version, since active contour is only possible
# with recent scipy version
import scipy
scipy_version = list(map(int, scipy.__version__.split('.')))
new_scipy = scipy_version[0] > 0 or \
            (scipy_version[0] == 0 and scipy_version[1] >= 14)

img = data.astronaut()
img = rgb2gray(img)

s = np.linspace(0, 2*np.pi, 400)
x = 220 + 100*np.cos(s)
y = 100 + 100*np.sin(s)
init = np.array([x, y]).T

if not new_scipy:
    print('You are using an old version of scipy. '
          'Active contours is implemented for scipy versions '
          '0.14.0 and above.')

if new_scipy:
    snake = active_contour(gaussian_filter(img, 3),
                           init, alpha=0.015, beta=10, gamma=0.001)

    fig = plt.figure(figsize=(7, 7))
    ax = fig.add_subplot(111)
    plt.gray()
    ax.imshow(img)
    ax.plot(init[:, 0], init[:, 1], '--r', lw=3)
    ax.plot(snake[:, 0], snake[:, 1], '-b', lw=3)
    ax.set_xticks([]), ax.set_yticks([])
    ax.axis([0, img.shape[1], img.shape[0], 0])

"""
.. image:: PLOT2RST.current_figure

Here we initialize a straight line between two points, `(5, 136)` and
`(424, 50)`, and require that the spline has its end points there by giving
the boundary condition `bc='fixed'`. We furthermore make the algorithm search
for dark lines by giving a negative `w_line` value.
"""

img = data.text()

x = np.linspace(5, 424, 100)
y = np.linspace(136, 50, 100)
init = np.array([x, y]).T

if new_scipy:
    snake = active_contour(gaussian_filter(img, 1), init, bc='fixed',
                           alpha=0.1, beta=1.0, w_line=-5, w_edge=0, gamma=0.1)

    fig = plt.figure(figsize=(9, 5))
    ax = fig.add_subplot(111)
    plt.gray()
    ax.imshow(img)
    ax.plot(init[:, 0], init[:, 1], '--r', lw=3)
    ax.plot(snake[:, 0], snake[:, 1], '-b', lw=3)
    ax.set_xticks([]), ax.set_yticks([])
    ax.axis([0, img.shape[1], img.shape[0], 0])

plt.show()

"""
.. image:: PLOT2RST.current_figure
"""


#!/usr/bin/python

###########################################################################
# PreProcessing.py :
#
# Iterate through the expert labelmap and create 65x65 patches around the
# central pixel. All positive pixels are used as positives input cases.
# The same amount of negatives is randomly picked. For each input patch,
# the corresponding filename and expected output are written to a text file
# and will be used later to create the database.
#
###########################################################################

import os
import glob
import json

import matplotlib.image as mpimg
import matplotlib.pyplot as plt

import random
import math

# Create image set and save expected output
def createImgSet(expertImg, inputImg, filenamePrefix, fileOutputDir, textFile):

    w = 32  # Patch size
    count = 0  # Input count
    negativeImageIndex = []  # Whole image index giving negative output
    negativeIndex = []  # Vessel bound index giving negative output

    # Write filename and expected output
    textFile = open(textFile, "a")
    textFile.truncate()  # Erase file

    resample = 1  # WARNING: resample pixels to reduce training size for Debug

    # Iterate through the expert label map
    for i in range(0, expertImg.shape[0], resample):
        for j in range(0, expertImg.shape[1], resample):

            if j > w and j + w + 1 < inputImg.shape[1]:
                if i > w and i + w + 1 < inputImg.shape[0]:

                    # Centerline pixel (positive)
                    if expertImg[i, j] > 0.5:
                        count += 1
                        filename = os.path.join(
                            "1", filenamePrefix + "_" + str(i) + "_" + str(j) + ".png")
                        textFile.write(filename + " " + str(1) + "\n")
                        plt.imsave(os.path.join(fileOutputDir, filename),
                                   inputImg[i - w:i + w + 1, j - w:j + w + 1], cmap='Greys_r')

                    # Vessel bound pixel (negative)
                    elif expertImg[i, j] > 0:
                        negativeIndex.append([i, j])

                    # Background pixel (negative)
                    else:
                        negativeImageIndex.append([i, j])

    # Pick random negatives from vessel bound
    rndmNegativeInd = random.sample(negativeIndex, int(math.ceil(0.8 * count)))
    for [i, j] in rndmNegativeInd:
        filename = os.path.join("0", filenamePrefix +
                                "_" + str(i) + "_" + str(j) + ".png")
        textFile.write(filename + " " + str(0) + "\n")
        plt.imsave(os.path.join(fileOutputDir, filename),
                   inputImg[i - w:i + w + 1, j - w:j + w + 1], cmap='Greys_r')

    # Pick random negatives from the entire image
    rndmNegativeImageInd = random.sample(
        negativeImageIndex, int(math.ceil(0.2 * count)))
    for [i, j] in rndmNegativeImageInd:
        filename = os.path.join("0", filenamePrefix +
                                "_" + str(i) + "_" + str(j) + ".png")
        textFile.write(filename + " " + str(0) + "\n")
        plt.imsave(os.path.join(fileOutputDir, filename),
                   inputImg[i - w:i + w + 1, j - w:j + w + 1], cmap='Greys_r')

    textFile.close()
    print(count)

########
# Main #
########

# Path variable
script_params = json.load(open('params.json'))
caffe_root = script_params['CAFFE_SRC_ROOT']
hardDrive_root = script_params['CNN_DATA_ROOT']
proj_rel_path = script_params['PROJECT_REL_PATH']

caffe_proj_root = os.path.join(caffe_root, "data", proj_rel_path)
hardDrive_proj_root = os.path.join(hardDrive_root, proj_rel_path)

trainDataDir = os.path.join(hardDrive_proj_root, "training")
valDataDir = os.path.join(hardDrive_proj_root, "testing")

# Text file
trainFilename = os.path.join(caffe_proj_root, "train.txt")
trainFile = open(trainFilename, "w+")
trainFile.truncate()  # Erase file
trainFile.close()

valFilename = os.path.join(caffe_proj_root, "val.txt")
valFile = open(valFilename, "w+")
valFile.truncate()  # Erase file
valFile.close()

# Output patches directories
trainFileOutputDir = os.path.join(trainDataDir, "out")
if not os.path.exists(trainFileOutputDir):
    os.mkdir(trainFileOutputDir)

for label in range(2):
    curLabelOutputDir = os.path.join(trainFileOutputDir, str(label))
    if not os.path.exists(curLabelOutputDir):
        os.mkdir(curLabelOutputDir)

valFileOutputDir = os.path.join(valDataDir, "out")
if not os.path.exists(valFileOutputDir):
    os.mkdir(valFileOutputDir)

for label in range(2):
    curLabelOutputDir = os.path.join(valFileOutputDir, str(label))
    if not os.path.exists(curLabelOutputDir):
        os.mkdir(curLabelOutputDir)

# Images directories
trainExpertDir = os.path.join(trainDataDir, "expert")
trainImgDir = os.path.join(trainDataDir, "images")

valExpertDir = os.path.join(valDataDir, "expert")
valImgDir = os.path.join(valDataDir, "images")

# Create train set
trainImages = glob.glob(os.path.join(trainImgDir, "*.png"))

for trainImage in trainImages:

    print(trainImage)

    # Get image ID
    trainImagePrefix = os.path.basename(os.path.splitext(trainImage)[0])

    # Set filename
    trainExpertFile = os.path.join(
        trainExpertDir, trainImagePrefix + "_expert.png")
    trainImageFile = os.path.join(trainImgDir, trainImagePrefix + ".png")

    # Load images
    trainExpert = mpimg.imread(trainExpertFile)

    # print trainExpert.shape
    # trainExpert=trainExpert[:,:,0]

    trainImg = mpimg.imread(trainImageFile)
    # trainImg=trainImg[:,:,0]

    # Write images and text files
    createImgSet(trainExpert, trainImg, trainImagePrefix,
                 trainFileOutputDir, trainFilename)

# Create validation set
valImages = glob.glob(os.path.join(valImgDir, "*.png"))
for valImage in valImages:

    print(valImage)

    # Get image ID
    valImagePrefix = os.path.basename(os.path.splitext(valImage)[0])

    # Set filename
    valExpertFilename = os.path.join(
        valExpertDir, valImagePrefix + "_expert.png")
    valImgFilename = os.path.join(valImgDir, valImagePrefix + ".png")

    # Load images
    valExpert = mpimg.imread(valExpertFilename)

    # valExpert=valExpert[:,:,0]
    valImg = mpimg.imread(valImgFilename)
    # valImg=valImg[:,:,0]

    # Write images and text files
    createImgSet(valExpert, valImg, valImagePrefix,
                 valFileOutputDir, valFilename)


"""
=================================================================
Test with permutations the significance of a classification score
=================================================================

In order to test if a classification score is significative a technique
in repeating the classification procedure after randomizing, permuting,
the labels. The p-value is then given by the percentage of runs for
which the score obtained is greater than the classification score
obtained in the first place.

"""

# Author:  Alexandre Gramfort <alexandre.gramfort@inria.fr>
# License: BSD

print __doc__

import numpy as np
import pylab as pl

from sklearn.svm import SVC
from sklearn.cross_validation import StratifiedKFold, permutation_test_score
from sklearn import datasets
from sklearn.metrics import zero_one_score

##############################################################################
# Loading a dataset
iris = datasets.load_iris()
X = iris.data
y = iris.target
n_classes = np.unique(y).size

# Some noisy data not correlated
random = np.random.RandomState(seed=0)
E = random.normal(size=(len(X), 2200))

# Add noisy data to the informative features for make the task harder
X = np.c_[X, E]

svm = SVC(kernel='linear')
cv = StratifiedKFold(y, 2)

score, permutation_scores, pvalue = permutation_test_score(
    svm, X, y, zero_one_score, cv=cv, n_permutations=100, n_jobs=1)

print "Classification score %s (pvalue : %s)" % (score, pvalue)

###############################################################################
# View histogram of permutation scores
pl.hist(permutation_scores, 20, label='Permutation scores')
ylim = pl.ylim()
# BUG: vlines(..., linestyle='--') fails on older versions of matplotlib
#pl.vlines(score, ylim[0], ylim[1], linestyle='--',
#          color='g', linewidth=3, label='Classification Score'
#          ' (pvalue %s)' % pvalue)
#pl.vlines(1.0 / n_classes, ylim[0], ylim[1], linestyle='--',
#          color='k', linewidth=3, label='Luck')
pl.plot(2 * [score], ylim, '--g', linewidth=3,
        label='Classification Score'
        ' (pvalue %s)' % pvalue)
pl.plot(2 * [1. / n_classes], ylim, '--k', linewidth=3, label='Luck')

pl.ylim(ylim)
pl.legend()
pl.xlabel('Score')
pl.show()


import os
import random
import shutil

import numpy as np

import autosklearn.automl
from autosklearn.constants import *

class AutoSklearnClassifier(autosklearn.automl.AutoML):
    """This class implements the classification task. It must not be pickled!

    Parameters
    ----------
    time_left_for_this_task : int, optional (default=3600)
        Time limit in seconds for the search for appropriate classification
        models. By increasing this value, *auto-sklearn* will find better
        configurations.

    per_run_time_limit : int, optional (default=360)
        Time limit for a single call to machine learning model.

    initial_configurations_via_metalearning : int, optional (default=25)

    ensemble_size : int, optional (default=50)

    ensemble_nbest : int, optional (default=50)

    seed : int, optional (default=1)

    ml_memory_limit : int, optional (3000)
        Memory limit for the machine learning algorithm. If the machine
        learning algorithm allocates tries to allocate more memory,
        its evaluation will be stopped.
    """

    def __init__(self, time_left_for_this_task=3600,
                 per_run_time_limit=360,
                 initial_configurations_via_metalearning=25,
                 ensemble_size=50, ensemble_nbest=50, seed=1,
                 ml_memory_limit=3000):
        random_number = random.randint(0, 10000)

        pid = os.getpid()
        output_dir = "/tmp/autosklearn_output_%d_%d" % (pid, random_number)
        tmp_dir = "/tmp/autosklearn_tmp_%d_%d" % (pid, random_number)
        os.makedirs(output_dir)
        os.makedirs(tmp_dir)

        super(AutoSklearnClassifier, self).__init__(
            tmp_dir, output_dir, time_left_for_this_task, per_run_time_limit,
            log_dir=tmp_dir,
            initial_configurations_via_metalearning=initial_configurations_via_metalearning,
            ensemble_size=ensemble_size, ensemble_nbest=ensemble_nbest,
            seed=seed, ml_memory_limit=ml_memory_limit)

    def __del__(self):
        self._delete_output_directories()

    def _create_output_directories(self):
        os.makedirs(self.output_dir)
        os.makedirs(self.tmp_dir)

    def _delete_output_directories(self):
        shutil.rmtree(self.tmp_dir)
        shutil.rmtree(self.output_dir)

    def fit(self, X, y, metric='acc_metric', feat_type=None):
        """Fit *autosklearn* to given training set (X, y).

        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            The target classes.

        metric : str, optional (default='acc_metric')
            The metric to optimize for. Can be one of: ['acc_metric',
            'auc_metric', 'bac_metric', 'f1_metric', 'pac_metric']

        feat_type : list, optional (default=None)
            List of :python:`len(X.shape[1])` describing if an attribute is
            continuous or categorical. Categorical attributes will
            automatically 1Hot encoded.
        """
        # Fit is supposed to be idempotent!
        self._delete_output_directories()
        self._create_output_directories()

        y = np.atleast_1d(y)

        if y.ndim == 1:
            # reshape is necessary to preserve the data contiguity against vs
            # [:, np.newaxis] that does not.
            y = np.reshape(y, (-1, 1))

        self.n_outputs_ = y.shape[1]

        y = np.copy(y)

        self.classes_ = []
        self.n_classes_ = []

        for k in xrange(self.n_outputs_):
            classes_k, y[:, k] = np.unique(y[:, k], return_inverse=True)
            self.classes_.append(classes_k)
            self.n_classes_.append(classes_k.shape[0])

        self.n_classes_ = np.array(self.n_classes_, dtype=np.int)

        if self.n_outputs_ > 1:
            task = MULTILABEL_CLASSIFICATION
        else:
            if len(self.classes_[0]) == 2:
                task = BINARY_CLASSIFICATION
            else:
                task = MULTICLASS_CLASSIFICATION

        # TODO: fix metafeatures calculation to allow this!
        if y.shape[1] == 1:
            y = y.flatten()

        return super(AutoSklearnClassifier, self).fit(X, y, task, metric,
                                                   feat_type)

    def predict(self, X):
        """Predict class for X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]

        Returns
        -------
        y : array of shape = [n_samples] or [n_samples, n_outputs]
            The predicted classes.
        """
        return super(AutoSklearnClassifier, self).predict(X)

class AutoSklearnRegressor(autosklearn.automl.AutoML):
    def __init__(self, **kwargs):
        raise NotImplementedError()


"""
========================================================================
Concentration Prior Type Analysis of Variation Bayesian Gaussian Mixture
========================================================================

This example plots the ellipsoids obtained from a toy dataset (mixture of three
Gaussians) fitted by the ``BayesianGaussianMixture`` class models with a
Dirichlet distribution prior
(``weight_concentration_prior_type='dirichlet_distribution'``) and a Dirichlet
process prior (``weight_concentration_prior_type='dirichlet_process'``). On
each figure, we plot the results for three different values of the weight
concentration prior.

The ``BayesianGaussianMixture`` class can adapt its number of mixture
componentsautomatically. The parameter ``weight_concentration_prior`` has a
direct link with the resulting number of components with non-zero weights.
Specifying a low value for the concentration prior will make the model put most
of the weight on few components set the remaining components weights very close
to zero. High values of the concentration prior will allow a larger number of
components to be active in the mixture.

The Dirichlet process prior allows to define an infinite number of components
and automatically selects the correct number of components: it activates a
component only if it is necessary.

On the contrary the classical finite mixture model with a Dirichlet
distribution prior will favor more uniformly weighted components and therefore
tends to divide natural clusters into unnecessary sub-components.
"""
# Author: Thierry Guillemot <thierry.guillemot.work@gmail.com>
# License: BSD 3 clause

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

from sklearn.mixture import BayesianGaussianMixture

print(__doc__)

def plot_ellipses(ax, weights, means, covars):
    for n in range(means.shape[0]):
        eig_vals, eig_vecs = np.linalg.eigh(covars[n])
        unit_eig_vec = eig_vecs[0] / np.linalg.norm(eig_vecs[0])
        angle = np.arctan2(unit_eig_vec[1], unit_eig_vec[0])
        # Ellipse needs degrees
        angle = 180 * angle / np.pi
        # eigenvector normalization
        eig_vals = 2 * np.sqrt(2) * np.sqrt(eig_vals)
        ell = mpl.patches.Ellipse(means[n], eig_vals[0], eig_vals[1],
                                  180 + angle)
        ell.set_clip_box(ax.bbox)
        ell.set_alpha(weights[n])
        ell.set_facecolor('#56B4E9')
        ax.add_artist(ell)

def plot_results(ax1, ax2, estimator, X, y, title, plot_title=False):
    ax1.set_title(title)
    ax1.scatter(X[:, 0], X[:, 1], s=5, marker='o', color=colors[y], alpha=0.8)
    ax1.set_xlim(-2., 2.)
    ax1.set_ylim(-3., 3.)
    ax1.set_xticks(())
    ax1.set_yticks(())
    plot_ellipses(ax1, estimator.weights_, estimator.means_,
                  estimator.covariances_)

    ax2.get_xaxis().set_tick_params(direction='out')
    ax2.yaxis.grid(True, alpha=0.7)
    for k, w in enumerate(estimator.weights_):
        ax2.bar(k - .45, w, width=0.9, color='#56B4E9', zorder=3)
        ax2.text(k, w + 0.007, "%.1f%%" % (w * 100.),
                 horizontalalignment='center')
    ax2.set_xlim(-.6, 2 * n_components - .4)
    ax2.set_ylim(0., 1.1)
    ax2.tick_params(axis='y', which='both', left='off',
                    right='off', labelleft='off')
    ax2.tick_params(axis='x', which='both', top='off')

    if plot_title:
        ax1.set_ylabel('Estimated Mixtures')
        ax2.set_ylabel('Weight of each component')

# Parameters of the dataset
random_state, n_components, n_features = 2, 3, 2
colors = np.array(['#0072B2', '#F0E442', '#D55E00'])

covars = np.array([[[.7, .0], [.0, .1]],
                   [[.5, .0], [.0, .1]],
                   [[.5, .0], [.0, .1]]])
samples = np.array([200, 500, 200])
means = np.array([[.0, -.70],
                  [.0, .0],
                  [.0, .70]])

# mean_precision_prior= 0.8 to minimize the influence of the prior
estimators = [
    ("Finite mixture with a Dirichlet distribution\nprior and "
     r"$\gamma_0=$", BayesianGaussianMixture(
        weight_concentration_prior_type="dirichlet_distribution",
        n_components=2 * n_components, reg_covar=0, init_params='random',
        max_iter=1500, mean_precision_prior=.8,
        random_state=random_state), [0.001, 1, 1000]),
    ("Infinite mixture with a Dirichlet process\n prior and" r"$\gamma_0=$",
     BayesianGaussianMixture(
        weight_concentration_prior_type="dirichlet_process",
        n_components=2 * n_components, reg_covar=0, init_params='random',
        max_iter=1500, mean_precision_prior=.8,
        random_state=random_state), [1, 1000, 100000])]

# Generate data
rng = np.random.RandomState(random_state)
X = np.vstack([
    rng.multivariate_normal(means[j], covars[j], samples[j])
    for j in range(n_components)])
y = np.concatenate([j * np.ones(samples[j], dtype=int)
                    for j in range(n_components)])

# Plot results in two different figures
for (title, estimator, concentrations_prior) in estimators:
    plt.figure(figsize=(4.7 * 3, 8))
    plt.subplots_adjust(bottom=.04, top=0.90, hspace=.05, wspace=.05,
                        left=.03, right=.99)

    gs = gridspec.GridSpec(3, len(concentrations_prior))
    for k, concentration in enumerate(concentrations_prior):
        estimator.weight_concentration_prior = concentration
        estimator.fit(X)
        plot_results(plt.subplot(gs[0:2, k]), plt.subplot(gs[2, k]), estimator,
                     X, y, r"%s$%.1e$" % (title, concentration),
                     plot_title=k == 0)

plt.show()


# -*- coding: utf-8 -*-

from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function

from nose.tools import assert_raises
import sys
import pandas as pd
import numpy as np

from sktracker.tracker.cost_function import AbstractCostFunction

def test_abstract_cost_function():

    cost_func = AbstractCostFunction(context={}, parameters={})

    assert cost_func.get_block() == None

def test_abstract_cost_function_check_context():

    cost_func = AbstractCostFunction(context={'cost': 1}, parameters={})

    assert_raises(ValueError, cost_func.check_context, 'test_string', str)

    cost_func.context['test_string'] = 5

    assert_raises(TypeError, cost_func.check_context, 'test_string', str)

    cost_func.context['test_string'] = "i am a string"
    ### This fails in py2.7
    if sys.version_info[0] > 2:
        cost_func.check_context('test_string', str)

    assert True

def test_abstract_cost_function_check_columns():

    cost_func = AbstractCostFunction(context={}, parameters={})

    df = pd.DataFrame([np.arange(0, 5), np.arange(20, 25)],
                      columns=['x', 'y', 'z', 'w', 't'])

    cost_func.check_columns(df, ['t', 'z', 'y'])
    cost_func.check_columns([df], ['t', 'z', 'y'])

    df = pd.DataFrame([np.arange(0, 4), np.arange(20, 24)],
                      columns=['x', 'y', 'w', 't'])

    assert_raises(ValueError, cost_func.check_columns, df, ['t', 'z', 'y'])


import numpy as np

from brew.base import Ensemble
from brew.metrics.diversity.paired import kuncheva_double_fault_measure
from .base import DCS

class DSKNN(DCS):
    """DS-KNN

    The DS-KNN selects an ensemble of classifiers based on
    their accuracy and diversity in the neighborhood of the
    test sample.

    Attributes
    ----------
    `Xval` : array-like, shape = [indeterminated, n_features]
        Validation set.

    `yval` : array-like, shape = [indeterminated]
        Labels of the validation set.

    `knn`  : sklearn KNeighborsClassifier,
        Classifier used to find neighborhood.

    Examples
    --------
    >>> from brew.selection.dynamic import DSKNN
    >>> from brew.generation.bagging import Bagging
    >>> from brew.base import EnsembleClassifier
    >>>
    >>> from sklearn.tree import DecisionTreeClassifier
    >>> import numpy as np
    >>>
    >>> X = np.array([[-1, 0], [-0.8, 1], [-0.8, -1], [-0.5, 0],
                      [0.5, 0], [1, 0], [0.8, 1], [0.8, -1]])
    >>> y = np.array([1, 1, 1, 2, 1, 2, 2, 2])
    >>> tree = DecisionTreeClassifier(max_depth=1, min_samples_leaf=1)
    >>> bag = Bagging(base_classifier=tree, n_classifiers=10)
    >>> bag.fit(X, y)
    >>>
    >>> sel = DSKNN(X, y, K=3)
    >>>
    >>> clf = EnsembleClassifier(bag.ensemble, selector=sel)
    >>> clf.predict([-1.1,-0.5])
    [1]

    See also
    --------
    brew.selection.dynamic.lca.OLA: Overall Local Accuracy.
    brew.selection.dynamic.lca.LCA: Local Class Accuracy.

    References
    ----------
    Santana, Alixandre, et al. "A dynamic classifier selection method
    to build ensembles using accuracy and diversity." 2006 Ninth
    Brazilian Symposium on Neural Networks (SBRN'06). IEEE, 2006.
    """

    def __init__(self, Xval, yval, K=5, weighted=False, knn=None,
                 n_1=0.7, n_2=0.3):
        if n_1 < 0 or n_2 < 0 or n_1 <= n_2:
            raise Exception

        self.n_1 = n_1
        self.n_2 = n_2
        super(DSKNN, self).__init__(
            Xval, yval, K=K, weighted=weighted, knn=knn)

    def select(self, ensemble, x):
        if ensemble.in_agreement(x):
            return Ensemble([ensemble.classifiers[0]]), None

        n_sel_1, n_sel_2 = self.n_1, self.n_2
        if isinstance(self.n_1, float):
            n_sel_1 = int(n_sel_1 * len(ensemble))

        if isinstance(self.n_2, float):
            n_sel_2 = int(n_sel_2 * len(ensemble))

        n_sel_1 = max(n_sel_1, 1)
        n_sel_2 = max(n_sel_2, 1)

        # intialize variables
        # the the indexes of the KNN of x
        classifiers = ensemble.classifiers
        [idx] = self.knn.kneighbors(x, return_distance=False)
        X, y = self.Xval[idx], self.yval[idx]

        acc_scores = np.array([clf.score(X, y) for clf in classifiers])

        out = ensemble.output(X, mode='labels')
        oracle = np.equal(out, y[:, np.newaxis])
        div_scores = np.zeros(len(ensemble), dtype=float)

        for i in range(len(ensemble)):
            tmp = []
            for j in range(len(ensemble)):
                if i != j:
                    d = kuncheva_double_fault_measure(oracle[:, [i, j]])
                    tmp.append(d)
            div_scores[i] = np.mean(tmp)

        z = zip(np.arange(len(ensemble)), acc_scores, div_scores)
        z = sorted(z, key=lambda e: e[1], reverse=True)[:n_sel_1]
        z = sorted(z, key=lambda e: e[2], reverse=False)[:n_sel_2]
        z = zip(*z)[0]

        classifiers = [classifiers[i] for i in z]
        return Ensemble(classifiers=classifiers), None


import numpy as np
from matplotlib import lines
from ...viewer.canvastools.base import CanvasToolBase, ToolHandles

__all__ = ['LineTool', 'ThickLineTool']

class LineTool(CanvasToolBase):
    """Widget for line selection in a plot.

    Parameters
    ----------
    manager : Viewer or PlotPlugin.
        Skimage viewer or plot plugin object.
    on_move : function
        Function called whenever a control handle is moved.
        This function must accept the end points of line as the only argument.
    on_release : function
        Function called whenever the control handle is released.
    on_enter : function
        Function called whenever the "enter" key is pressed.
    maxdist : float
        Maximum pixel distance allowed when selecting control handle.
    line_props : dict
        Properties for :class:`matplotlib.lines.Line2D`.
    handle_props : dict
        Marker properties for the handles (also see
        :class:`matplotlib.lines.Line2D`).

    Attributes
    ----------
    end_points : 2D array
        End points of line ((x1, y1), (x2, y2)).
    """
    def __init__(self, manager, on_move=None, on_release=None, on_enter=None,
                 maxdist=10, line_props=None, handle_props=None,
                 **kwargs):
        super(LineTool, self).__init__(manager, on_move=on_move,
                                       on_enter=on_enter,
                                       on_release=on_release, **kwargs)

        props = dict(color='r', linewidth=1, alpha=0.4, solid_capstyle='butt')
        props.update(line_props if line_props is not None else {})
        self.linewidth = props['linewidth']
        self.maxdist = maxdist
        self._active_pt = None

        x = (0, 0)
        y = (0, 0)
        self._end_pts = np.transpose([x, y])

        self._line = lines.Line2D(x, y, visible=False, animated=True, **props)
        self.ax.add_line(self._line)

        self._handles = ToolHandles(self.ax, x, y,
                                    marker_props=handle_props)
        self._handles.set_visible(False)
        self.artists = [self._line, self._handles.artist]

        if on_enter is None:
            def on_enter(pts):
                x, y = np.transpose(pts)
                print("length = %0.2f" %
                      np.sqrt(np.diff(x)**2 + np.diff(y)**2))
        self.callback_on_enter = on_enter
        self.manager.add_tool(self)

    @property
    def end_points(self):
        return self._end_pts.astype(int)

    @end_points.setter
    def end_points(self, pts):
        self._end_pts = np.asarray(pts)

        self._line.set_data(np.transpose(pts))
        self._handles.set_data(np.transpose(pts))
        self._line.set_linewidth(self.linewidth)

        self.set_visible(True)
        self.redraw()

    def hit_test(self, event):
        if event.button != 1 or not self.ax.in_axes(event):
            return False
        idx, px_dist = self._handles.closest(event.x, event.y)
        if px_dist < self.maxdist:
            self._active_pt = idx
            return True
        else:
            self._active_pt = None
            return False

    def on_mouse_press(self, event):
        self.set_visible(True)
        if self._active_pt is None:
            self._active_pt = 0
            x, y = event.xdata, event.ydata
            self._end_pts = np.array([[x, y], [x, y]])

    def on_mouse_release(self, event):
        if event.button != 1:
            return
        self._active_pt = None
        self.callback_on_release(self.geometry)
        self.redraw()

    def on_move(self, event):
        if event.button != 1 or self._active_pt is None:
            return
        if not self.ax.in_axes(event):
            return
        self.update(event.xdata, event.ydata)
        self.callback_on_move(self.geometry)

    def update(self, x=None, y=None):
        if x is not None:
            self._end_pts[self._active_pt, :] = x, y
        self.end_points = self._end_pts

    @property
    def geometry(self):
        return self.end_points

class ThickLineTool(LineTool):
    """Widget for line selection in a plot.

    The thickness of the line can be varied using the mouse scroll wheel, or
    with the '+' and '-' keys.

    Parameters
    ----------
    manager : Viewer or PlotPlugin.
        Skimage viewer or plot plugin object.
    on_move : function
        Function called whenever a control handle is moved.
        This function must accept the end points of line as the only argument.
    on_release : function
        Function called whenever the control handle is released.
    on_enter : function
        Function called whenever the "enter" key is pressed.
    on_change : function
        Function called whenever the line thickness is changed.
    maxdist : float
        Maximum pixel distance allowed when selecting control handle.
    line_props : dict
        Properties for :class:`matplotlib.lines.Line2D`.
    handle_props : dict
        Marker properties for the handles (also see
        :class:`matplotlib.lines.Line2D`).

    Attributes
    ----------
    end_points : 2D array
        End points of line ((x1, y1), (x2, y2)).
    """

    def __init__(self, manager, on_move=None, on_enter=None, on_release=None,
                 on_change=None, maxdist=10, line_props=None, handle_props=None):
        super(ThickLineTool, self).__init__(manager,
                                            on_move=on_move,
                                            on_enter=on_enter,
                                            on_release=on_release,
                                            maxdist=maxdist,
                                            line_props=line_props,
                                            handle_props=handle_props)

        if on_change is None:
            def on_change(*args):
                pass
        self.callback_on_change = on_change

    def on_scroll(self, event):
        if not event.inaxes:
            return
        if event.button == 'up':
            self._thicken_scan_line()
        elif event.button == 'down':
            self._shrink_scan_line()

    def on_key_press(self, event):
        if event.key == '+':
            self._thicken_scan_line()
        elif event.key == '-':
            self._shrink_scan_line()

    def _thicken_scan_line(self):
        self.linewidth += 1
        self.update()
        self.callback_on_change(self.geometry)

    def _shrink_scan_line(self):
        if self.linewidth > 1:
            self.linewidth -= 1
            self.update()
            self.callback_on_change(self.geometry)

if __name__ == '__main__':  # pragma: no cover
    from ... import data
    from ...viewer import ImageViewer

    image = data.camera()

    viewer = ImageViewer(image)
    h, w = image.shape

    line_tool = ThickLineTool(viewer)
    line_tool.end_points = ([w/3, h/2], [2*w/3, h/2])
    viewer.show()


'''
NAME
    NetCDF with Python
PURPOSE
    To demonstrate how to read and write data with NetCDF files using
    a NetCDF file from the NCEP/NCAR Reanalysis.
    Plotting using Matplotlib and Basemap is also shown.
PROGRAMMER(S)
    Chris Slocum
REVISION HISTORY
    20140320 -- Initial version created and posted online
    20140722 -- Added basic error handling to ncdump
                Thanks to K.-Michael Aye for highlighting the issue
REFERENCES
    netcdf4-python -- http://code.google.com/p/netcdf4-python/
    NCEP/NCAR Reanalysis -- Kalnay et al. 1996
        http://dx.doi.org/10.1175/1520-0477(1996)077<0437:TNYRP>2.0.CO;2
'''
import datetime as dt  # Python standard library datetime  module
import numpy as np
from netCDF4 import Dataset  # http://code.google.com/p/netcdf4-python/
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, addcyclic, shiftgrid

def ncdump(nc_fid, verb=True):
    '''
    ncdump outputs dimensions, variables and their attribute information.
    The information is similar to that of NCAR's ncdump utility.
    ncdump requires a valid instance of Dataset.

    Parameters
    ----------
    nc_fid : netCDF4.Dataset
        A netCDF4 dateset object
    verb : Boolean
        whether or not nc_attrs, nc_dims, and nc_vars are printed

    Returns
    -------
    nc_attrs : list
        A Python list of the NetCDF file global attributes
    nc_dims : list
        A Python list of the NetCDF file dimensions
    nc_vars : list
        A Python list of the NetCDF file variables
    '''
    def print_ncattr(key):
        """
        Prints the NetCDF file attributes for a given key

        Parameters
        ----------
        key : unicode
            a valid netCDF4.Dataset.variables key
        """
        try:
            print "\t\ttype:", repr(nc_fid.variables[key].dtype)
            for ncattr in nc_fid.variables[key].ncattrs():
                print '\t\t%s:' % ncattr,\
                      repr(nc_fid.variables[key].getncattr(ncattr))
        except KeyError:
            print "\t\tWARNING: %s does not contain variable attributes" % key

    # NetCDF global attributes
    nc_attrs = nc_fid.ncattrs()
    if verb:
        print "NetCDF Global Attributes:"
        for nc_attr in nc_attrs:
            print '\t%s:' % nc_attr, repr(nc_fid.getncattr(nc_attr))
    nc_dims = [dim for dim in nc_fid.dimensions]  # list of nc dimensions
    # Dimension shape information.
    if verb:
        print "NetCDF dimension information:"
        for dim in nc_dims:
            print "\tName:", dim
            print "\t\tsize:", len(nc_fid.dimensions[dim])
            print_ncattr(dim)
    # Variable information.
    nc_vars = [var for var in nc_fid.variables]  # list of nc variables
    if verb:
        print "NetCDF variable information:"
        for var in nc_vars:
            if var not in nc_dims:
                print '\tName:', var
                print "\t\tdimensions:", nc_fid.variables[var].dimensions
                print "\t\tsize:", nc_fid.variables[var].size
                print_ncattr(var)
    return nc_attrs, nc_dims, nc_vars

nc_f = './CLM45_Micro_UW_SRF.2005120100.for.test.nc'  # Your filename
nc_fid = Dataset(nc_f, 'r')  # Dataset is the class behavior to open the file
                             # and create an instance of the ncCDF4 class
nc_attrs, nc_dims, nc_vars = ncdump(nc_fid)
# Extract data from NetCDF file
lats = nc_fid.variables['xlat'][:]  # extract/copy the data
lons = nc_fid.variables['xlon'][:]
time = nc_fid.variables['time'][:]
rsds = nc_fid.variables['rsds'][:]  # shape is time, lat, lon as shown above

time_idx = 237  # some random day in 2012
# Python and the renalaysis are slightly off in time so this fixes that problem
offset = dt.timedelta(hours=48)
# List of all times in the file as datetime objects
dt_time = [dt.date(1, 1, 1) + dt.timedelta(hours=t/20) - offset\
           for t in time]
cur_time = dt_time[time_idx]

# Plot of global temperature on our random day
fig = plt.figure()
fig.subplots_adjust(left=0., right=1., bottom=0., top=0.9)
# Setup the map. See http://matplotlib.org/basemap/users/mapsetup.html
# for other projections.
m = Basemap(projection='moll', llcrnrlat=-90, urcrnrlat=90,\
            llcrnrlon=0, urcrnrlon=360, resolution='c', lon_0=0)
m.drawcoastlines()
m.drawmapboundary()
# Make the plot continuous

test=rsds[0,:,:]
print test.shape
print rsds.shape
print lons.shape

rsds_cyclic, lons_cyclic = addcyclic(rsds[time_idx,:,:], lons)
# Shift the grid so lons go from -180 to 180 instead of 0 to 360.
rsds_cyclic, lons_cyclic = shiftgrid(180., rsds_cyclic, lons_cyclic, start=False)
# Create 2D lat/lon arrays for Basemap
lon2d, lat2d = np.meshgrid(lons_cyclic, lats)
# Transforms lat/lon into plotting coordinates for projection
x, y = m(lon2d, lat2d)
# Plot of rsds temperature with 11 contour intervals
cs = m.contourf(x, y, rsds_cyclic, 11, cmap=plt.cm.Spectral_r)
cbar = plt.colorbar(cs, orientation='horizontal', shrink=0.5)
cbar.set_label("%s (%s)" % (nc_fid.variables['rsds'].var_desc,\
                            nc_fid.variables['rsds'].units))
plt.title("%s on %s" % (nc_fid.variables['rsds'].var_desc, cur_time))

# Writing NetCDF files
# For this example, we will create two NetCDF4 files. One with the global rsds
# temperature departure from its value at Darwin, Australia. The other with
# the temperature profile for the entire year at Darwin.
darwin = {'name': 'Darwin, Australia', 'lat': -12.45, 'lon': 130.83}

# Find the nearest latitude and longitude for Darwin
lat_idx = np.abs(lats - darwin['lat']).argmin()
lon_idx = np.abs(lons - darwin['lon']).argmin()

# Simple example: temperature profile for the entire year at Darwin.
# Open a new NetCDF file to write the data to. For format, you can choose from
# 'NETCDF3_CLASSIC', 'NETCDF3_64BIT', 'NETCDF4_CLASSIC', and 'NETCDF4'
w_nc_fid = Dataset('darwin_2012.nc', 'w', format='NETCDF4')
w_nc_fid.description = "NCEP/NCAR Reanalysis %s from its value at %s. %s" %\
                      (nc_fid.variables['rsds'].var_desc.lower(),\
                       darwin['name'], nc_fid.description)
# Using our previous dimension info, we can create the new time dimension
# Even though we know the size, we are going to set the size to unknown
w_nc_fid.createDimension('time', None)
w_nc_dim = w_nc_fid.createVariable('time', nc_fid.variables['time'].dtype,\
                                   ('time',))
# You can do this step yourself but someone else did the work for us.
for ncattr in nc_fid.variables['time'].ncattrs():
    w_nc_dim.setncattr(ncattr, nc_fid.variables['time'].getncattr(ncattr))
# Assign the dimension data to the new NetCDF file.
w_nc_fid.variables['time'][:] = time
w_nc_var = w_nc_fid.createVariable('rsds', 'f8', ('time'))
w_nc_var.setncatts({'long_name': u"mean Daily Air temperature",\
                    'units': u"degK", 'level_desc': u'Surface',\
                    'var_desc': u"Air temperature",\
                    'statistic': u'Mean\nM'})
w_nc_fid.variables['rsds'][:] = rsds[time_idx, lat_idx, lon_idx]
w_nc_fid.close()  # close the new file

# A plot of the temperature profile for Darwin in 2012
fig = plt.figure()
plt.plot(dt_time, rsds[:, lat_idx, lon_idx], c='r')
plt.plot(dt_time[time_idx], rsds[time_idx, lat_idx, lon_idx], c='b', marker='o')
plt.text(dt_time[time_idx], rsds[time_idx, lat_idx, lon_idx], cur_time,\
         ha='right')
fig.autofmt_xdate()
plt.ylabel("%s (%s)" % (nc_fid.variables['rsds'].var_desc,\
                        nc_fid.variables['rsds'].units))
plt.xlabel("Time")
plt.title("%s from\n%s for %s" % (nc_fid.variables['rsds'].var_desc,\
                                  darwin['name'], cur_time.year))

# Complex example: global temperature departure from its value at Darwin
departure = rsds[:, :, :] - rsds[:, lat_idx, lon_idx].reshape((time.shape[0],\
                                                             1, 1))

# Open a new NetCDF file to write the data to. For format, you can choose from
# 'NETCDF3_CLASSIC', 'NETCDF3_64BIT', 'NETCDF4_CLASSIC', and 'NETCDF4'
w_nc_fid = Dataset('rsds.departure.sig995.2012.nc', 'w', format='NETCDF4')
w_nc_fid.description = "The departure of the NCEP/NCAR Reanalysis " +\
                      "%s from its value at %s. %s" %\
                      (nc_fid.variables['rsds'].var_desc.lower(),\
                       darwin['name'], nc_fid.description)
# Using our previous dimension information, we can create the new dimensions
data = {}
for dim in nc_dims:
    w_nc_fid.createDimension(dim, nc_fid.variables[dim].size)
    data[dim] = w_nc_fid.createVariable(dim, nc_fid.variables[dim].dtype,\
                                        (dim,))
    # You can do this step yourself but someone else did the work for us.
    for ncattr in nc_fid.variables[dim].ncattrs():
        data[dim].setncattr(ncattr, nc_fid.variables[dim].getncattr(ncattr))
# Assign the dimension data to the new NetCDF file.
w_nc_fid.variables['time'][:] = time
w_nc_fid.variables['lat'][:] = lats
w_nc_fid.variables['lon'][:] = lons

# Ok, time to create our departure variable
w_nc_var = w_nc_fid.createVariable('rsds_dep', 'f8', ('time', 'lat', 'lon'))
w_nc_var.setncatts({'long_name': u"mean Daily Air temperature departure",\
                    'units': u"degK", 'level_desc': u'Surface',\
                    'var_desc': u"Air temperature departure",\
                    'statistic': u'Mean\nM'})
w_nc_fid.variables['rsds_dep'][:] = departure
w_nc_fid.close()  # close the new file

# Rounded maximum absolute value of the departure used for contouring
max_dep = np.round(np.abs(departure[time_idx, :, :]).max()+5., decimals=-1)

# Generate a figure of the departure for a single day
fig = plt.figure()
fig.subplots_adjust(left=0., right=1., bottom=0., top=0.9)
m = Basemap(projection='moll', llcrnrlat=-90, urcrnrlat=90,\
            llcrnrlon=0, urcrnrlon=360, resolution='c', lon_0=0)
m.drawcoastlines()
m.drawmapboundary()
dep_cyclic, lons_cyclic = addcyclic(departure[time_idx, :, :], lons)
dep_cyclic, lons_cyclic = shiftgrid(180., dep_cyclic, lons_cyclic, start=False)
lon2d, lat2d = np.meshgrid(lons_cyclic, lats)
x, y = m(lon2d, lat2d)
levels = np.linspace(-max_dep, max_dep, 11)
cs = m.contourf(x, y, dep_cyclic, levels=levels, cmap=plt.cm.bwr)
x, y = m(darwin['lon'], darwin['lat'])
plt.plot(x, y, c='c', marker='o')
plt.text(x, y, 'Darwin,\nAustralia', color='r', weight='semibold')
cbar = plt.colorbar(cs, orientation='horizontal', shrink=0.5)
cbar.set_label("%s departure (%s)" % (nc_fid.variables['rsds'].var_desc,\
                            nc_fid.variables['rsds'].units))
plt.title("Departure of Global %s from\n%s for %s" %\
          (nc_fid.variables['rsds'].var_desc, darwin['name'], cur_time))
plt.show()

# Close original NetCDF file.
nc_fid.close()


"""
The :mod:`sklearn.covariance` module includes methods and algorithms to
robustly estimate the covariance of features given a set of points. The
precision matrix defined as the inverse of the covariance is also estimated.
Covariance estimation is closely related to the theory of Gaussian Graphical
Models.
"""

from .empirical_covariance_ import empirical_covariance, EmpiricalCovariance, \
    log_likelihood
from .shrunk_covariance_ import shrunk_covariance, ShrunkCovariance, \
    ledoit_wolf, ledoit_wolf_shrinkage, \
    LedoitWolf, oas, OAS
from .robust_covariance import fast_mcd, MinCovDet
from .graph_lasso_ import graph_lasso, GraphLasso, GraphLassoCV
from .outlier_detection import EllipticEnvelope

__all__ = ['EllipticEnvelope',
           'EmpiricalCovariance',
           'GraphLasso',
           'GraphLassoCV',
           'LedoitWolf',
           'MinCovDet',
           'OAS',
           'ShrunkCovariance',
           'empirical_covariance',
           'fast_mcd',
           'graph_lasso',
           'ledoit_wolf',
           'ledoit_wolf_shrinkage',
           'log_likelihood',
           'oas',
           'shrunk_covariance']