import pickle
import unittest

import numpy as np
import scipy.sparse as sp

from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_greater
from sklearn.utils.testing import assert_less
from sklearn.utils.testing import raises
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_false, assert_true
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_raises_regexp

from sklearn import linear_model, datasets, metrics
from sklearn.base import clone
from sklearn.linear_model import SGDClassifier, SGDRegressor
from sklearn.preprocessing import LabelEncoder, scale, MinMaxScaler

class SparseSGDClassifier(SGDClassifier):

    def fit(self, X, y, *args, **kw):
        X = sp.csr_matrix(X)
        return super(SparseSGDClassifier, self).fit(X, y, *args, **kw)

    def partial_fit(self, X, y, *args, **kw):
        X = sp.csr_matrix(X)
        return super(SparseSGDClassifier, self).partial_fit(X, y, *args, **kw)

    def decision_function(self, X):
        X = sp.csr_matrix(X)
        return super(SparseSGDClassifier, self).decision_function(X)

    def predict_proba(self, X):
        X = sp.csr_matrix(X)
        return super(SparseSGDClassifier, self).predict_proba(X)

class SparseSGDRegressor(SGDRegressor):

    def fit(self, X, y, *args, **kw):
        X = sp.csr_matrix(X)
        return SGDRegressor.fit(self, X, y, *args, **kw)

    def partial_fit(self, X, y, *args, **kw):
        X = sp.csr_matrix(X)
        return SGDRegressor.partial_fit(self, X, y, *args, **kw)

    def decision_function(self, X, *args, **kw):
        X = sp.csr_matrix(X)
        return SGDRegressor.decision_function(self, X, *args, **kw)

# Test Data

# test sample 1
X = np.array([[-2, -1], [-1, -1], [-1, -2], [1, 1], [1, 2], [2, 1]])
Y = [1, 1, 1, 2, 2, 2]
T = np.array([[-1, -1], [2, 2], [3, 2]])
true_result = [1, 2, 2]

# test sample 2; string class labels
X2 = np.array([[-1, 1], [-0.75, 0.5], [-1.5, 1.5],
               [1, 1], [0.75, 0.5], [1.5, 1.5],
               [-1, -1], [0, -0.5], [1, -1]])
Y2 = ["one"] * 3 + ["two"] * 3 + ["three"] * 3
T2 = np.array([[-1.5, 0.5], [1, 2], [0, -2]])
true_result2 = ["one", "two", "three"]

# test sample 3
X3 = np.array([[1, 1, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0],
               [0, 0, 1, 0, 0, 0], [0, 0, 1, 0, 0, 0],
               [0, 0, 0, 0, 1, 1], [0, 0, 0, 0, 1, 1],
               [0, 0, 0, 1, 0, 0], [0, 0, 0, 1, 0, 0]])
Y3 = np.array([1, 1, 1, 1, 2, 2, 2, 2])

# test sample 4 - two more or less redundent feature groups
X4 = np.array([[1, 0.9, 0.8, 0, 0, 0], [1, .84, .98, 0, 0, 0],
               [1, .96, .88, 0, 0, 0], [1, .91, .99, 0, 0, 0],
               [0, 0, 0, .89, .91, 1], [0, 0, 0, .79, .84, 1],
               [0, 0, 0, .91, .95, 1], [0, 0, 0, .93, 1, 1]])
Y4 = np.array([1, 1, 1, 1, 2, 2, 2, 2])

iris = datasets.load_iris()

# test sample 5 - test sample 1 as binary classification problem
X5 = np.array([[-2, -1], [-1, -1], [-1, -2], [1, 1], [1, 2], [2, 1]])
Y5 = [1, 1, 1, 2, 2, 2]
true_result5 = [0, 1, 1]

# Classification Test Case

class CommonTest(object):

    def factory(self, **kwargs):
        if "random_state" not in kwargs:
            kwargs["random_state"] = 42
        return self.factory_class(**kwargs)

    # a simple implementation of ASGD to use for testing
    # uses squared loss to find the gradient
    def asgd(self, X, y, eta, alpha, weight_init=None, intercept_init=0.0):
        if weight_init is None:
            weights = np.zeros(X.shape[1])
        else:
            weights = weight_init

        average_weights = np.zeros(X.shape[1])
        intercept = intercept_init
        average_intercept = 0.0
        decay = 1.0

        # sparse data has a fixed decay of .01
        if (isinstance(self, SparseSGDClassifierTestCase) or
                isinstance(self, SparseSGDRegressorTestCase)):
            decay = .01

        for i, entry in enumerate(X):
            p = np.dot(entry, weights)
            p += intercept
            gradient = p - y[i]
            weights *= 1.0 - (eta * alpha)
            weights += -(eta * gradient * entry)
            intercept += -(eta * gradient) * decay

            average_weights *= i
            average_weights += weights
            average_weights /= i + 1.0

            average_intercept *= i
            average_intercept += intercept
            average_intercept /= i + 1.0

        return average_weights, average_intercept

    def _test_warm_start(self, X, Y, lr):
        # Test that explicit warm restart...
        clf = self.factory(alpha=0.01, eta0=0.01, n_iter=5, shuffle=False,
                           learning_rate=lr)
        clf.fit(X, Y)

        clf2 = self.factory(alpha=0.001, eta0=0.01, n_iter=5, shuffle=False,
                            learning_rate=lr)
        clf2.fit(X, Y,
                 coef_init=clf.coef_.copy(),
                 intercept_init=clf.intercept_.copy())

        # ... and implicit warm restart are equivalent.
        clf3 = self.factory(alpha=0.01, eta0=0.01, n_iter=5, shuffle=False,
                            warm_start=True, learning_rate=lr)
        clf3.fit(X, Y)

        assert_equal(clf3.t_, clf.t_)
        assert_array_almost_equal(clf3.coef_, clf.coef_)

        clf3.set_params(alpha=0.001)
        clf3.fit(X, Y)

        assert_equal(clf3.t_, clf2.t_)
        assert_array_almost_equal(clf3.coef_, clf2.coef_)

    def test_warm_start_constant(self):
        self._test_warm_start(X, Y, "constant")

    def test_warm_start_invscaling(self):
        self._test_warm_start(X, Y, "invscaling")

    def test_warm_start_optimal(self):
        self._test_warm_start(X, Y, "optimal")

    def test_input_format(self):
        # Input format tests.
        clf = self.factory(alpha=0.01, n_iter=5,
                           shuffle=False)
        clf.fit(X, Y)
        Y_ = np.array(Y)[:, np.newaxis]

        Y_ = np.c_[Y_, Y_]
        assert_raises(ValueError, clf.fit, X, Y_)

    def test_clone(self):
        # Test whether clone works ok.
        clf = self.factory(alpha=0.01, n_iter=5, penalty='l1')
        clf = clone(clf)
        clf.set_params(penalty='l2')
        clf.fit(X, Y)

        clf2 = self.factory(alpha=0.01, n_iter=5, penalty='l2')
        clf2.fit(X, Y)

        assert_array_equal(clf.coef_, clf2.coef_)

    def test_plain_has_no_average_attr(self):
        clf = self.factory(average=True, eta0=.01)
        clf.fit(X, Y)

        assert_true(hasattr(clf, 'average_coef_'))
        assert_true(hasattr(clf, 'average_intercept_'))
        assert_true(hasattr(clf, 'standard_intercept_'))
        assert_true(hasattr(clf, 'standard_coef_'))

        clf = self.factory()
        clf.fit(X, Y)

        assert_false(hasattr(clf, 'average_coef_'))
        assert_false(hasattr(clf, 'average_intercept_'))
        assert_false(hasattr(clf, 'standard_intercept_'))
        assert_false(hasattr(clf, 'standard_coef_'))

    def test_late_onset_averaging_not_reached(self):
        clf1 = self.factory(average=600)
        clf2 = self.factory()
        for _ in range(100):
            if isinstance(clf1, SGDClassifier):
                clf1.partial_fit(X, Y, classes=np.unique(Y))
                clf2.partial_fit(X, Y, classes=np.unique(Y))
            else:
                clf1.partial_fit(X, Y)
                clf2.partial_fit(X, Y)

        assert_array_almost_equal(clf1.coef_, clf2.coef_, decimal=16)
        assert_almost_equal(clf1.intercept_, clf2.intercept_, decimal=16)

    def test_late_onset_averaging_reached(self):
        eta0 = .001
        alpha = .0001
        Y_encode = np.array(Y)
        Y_encode[Y_encode == 1] = -1.0
        Y_encode[Y_encode == 2] = 1.0

        clf1 = self.factory(average=7, learning_rate="constant",
                            loss='squared_loss', eta0=eta0,
                            alpha=alpha, n_iter=2, shuffle=False)
        clf2 = self.factory(average=0, learning_rate="constant",
                            loss='squared_loss', eta0=eta0,
                            alpha=alpha, n_iter=1, shuffle=False)

        clf1.fit(X, Y_encode)
        clf2.fit(X, Y_encode)

        average_weights, average_intercept = \
            self.asgd(X, Y_encode, eta0, alpha,
                      weight_init=clf2.coef_.ravel(),
                      intercept_init=clf2.intercept_)

        assert_array_almost_equal(clf1.coef_.ravel(),
                                  average_weights.ravel(),
                                  decimal=16)
        assert_almost_equal(clf1.intercept_, average_intercept, decimal=16)

class DenseSGDClassifierTestCase(unittest.TestCase, CommonTest):
    """Test suite for the dense representation variant of SGD"""
    factory_class = SGDClassifier
    def test_sgd(self):
        # Check that SGD gives any results :-)

        for loss in ("hinge", "squared_hinge", "log", "modified_huber"):
            clf = self.factory(penalty='l2', alpha=0.01, fit_intercept=True,
                               loss=loss, n_iter=10, shuffle=True)
            clf.fit(X, Y)
            # assert_almost_equal(clf.coef_[0], clf.coef_[1], decimal=7)
            assert_array_equal(clf.predict(T), true_result)

    @raises(ValueError)
    def test_sgd_bad_l1_ratio(self):
        # Check whether expected ValueError on bad l1_ratio
        self.factory(l1_ratio=1.1)

    @raises(ValueError)
    def test_sgd_bad_learning_rate_schedule(self):
        # Check whether expected ValueError on bad learning_rate
        self.factory(learning_rate="<unknown>")

    @raises(ValueError)
    def test_sgd_bad_eta0(self):
        # Check whether expected ValueError on bad eta0
        self.factory(eta0=0, learning_rate="constant")

    @raises(ValueError)
    def test_sgd_bad_alpha(self):
        # Check whether expected ValueError on bad alpha
        self.factory(alpha=-.1)

    @raises(ValueError)
    def test_sgd_bad_penalty(self):
        # Check whether expected ValueError on bad penalty
        self.factory(penalty='foobar', l1_ratio=0.85)

    @raises(ValueError)
    def test_sgd_bad_loss(self):
        # Check whether expected ValueError on bad loss
        self.factory(loss="foobar")

    @raises(ValueError)
    def test_sgd_n_iter_param(self):
        # Test parameter validity check
        self.factory(n_iter=-10000)

    @raises(ValueError)
    def test_sgd_shuffle_param(self):
        # Test parameter validity check
        self.factory(shuffle="false")

    @raises(TypeError)
    def test_argument_coef(self):
        # Checks coef_init not allowed as model argument (only fit)
        # Provided coef_ does not match dataset.
        self.factory(coef_init=np.zeros((3,))).fit(X, Y)

    @raises(ValueError)
    def test_provide_coef(self):
        # Checks coef_init shape for the warm starts
        # Provided coef_ does not match dataset.
        self.factory().fit(X, Y, coef_init=np.zeros((3,)))

    @raises(ValueError)
    def test_set_intercept(self):
        # Checks intercept_ shape for the warm starts
        # Provided intercept_ does not match dataset.
        self.factory().fit(X, Y, intercept_init=np.zeros((3,)))

    def test_set_intercept_binary(self):
        # Checks intercept_ shape for the warm starts in binary case
        self.factory().fit(X5, Y5, intercept_init=0)

    def test_average_binary_computed_correctly(self):
        # Checks the SGDClassifier correctly computes the average weights
        eta = .1
        alpha = 2.
        n_samples = 20
        n_features = 10
        rng = np.random.RandomState(0)
        X = rng.normal(size=(n_samples, n_features))
        w = rng.normal(size=n_features)

        clf = self.factory(loss='squared_loss',
                           learning_rate='constant',
                           eta0=eta, alpha=alpha,
                           fit_intercept=True,
                           n_iter=1, average=True, shuffle=False)

        # simple linear function without noise
        y = np.dot(X, w)
        y = np.sign(y)

        clf.fit(X, y)

        average_weights, average_intercept = self.asgd(X, y, eta, alpha)
        average_weights = average_weights.reshape(1, -1)
        assert_array_almost_equal(clf.coef_,
                                  average_weights,
                                  decimal=14)
        assert_almost_equal(clf.intercept_, average_intercept, decimal=14)

    def test_set_intercept_to_intercept(self):
        # Checks intercept_ shape consistency for the warm starts
        # Inconsistent intercept_ shape.
        clf = self.factory().fit(X5, Y5)
        self.factory().fit(X5, Y5, intercept_init=clf.intercept_)
        clf = self.factory().fit(X, Y)
        self.factory().fit(X, Y, intercept_init=clf.intercept_)

    @raises(ValueError)
    def test_sgd_at_least_two_labels(self):
        # Target must have at least two labels
        self.factory(alpha=0.01, n_iter=20).fit(X2, np.ones(9))

    def test_partial_fit_weight_class_auto(self):
        # partial_fit with class_weight='auto' not supported
        assert_raises_regexp(ValueError,
                             "class_weight 'auto' is not supported for "
                             "partial_fit. In order to use 'auto' weights, "
                             "use compute_class_weight\('auto', classes, y\). "
                             "In place of y you can us a large enough sample "
                             "of the full training set target to properly "
                             "estimate the class frequency distributions. "
                             "Pass the resulting weights as the class_weight "
                             "parameter.",
                             self.factory(class_weight='auto').partial_fit,
                             X, Y, classes=np.unique(Y))

    def test_sgd_multiclass(self):
        # Multi-class test case
        clf = self.factory(alpha=0.01, n_iter=20).fit(X2, Y2)
        assert_equal(clf.coef_.shape, (3, 2))
        assert_equal(clf.intercept_.shape, (3,))
        assert_equal(clf.decision_function([0, 0]).shape, (1, 3))
        pred = clf.predict(T2)
        assert_array_equal(pred, true_result2)

    def test_sgd_multiclass_average(self):
        eta = .001
        alpha = .01
        # Multi-class average test case
        clf = self.factory(loss='squared_loss',
                           learning_rate='constant',
                           eta0=eta, alpha=alpha,
                           fit_intercept=True,
                           n_iter=1, average=True, shuffle=False)

        np_Y2 = np.array(Y2)
        clf.fit(X2, np_Y2)
        classes = np.unique(np_Y2)

        for i, cl in enumerate(classes):
            y_i = np.ones(np_Y2.shape[0])
            y_i[np_Y2 != cl] = -1
            average_coef, average_intercept = self.asgd(X2, y_i, eta, alpha)
            assert_array_almost_equal(average_coef, clf.coef_[i], decimal=16)
            assert_almost_equal(average_intercept,
                                clf.intercept_[i],
                                decimal=16)

    def test_sgd_multiclass_with_init_coef(self):
        # Multi-class test case
        clf = self.factory(alpha=0.01, n_iter=20)
        clf.fit(X2, Y2, coef_init=np.zeros((3, 2)),
                intercept_init=np.zeros(3))
        assert_equal(clf.coef_.shape, (3, 2))
        assert_true(clf.intercept_.shape, (3,))
        pred = clf.predict(T2)
        assert_array_equal(pred, true_result2)

    def test_sgd_multiclass_njobs(self):
        # Multi-class test case with multi-core support
        clf = self.factory(alpha=0.01, n_iter=20, n_jobs=2).fit(X2, Y2)
        assert_equal(clf.coef_.shape, (3, 2))
        assert_equal(clf.intercept_.shape, (3,))
        assert_equal(clf.decision_function([0, 0]).shape, (1, 3))
        pred = clf.predict(T2)
        assert_array_equal(pred, true_result2)

    def test_set_coef_multiclass(self):
        # Checks coef_init and intercept_init shape for for multi-class
        # problems
        # Provided coef_ does not match dataset
        clf = self.factory()
        assert_raises(ValueError, clf.fit, X2, Y2, coef_init=np.zeros((2, 2)))

        # Provided coef_ does match dataset
        clf = self.factory().fit(X2, Y2, coef_init=np.zeros((3, 2)))

        # Provided intercept_ does not match dataset
        clf = self.factory()
        assert_raises(ValueError, clf.fit, X2, Y2,
                      intercept_init=np.zeros((1,)))

        # Provided intercept_ does match dataset.
        clf = self.factory().fit(X2, Y2, intercept_init=np.zeros((3,)))

    def test_sgd_proba(self):
        # Check SGD.predict_proba

        # Hinge loss does not allow for conditional prob estimate.
        # We cannot use the factory here, because it defines predict_proba
        # anyway.
        clf = SGDClassifier(loss="hinge", alpha=0.01, n_iter=10).fit(X, Y)
        assert_false(hasattr(clf, "predict_proba"))
        assert_false(hasattr(clf, "predict_log_proba"))

        # log and modified_huber losses can output probability estimates
        # binary case
        for loss in ["log", "modified_huber"]:
            clf = self.factory(loss="modified_huber", alpha=0.01, n_iter=10)
            clf.fit(X, Y)
            p = clf.predict_proba([3, 2])
            assert_true(p[0, 1] > 0.5)
            p = clf.predict_proba([-1, -1])
            assert_true(p[0, 1] < 0.5)

            p = clf.predict_log_proba([3, 2])
            assert_true(p[0, 1] > p[0, 0])
            p = clf.predict_log_proba([-1, -1])
            assert_true(p[0, 1] < p[0, 0])

        # log loss multiclass probability estimates
        clf = self.factory(loss="log", alpha=0.01, n_iter=10).fit(X2, Y2)

        d = clf.decision_function([[.1, -.1], [.3, .2]])
        p = clf.predict_proba([[.1, -.1], [.3, .2]])
        assert_array_equal(np.argmax(p, axis=1), np.argmax(d, axis=1))
        assert_almost_equal(p[0].sum(), 1)
        assert_true(np.all(p[0] >= 0))

        p = clf.predict_proba([-1, -1])
        d = clf.decision_function([-1, -1])
        assert_array_equal(np.argsort(p[0]), np.argsort(d[0]))

        l = clf.predict_log_proba([3, 2])
        p = clf.predict_proba([3, 2])
        assert_array_almost_equal(np.log(p), l)

        l = clf.predict_log_proba([-1, -1])
        p = clf.predict_proba([-1, -1])
        assert_array_almost_equal(np.log(p), l)

        # Modified Huber multiclass probability estimates; requires a separate
        # test because the hard zero/one probabilities may destroy the
        # ordering present in decision_function output.
        clf = self.factory(loss="modified_huber", alpha=0.01, n_iter=10)
        clf.fit(X2, Y2)
        d = clf.decision_function([3, 2])
        p = clf.predict_proba([3, 2])
        if not isinstance(self, SparseSGDClassifierTestCase):
            assert_equal(np.argmax(d, axis=1), np.argmax(p, axis=1))
        else:   # XXX the sparse test gets a different X2 (?)
            assert_equal(np.argmin(d, axis=1), np.argmin(p, axis=1))

        # the following sample produces decision_function values < -1,
        # which would cause naive normalization to fail (see comment
        # in SGDClassifier.predict_proba)
        x = X.mean(axis=0)
        d = clf.decision_function(x)
        if np.all(d < -1):  # XXX not true in sparse test case (why?)
            p = clf.predict_proba(x)
            assert_array_almost_equal(p[0], [1 / 3.] * 3)

    def test_sgd_l1(self):
        # Test L1 regularization
        n = len(X4)
        rng = np.random.RandomState(13)
        idx = np.arange(n)
        rng.shuffle(idx)

        X = X4[idx, :]
        Y = Y4[idx]

        clf = self.factory(penalty='l1', alpha=.2, fit_intercept=False,
                           n_iter=2000, shuffle=False)
        clf.fit(X, Y)
        assert_array_equal(clf.coef_[0, 1:-1], np.zeros((4,)))
        pred = clf.predict(X)
        assert_array_equal(pred, Y)

        # test sparsify with dense inputs
        clf.sparsify()
        assert_true(sp.issparse(clf.coef_))
        pred = clf.predict(X)
        assert_array_equal(pred, Y)

        # pickle and unpickle with sparse coef_
        clf = pickle.loads(pickle.dumps(clf))
        assert_true(sp.issparse(clf.coef_))
        pred = clf.predict(X)
        assert_array_equal(pred, Y)

    def test_class_weights(self):
        # Test class weights.
        X = np.array([[-1.0, -1.0], [-1.0, 0], [-.8, -1.0],
                      [1.0, 1.0], [1.0, 0.0]])
        y = [1, 1, 1, -1, -1]

        clf = self.factory(alpha=0.1, n_iter=1000, fit_intercept=False,
                           class_weight=None)
        clf.fit(X, y)
        assert_array_equal(clf.predict([[0.2, -1.0]]), np.array([1]))

        # we give a small weights to class 1
        clf = self.factory(alpha=0.1, n_iter=1000, fit_intercept=False,
                           class_weight={1: 0.001})
        clf.fit(X, y)

        # now the hyperplane should rotate clock-wise and
        # the prediction on this point should shift
        assert_array_equal(clf.predict([[0.2, -1.0]]), np.array([-1]))

    def test_equal_class_weight(self):
        # Test if equal class weights approx. equals no class weights.
        X = [[1, 0], [1, 0], [0, 1], [0, 1]]
        y = [0, 0, 1, 1]
        clf = self.factory(alpha=0.1, n_iter=1000, class_weight=None)
        clf.fit(X, y)

        X = [[1, 0], [0, 1]]
        y = [0, 1]
        clf_weighted = self.factory(alpha=0.1, n_iter=1000,
                                    class_weight={0: 0.5, 1: 0.5})
        clf_weighted.fit(X, y)

        # should be similar up to some epsilon due to learning rate schedule
        assert_almost_equal(clf.coef_, clf_weighted.coef_, decimal=2)

    @raises(ValueError)
    def test_wrong_class_weight_label(self):
        # ValueError due to not existing class label.
        clf = self.factory(alpha=0.1, n_iter=1000, class_weight={0: 0.5})
        clf.fit(X, Y)

    @raises(ValueError)
    def test_wrong_class_weight_format(self):
        # ValueError due to wrong class_weight argument type.
        clf = self.factory(alpha=0.1, n_iter=1000, class_weight=[0.5])
        clf.fit(X, Y)

    def test_weights_multiplied(self):
        # Tests that class_weight and sample_weight are multiplicative
        class_weights = {1: .6, 2: .3}
        sample_weights = np.random.random(Y4.shape[0])
        multiplied_together = np.copy(sample_weights)
        multiplied_together[Y4 == 1] *= class_weights[1]
        multiplied_together[Y4 == 2] *= class_weights[2]

        clf1 = self.factory(alpha=0.1, n_iter=20, class_weight=class_weights)
        clf2 = self.factory(alpha=0.1, n_iter=20)

        clf1.fit(X4, Y4, sample_weight=sample_weights)
        clf2.fit(X4, Y4, sample_weight=multiplied_together)

        assert_almost_equal(clf1.coef_, clf2.coef_)

    def test_auto_weight(self):
        # Test class weights for imbalanced data
        # compute reference metrics on iris dataset that is quite balanced by
        # default
        X, y = iris.data, iris.target
        X = scale(X)
        idx = np.arange(X.shape[0])
        rng = np.random.RandomState(6)
        rng.shuffle(idx)
        X = X[idx]
        y = y[idx]
        clf = self.factory(alpha=0.0001, n_iter=1000,
                           class_weight=None, shuffle=False).fit(X, y)
        assert_almost_equal(metrics.f1_score(y, clf.predict(X), average='weighted'), 0.96,
                            decimal=1)

        # make the same prediction using automated class_weight
        clf_auto = self.factory(alpha=0.0001, n_iter=1000,
                                class_weight="auto", shuffle=False).fit(X, y)
        assert_almost_equal(metrics.f1_score(y, clf_auto.predict(X), average='weighted'), 0.96,
                            decimal=1)

        # Make sure that in the balanced case it does not change anything
        # to use "auto"
        assert_array_almost_equal(clf.coef_, clf_auto.coef_, 6)

        # build an very very imbalanced dataset out of iris data
        X_0 = X[y == 0, :]
        y_0 = y[y == 0]

        X_imbalanced = np.vstack([X] + [X_0] * 10)
        y_imbalanced = np.concatenate([y] + [y_0] * 10)

        # fit a model on the imbalanced data without class weight info
        clf = self.factory(n_iter=1000, class_weight=None, shuffle=False)
        clf.fit(X_imbalanced, y_imbalanced)
        y_pred = clf.predict(X)
        assert_less(metrics.f1_score(y, y_pred, average='weighted'), 0.96)

        # fit a model with auto class_weight enabled
        clf = self.factory(n_iter=1000, class_weight="auto", shuffle=False)
        clf.fit(X_imbalanced, y_imbalanced)
        y_pred = clf.predict(X)
        assert_greater(metrics.f1_score(y, y_pred, average='weighted'), 0.96)

        # fit another using a fit parameter override
        clf = self.factory(n_iter=1000, class_weight="auto", shuffle=False)
        clf.fit(X_imbalanced, y_imbalanced)
        y_pred = clf.predict(X)
        assert_greater(metrics.f1_score(y, y_pred, average='weighted'), 0.96)

    def test_sample_weights(self):
        # Test weights on individual samples
        X = np.array([[-1.0, -1.0], [-1.0, 0], [-.8, -1.0],
                      [1.0, 1.0], [1.0, 0.0]])
        y = [1, 1, 1, -1, -1]

        clf = self.factory(alpha=0.1, n_iter=1000, fit_intercept=False)
        clf.fit(X, y)
        assert_array_equal(clf.predict([[0.2, -1.0]]), np.array([1]))

        # we give a small weights to class 1
        clf.fit(X, y, sample_weight=[0.001] * 3 + [1] * 2)

        # now the hyperplane should rotate clock-wise and
        # the prediction on this point should shift
        assert_array_equal(clf.predict([[0.2, -1.0]]), np.array([-1]))

    @raises(ValueError)
    def test_wrong_sample_weights(self):
        # Test if ValueError is raised if sample_weight has wrong shape
        clf = self.factory(alpha=0.1, n_iter=1000, fit_intercept=False)
        # provided sample_weight too long
        clf.fit(X, Y, sample_weight=np.arange(7))

    @raises(ValueError)
    def test_partial_fit_exception(self):
        clf = self.factory(alpha=0.01)
        # classes was not specified
        clf.partial_fit(X3, Y3)

    def test_partial_fit_binary(self):
        third = X.shape[0] // 3
        clf = self.factory(alpha=0.01)
        classes = np.unique(Y)

        clf.partial_fit(X[:third], Y[:third], classes=classes)
        assert_equal(clf.coef_.shape, (1, X.shape[1]))
        assert_equal(clf.intercept_.shape, (1,))
        assert_equal(clf.decision_function([0, 0]).shape, (1, ))
        id1 = id(clf.coef_.data)

        clf.partial_fit(X[third:], Y[third:])
        id2 = id(clf.coef_.data)
        # check that coef_ haven't been re-allocated
        assert_true(id1, id2)

        y_pred = clf.predict(T)
        assert_array_equal(y_pred, true_result)

    def test_partial_fit_multiclass(self):
        third = X2.shape[0] // 3
        clf = self.factory(alpha=0.01)
        classes = np.unique(Y2)

        clf.partial_fit(X2[:third], Y2[:third], classes=classes)
        assert_equal(clf.coef_.shape, (3, X2.shape[1]))
        assert_equal(clf.intercept_.shape, (3,))
        assert_equal(clf.decision_function([0, 0]).shape, (1, 3))
        id1 = id(clf.coef_.data)

        clf.partial_fit(X2[third:], Y2[third:])
        id2 = id(clf.coef_.data)
        # check that coef_ haven't been re-allocated
        assert_true(id1, id2)

    def test_fit_then_partial_fit(self):
        # Partial_fit should work after initial fit in the multiclass case.
        # Non-regression test for #2496; fit would previously produce a
        # Fortran-ordered coef_ that subsequent partial_fit couldn't handle.
        clf = self.factory()
        clf.fit(X2, Y2)
        clf.partial_fit(X2, Y2)     # no exception here

    def _test_partial_fit_equal_fit(self, lr):
        for X_, Y_, T_ in ((X, Y, T), (X2, Y2, T2)):
            clf = self.factory(alpha=0.01, eta0=0.01, n_iter=2,
                               learning_rate=lr, shuffle=False)
            clf.fit(X_, Y_)
            y_pred = clf.decision_function(T_)
            t = clf.t_

            classes = np.unique(Y_)
            clf = self.factory(alpha=0.01, eta0=0.01, learning_rate=lr,
                               shuffle=False)
            for i in range(2):
                clf.partial_fit(X_, Y_, classes=classes)
            y_pred2 = clf.decision_function(T_)

            assert_equal(clf.t_, t)
            assert_array_almost_equal(y_pred, y_pred2, decimal=2)

    def test_partial_fit_equal_fit_constant(self):
        self._test_partial_fit_equal_fit("constant")

    def test_partial_fit_equal_fit_optimal(self):
        self._test_partial_fit_equal_fit("optimal")

    def test_partial_fit_equal_fit_invscaling(self):
        self._test_partial_fit_equal_fit("invscaling")

    def test_regression_losses(self):
        clf = self.factory(alpha=0.01, learning_rate="constant",
                           eta0=0.1, loss="epsilon_insensitive")
        clf.fit(X, Y)
        assert_equal(1.0, np.mean(clf.predict(X) == Y))

        clf = self.factory(alpha=0.01, learning_rate="constant",
                           eta0=0.1, loss="squared_epsilon_insensitive")
        clf.fit(X, Y)
        assert_equal(1.0, np.mean(clf.predict(X) == Y))

        clf = self.factory(alpha=0.01, loss="huber")
        clf.fit(X, Y)
        assert_equal(1.0, np.mean(clf.predict(X) == Y))

        clf = self.factory(alpha=0.01, learning_rate="constant", eta0=0.01,
                           loss="squared_loss")
        clf.fit(X, Y)
        assert_equal(1.0, np.mean(clf.predict(X) == Y))

    def test_warm_start_multiclass(self):
        self._test_warm_start(X2, Y2, "optimal")

    def test_multiple_fit(self):
        # Test multiple calls of fit w/ different shaped inputs.
        clf = self.factory(alpha=0.01, n_iter=5,
                           shuffle=False)
        clf.fit(X, Y)
        assert_true(hasattr(clf, "coef_"))

        # Non-regression test: try fitting with a different label set.
        y = [["ham", "spam"][i] for i in LabelEncoder().fit_transform(Y)]
        clf.fit(X[:, :-1], y)

class SparseSGDClassifierTestCase(DenseSGDClassifierTestCase):
    """Run exactly the same tests using the sparse representation variant"""

    factory_class = SparseSGDClassifier

###############################################################################
# Regression Test Case

class DenseSGDRegressorTestCase(unittest.TestCase, CommonTest):
    """Test suite for the dense representation variant of SGD"""

    factory_class = SGDRegressor

    def test_sgd(self):
        # Check that SGD gives any results.
        clf = self.factory(alpha=0.1, n_iter=2,
                           fit_intercept=False)
        clf.fit([[0, 0], [1, 1], [2, 2]], [0, 1, 2])
        assert_equal(clf.coef_[0], clf.coef_[1])

    @raises(ValueError)
    def test_sgd_bad_penalty(self):
        # Check whether expected ValueError on bad penalty
        self.factory(penalty='foobar', l1_ratio=0.85)

    @raises(ValueError)
    def test_sgd_bad_loss(self):
        # Check whether expected ValueError on bad loss
        self.factory(loss="foobar")

    def test_sgd_averaged_computed_correctly(self):
        # Tests the average regressor matches the naive implementation

        eta = .001
        alpha = .01
        n_samples = 20
        n_features = 10
        rng = np.random.RandomState(0)
        X = rng.normal(size=(n_samples, n_features))
        w = rng.normal(size=n_features)

        # simple linear function without noise
        y = np.dot(X, w)

        clf = self.factory(loss='squared_loss',
                           learning_rate='constant',
                           eta0=eta, alpha=alpha,
                           fit_intercept=True,
                           n_iter=1, average=True, shuffle=False)

        clf.fit(X, y)
        average_weights, average_intercept = self.asgd(X, y, eta, alpha)

        assert_array_almost_equal(clf.coef_,
                                  average_weights,
                                  decimal=16)
        assert_almost_equal(clf.intercept_, average_intercept, decimal=16)

    def test_sgd_averaged_partial_fit(self):
        # Tests whether the partial fit yields the same average as the fit
        eta = .001
        alpha = .01
        n_samples = 20
        n_features = 10
        rng = np.random.RandomState(0)
        X = rng.normal(size=(n_samples, n_features))
        w = rng.normal(size=n_features)

        # simple linear function without noise
        y = np.dot(X, w)

        clf = self.factory(loss='squared_loss',
                           learning_rate='constant',
                           eta0=eta, alpha=alpha,
                           fit_intercept=True,
                           n_iter=1, average=True, shuffle=False)

        clf.partial_fit(X[:int(n_samples / 2)][:], y[:int(n_samples / 2)])
        clf.partial_fit(X[int(n_samples / 2):][:], y[int(n_samples / 2):])
        average_weights, average_intercept = self.asgd(X, y, eta, alpha)

        assert_array_almost_equal(clf.coef_,
                                  average_weights,
                                  decimal=16)
        assert_almost_equal(clf.intercept_[0], average_intercept, decimal=16)

    def test_average_sparse(self):
        # Checks the average weights on data with 0s

        eta = .001
        alpha = .01
        clf = self.factory(loss='squared_loss',
                           learning_rate='constant',
                           eta0=eta, alpha=alpha,
                           fit_intercept=True,
                           n_iter=1, average=True, shuffle=False)

        n_samples = Y3.shape[0]

        clf.partial_fit(X3[:int(n_samples / 2)][:], Y3[:int(n_samples / 2)])
        clf.partial_fit(X3[int(n_samples / 2):][:], Y3[int(n_samples / 2):])
        average_weights, average_intercept = self.asgd(X3, Y3, eta, alpha)

        assert_array_almost_equal(clf.coef_,
                                  average_weights,
                                  decimal=16)
        assert_almost_equal(clf.intercept_, average_intercept, decimal=16)

    def test_sgd_least_squares_fit(self):
        xmin, xmax = -5, 5
        n_samples = 100
        rng = np.random.RandomState(0)
        X = np.linspace(xmin, xmax, n_samples).reshape(n_samples, 1)

        # simple linear function without noise
        y = 0.5 * X.ravel()

        clf = self.factory(loss='squared_loss', alpha=0.1, n_iter=20,
                           fit_intercept=False)
        clf.fit(X, y)
        score = clf.score(X, y)
        assert_greater(score, 0.99)

        # simple linear function with noise
        y = 0.5 * X.ravel() + rng.randn(n_samples, 1).ravel()

        clf = self.factory(loss='squared_loss', alpha=0.1, n_iter=20,
                           fit_intercept=False)
        clf.fit(X, y)
        score = clf.score(X, y)
        assert_greater(score, 0.5)

    def test_sgd_epsilon_insensitive(self):
        xmin, xmax = -5, 5
        n_samples = 100
        X = np.linspace(xmin, xmax, n_samples).reshape(n_samples, 1)

        # simple linear function without noise
        y = 0.5 * X.ravel()

        clf = self.factory(loss='epsilon_insensitive', epsilon=0.01,
                           alpha=0.1, n_iter=20,
                           fit_intercept=False)
        clf.fit(X, y)
        score = clf.score(X, y)
        assert_true(score > 0.99)

        # simple linear function with noise
        y = 0.5 * X.ravel() \
            + np.random.randn(n_samples, 1).ravel()

        clf = self.factory(loss='epsilon_insensitive', epsilon=0.01,
                           alpha=0.1, n_iter=20,
                           fit_intercept=False)
        clf.fit(X, y)
        score = clf.score(X, y)
        assert_true(score > 0.5)

    def test_sgd_huber_fit(self):
        xmin, xmax = -5, 5
        n_samples = 100
        rng = np.random.RandomState(0)
        X = np.linspace(xmin, xmax, n_samples).reshape(n_samples, 1)

        # simple linear function without noise
        y = 0.5 * X.ravel()

        clf = self.factory(loss="huber", epsilon=0.1, alpha=0.1, n_iter=20,
                           fit_intercept=False)
        clf.fit(X, y)
        score = clf.score(X, y)
        assert_greater(score, 0.99)

        # simple linear function with noise
        y = 0.5 * X.ravel() + rng.randn(n_samples, 1).ravel()

        clf = self.factory(loss="huber", epsilon=0.1, alpha=0.1, n_iter=20,
                           fit_intercept=False)
        clf.fit(X, y)
        score = clf.score(X, y)
        assert_greater(score, 0.5)

    def test_elasticnet_convergence(self):
        # Check that the SGD output is consistent with coordinate descent

        n_samples, n_features = 1000, 5
        rng = np.random.RandomState(0)
        X = np.random.randn(n_samples, n_features)
        # ground_truth linear model that generate y from X and to which the
        # models should converge if the regularizer would be set to 0.0
        ground_truth_coef = rng.randn(n_features)
        y = np.dot(X, ground_truth_coef)

        # XXX: alpha = 0.1 seems to cause convergence problems
        for alpha in [0.01, 0.001]:
            for l1_ratio in [0.5, 0.8, 1.0]:
                cd = linear_model.ElasticNet(alpha=alpha, l1_ratio=l1_ratio,
                                             fit_intercept=False)
                cd.fit(X, y)
                sgd = self.factory(penalty='elasticnet', n_iter=50,
                                   alpha=alpha, l1_ratio=l1_ratio,
                                   fit_intercept=False)
                sgd.fit(X, y)
                err_msg = ("cd and sgd did not converge to comparable "
                           "results for alpha=%f and l1_ratio=%f"
                           % (alpha, l1_ratio))
                assert_almost_equal(cd.coef_, sgd.coef_, decimal=2,
                                    err_msg=err_msg)

    def test_partial_fit(self):
        third = X.shape[0] // 3
        clf = self.factory(alpha=0.01)

        clf.partial_fit(X[:third], Y[:third])
        assert_equal(clf.coef_.shape, (X.shape[1], ))
        assert_equal(clf.intercept_.shape, (1,))
        assert_equal(clf.decision_function([0, 0]).shape, (1, ))
        id1 = id(clf.coef_.data)

        clf.partial_fit(X[third:], Y[third:])
        id2 = id(clf.coef_.data)
        # check that coef_ haven't been re-allocated
        assert_true(id1, id2)

    def _test_partial_fit_equal_fit(self, lr):
        clf = self.factory(alpha=0.01, n_iter=2, eta0=0.01,
                           learning_rate=lr, shuffle=False)
        clf.fit(X, Y)
        y_pred = clf.predict(T)
        t = clf.t_

        clf = self.factory(alpha=0.01, eta0=0.01,
                           learning_rate=lr, shuffle=False)
        for i in range(2):
            clf.partial_fit(X, Y)
        y_pred2 = clf.predict(T)

        assert_equal(clf.t_, t)
        assert_array_almost_equal(y_pred, y_pred2, decimal=2)

    def test_partial_fit_equal_fit_constant(self):
        self._test_partial_fit_equal_fit("constant")

    def test_partial_fit_equal_fit_optimal(self):
        self._test_partial_fit_equal_fit("optimal")

    def test_partial_fit_equal_fit_invscaling(self):
        self._test_partial_fit_equal_fit("invscaling")

    def test_loss_function_epsilon(self):
        clf = self.factory(epsilon=0.9)
        clf.set_params(epsilon=0.1)
        assert clf.loss_functions['huber'][1] == 0.1

class SparseSGDRegressorTestCase(DenseSGDRegressorTestCase):
    # Run exactly the same tests using the sparse representation variant

    factory_class = SparseSGDRegressor

def test_l1_ratio():
    # Test if l1 ratio extremes match L1 and L2 penalty settings.
    X, y = datasets.make_classification(n_samples=1000,
                                        n_features=100, n_informative=20,
                                        random_state=1234)

    # test if elasticnet with l1_ratio near 1 gives same result as pure l1
    est_en = SGDClassifier(alpha=0.001, penalty='elasticnet',
                           l1_ratio=0.9999999999, random_state=42).fit(X, y)
    est_l1 = SGDClassifier(alpha=0.001, penalty='l1', random_state=42).fit(X, y)
    assert_array_almost_equal(est_en.coef_, est_l1.coef_)

    # test if elasticnet with l1_ratio near 0 gives same result as pure l2
    est_en = SGDClassifier(alpha=0.001, penalty='elasticnet',
                           l1_ratio=0.0000000001, random_state=42).fit(X, y)
    est_l2 = SGDClassifier(alpha=0.001, penalty='l2', random_state=42).fit(X, y)
    assert_array_almost_equal(est_en.coef_, est_l2.coef_)

def test_underflow_or_overlow():
    with np.errstate(all='raise'):
        # Generate some weird data with hugely unscaled features
        rng = np.random.RandomState(0)
        n_samples = 100
        n_features = 10

        X = rng.normal(size=(n_samples, n_features))
        X[:, :2] *= 1e300
        assert_true(np.isfinite(X).all())

        # Use MinMaxScaler to scale the data without introducing a numerical
        # instability (computing the standard deviation naively is not possible
        # on this data)
        X_scaled = MinMaxScaler().fit_transform(X)
        assert_true(np.isfinite(X_scaled).all())

        # Define a ground truth on the scaled data
        ground_truth = rng.normal(size=n_features)
        y = (np.dot(X_scaled, ground_truth) > 0.).astype(np.int32)
        assert_array_equal(np.unique(y), [0, 1])

        model = SGDClassifier(alpha=0.1, loss='squared_hinge', n_iter=500)

        # smoke test: model is stable on scaled data
        model.fit(X_scaled, y)
        assert_true(np.isfinite(model.coef_).all())

        # model is numerically unstable on unscaled data
        msg_regxp = (r"Floating-point under-/overflow occurred at epoch #.*"
                     " Scaling input data with StandardScaler or MinMaxScaler"
                     " might help.")
        assert_raises_regexp(ValueError, msg_regxp, model.fit, X, y)

def test_numerical_stability_large_gradient():
    # Non regression test case for numerical stability on scaled problems
    # where the gradient can still explode with some losses
    model = SGDClassifier(loss='squared_hinge', n_iter=10, shuffle=True,
                          penalty='elasticnet', l1_ratio=0.3, alpha=0.01,
                          eta0=0.001, random_state=0)
    with np.errstate(all='raise'):
        model.fit(iris.data, iris.target)
    assert_true(np.isfinite(model.coef_).all())

def test_large_regularization():
    # Non regression tests for numerical stability issues caused by large
    # regularization parameters
    for penalty in ['l2', 'l1', 'elasticnet']:
        model = SGDClassifier(alpha=1e5, learning_rate='constant', eta0=0.1,
                              n_iter=5, penalty=penalty, shuffle=False)
        with np.errstate(all='raise'):
            model.fit(iris.data, iris.target)
        assert_array_almost_equal(model.coef_, np.zeros_like(model.coef_))


"""Automatically download MLdata datasets."""

# Copyright (c) 2011 Pietro Berkes
# License: BSD 3 clause

import os
from os.path import join, exists
import re
import numbers
try:
    # Python 2
    from urllib2 import HTTPError
    from urllib2 import quote
    from urllib2 import urlopen
except ImportError:
    # Python 3+
    from urllib.error import HTTPError
    from urllib.parse import quote
    from urllib.request import urlopen

import numpy as np
import scipy as sp
from scipy import io
from shutil import copyfileobj

from .base import get_data_home, Bunch

MLDATA_BASE_URL = "http://mldata.org/repository/data/download/matlab/%s"

def mldata_filename(dataname):
    """Convert a raw name for a data set in a mldata.org filename."""
    dataname = dataname.lower().replace(' ', '-')
    return re.sub(r'[().]', '', dataname)

def fetch_mldata(dataname, target_name='label', data_name='data',
                 transpose_data=True, data_home=None):
    """Fetch an mldata.org data set

    If the file does not exist yet, it is downloaded from mldata.org .

    mldata.org does not have an enforced convention for storing data or
    naming the columns in a data set. The default behavior of this function
    works well with the most common cases:

      1) data values are stored in the column 'data', and target values in the
         column 'label'
      2) alternatively, the first column stores target values, and the second
         data values
      3) the data array is stored as `n_features x n_samples` , and thus needs
         to be transposed to match the `sklearn` standard

    Keyword arguments allow to adapt these defaults to specific data sets
    (see parameters `target_name`, `data_name`, `transpose_data`, and
    the examples below).

    mldata.org data sets may have multiple columns, which are stored in the
    Bunch object with their original name.

    Parameters
    ----------

    dataname:
        Name of the data set on mldata.org,
        e.g.: "leukemia", "Whistler Daily Snowfall", etc.
        The raw name is automatically converted to a mldata.org URL .

    target_name: optional, default: 'label'
        Name or index of the column containing the target values.

    data_name: optional, default: 'data'
        Name or index of the column containing the data.

    transpose_data: optional, default: True
        If True, transpose the downloaded data array.

    data_home: optional, default: None
        Specify another download and cache folder for the data sets. By default
        all scikit learn data is stored in '~/scikit_learn_data' subfolders.

    Returns
    -------

    data : Bunch
        Dictionary-like object, the interesting attributes are:
        'data', the data to learn, 'target', the classification labels,
        'DESCR', the full description of the dataset, and
        'COL_NAMES', the original names of the dataset columns.

    Examples
    --------
    Load the 'iris' dataset from mldata.org:

    >>> from sklearn.datasets.mldata import fetch_mldata
    >>> import tempfile
    >>> test_data_home = tempfile.mkdtemp()

    >>> iris = fetch_mldata('iris', data_home=test_data_home)
    >>> iris.target.shape
    (150,)
    >>> iris.data.shape
    (150, 4)

    Load the 'leukemia' dataset from mldata.org, which needs to be transposed
    to respects the sklearn axes convention:

    >>> leuk = fetch_mldata('leukemia', transpose_data=True,
    ...                     data_home=test_data_home)
    >>> leuk.data.shape
    (72, 7129)

    Load an alternative 'iris' dataset, which has different names for the
    columns:

    >>> iris2 = fetch_mldata('datasets-UCI iris', target_name=1,
    ...                      data_name=0, data_home=test_data_home)
    >>> iris3 = fetch_mldata('datasets-UCI iris',
    ...                      target_name='class', data_name='double0',
    ...                      data_home=test_data_home)

    >>> import shutil
    >>> shutil.rmtree(test_data_home)
    """

    # normalize dataset name
    dataname = mldata_filename(dataname)

    # check if this data set has been already downloaded
    data_home = get_data_home(data_home=data_home)
    data_home = join(data_home, 'mldata')
    if not exists(data_home):
        os.makedirs(data_home)

    matlab_name = dataname + '.mat'
    filename = join(data_home, matlab_name)

    # if the file does not exist, download it
    if not exists(filename):
        urlname = MLDATA_BASE_URL % quote(dataname)
        try:
            mldata_url = urlopen(urlname)
        except HTTPError as e:
            if e.code == 404:
                e.msg = "Dataset '%s' not found on mldata.org." % dataname
            raise
        # store Matlab file
        try:
            with open(filename, 'w+b') as matlab_file:
                copyfileobj(mldata_url, matlab_file)
        except:
            os.remove(filename)
            raise
        mldata_url.close()

    # load dataset matlab file
    with open(filename, 'rb') as matlab_file:
        matlab_dict = io.loadmat(matlab_file, struct_as_record=True)

    # -- extract data from matlab_dict

    # flatten column names
    col_names = [str(descr[0])
                 for descr in matlab_dict['mldata_descr_ordering'][0]]

    # if target or data names are indices, transform then into names
    if isinstance(target_name, numbers.Integral):
        target_name = col_names[target_name]
    if isinstance(data_name, numbers.Integral):
        data_name = col_names[data_name]

    # rules for making sense of the mldata.org data format
    # (earlier ones have priority):
    # 1) there is only one array => it is "data"
    # 2) there are multiple arrays
    #    a) copy all columns in the bunch, using their column name
    #    b) if there is a column called `target_name`, set "target" to it,
    #        otherwise set "target" to first column
    #    c) if there is a column called `data_name`, set "data" to it,
    #        otherwise set "data" to second column

    dataset = {'DESCR': 'mldata.org dataset: %s' % dataname,
               'COL_NAMES': col_names}

    # 1) there is only one array => it is considered data
    if len(col_names) == 1:
        data_name = col_names[0]
        dataset['data'] = matlab_dict[data_name]
    # 2) there are multiple arrays
    else:
        for name in col_names:
            dataset[name] = matlab_dict[name]

        if target_name in col_names:
            del dataset[target_name]
            dataset['target'] = matlab_dict[target_name]
        else:
            del dataset[col_names[0]]
            dataset['target'] = matlab_dict[col_names[0]]

        if data_name in col_names:
            del dataset[data_name]
            dataset['data'] = matlab_dict[data_name]
        else:
            del dataset[col_names[1]]
            dataset['data'] = matlab_dict[col_names[1]]

    # set axes to sklearn conventions
    if transpose_data:
        dataset['data'] = dataset['data'].T
    if 'target' in dataset:
        if not sp.sparse.issparse(dataset['target']):
            dataset['target'] = dataset['target'].squeeze()

    return Bunch(**dataset)

# The following is used by nosetests to setup the docstring tests fixture

def setup_module(module):
    # setup mock urllib2 module to avoid downloading from mldata.org
    from sklearn.utils.testing import install_mldata_mock
    install_mldata_mock({
        'iris': {
            'data': np.empty((150, 4)),
            'label': np.empty(150),
        },
        'datasets-uci-iris': {
            'double0': np.empty((150, 4)),
            'class': np.empty((150,)),
        },
        'leukemia': {
            'data': np.empty((72, 7129)),
        },
    })

def teardown_module(module):
    from sklearn.utils.testing import uninstall_mldata_mock
    uninstall_mldata_mock()


#!/usr/bin/python

import os
import sys
import time
import rospkg
import numpy as np
import matplotlib.pyplot as plt
import yaml
import argparse
from matplotlib import rc

# tell matplotlib to use latex font
rc('font',**{'family':'serif','serif':['Cardo']})
rc('text', usetex=True)

from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes
from mpl_toolkits.axes_grid1.inset_locator import mark_inset

def plot_trajectory(ax, filename, label, color, linewidth):
  file = open(filename)
  data = file.read()
  lines = data.replace(","," ").replace("\t"," ").split("\n")
  trajectory = np.array([[v.strip() for v in line.split(" ") if v.strip()!=""] for line in lines if len(line)>0 and line[0]!="#"], dtype=np.float64)
  ax.plot(trajectory[:,1], trajectory[:,2], label=label, color=color, linewidth=linewidth)

def compare_results(experiments, results_dir, comparison_dir,
                    plot_scale_drift = False):

  # ------------------------------------------------------------------------------
  # position error
  fig_poserr = plt.figure(figsize=(8,6))
  ax_poserr_x = fig_poserr.add_subplot(311, ylabel='x-error [m]')
  ax_poserr_y = fig_poserr.add_subplot(312, ylabel='y-error [m]')
  ax_poserr_z = fig_poserr.add_subplot(313, ylabel='z-error [m]', xlabel='time [s]')

  for exp in experiments:

    # load dataset parameters
    params_stream = open(os.path.join(results_dir, exp, 'params.yaml'))
    params = yaml.load(params_stream)

    # plot translation error
    trans_error = np.loadtxt(os.path.join(results_dir, exp, 'translation_error.txt'))
    trans_error[:,0] = trans_error[:,0]-trans_error[0,0]
    ax_poserr_x.plot(trans_error[:,0], trans_error[:,1], label=params['experiment_label'])
    ax_poserr_y.plot(trans_error[:,0], trans_error[:,2])
    ax_poserr_z.plot(trans_error[:,0], trans_error[:,3])

  ax_poserr_x.set_xlim([0, trans_error[-1,0]+4])
  ax_poserr_y.set_xlim([0, trans_error[-1,0]+4])
  ax_poserr_z.set_xlim([0, trans_error[-1,0]+4])
  ax_poserr_x.legend(bbox_to_anchor=[0, 0], loc='lower left', ncol=3)
  ax_poserr_x.grid()
  ax_poserr_y.grid()
  ax_poserr_z.grid()
  fig_poserr.tight_layout()
  fig_poserr.savefig(os.path.join(comparison_dir, 'translation_error.pdf'))

  # ------------------------------------------------------------------------------
  # orientation error
  fig_roterr = plt.figure(figsize=(8,6))
  ax_roterr_r = fig_roterr.add_subplot(311, ylabel='roll-error [rad]')
  ax_roterr_p = fig_roterr.add_subplot(312, ylabel='pitch-error [rad]')
  ax_roterr_y = fig_roterr.add_subplot(313, ylabel='yaw-error [rad]', xlabel='time [s]')

  for exp in experiments:

    # load dataset parameters
    params_stream = open(os.path.join(results_dir, exp, 'params.yaml'))
    params = yaml.load(params_stream)

    # plot translation error
    rot_error = np.loadtxt(os.path.join(results_dir, exp, 'orientation_error.txt'))
    rot_error[:,0] = rot_error[:,0]-rot_error[0,0]
    ax_roterr_r.plot(rot_error[:,0], rot_error[:,3], label=params['experiment_label'])
    ax_roterr_p.plot(rot_error[:,0], rot_error[:,2])
    ax_roterr_y.plot(rot_error[:,0], rot_error[:,1])

  ax_roterr_r.set_xlim([0, rot_error[-1,0]+4])
  ax_roterr_p.set_xlim([0, rot_error[-1,0]+4])
  ax_roterr_y.set_xlim([0, rot_error[-1,0]+4])
  ax_roterr_r.legend(bbox_to_anchor=[0, 1], loc='upper left', ncol=3)
  ax_roterr_r.grid()
  ax_roterr_p.grid()
  ax_roterr_y.grid()
  fig_roterr.tight_layout()
  fig_roterr.savefig(os.path.join(comparison_dir, 'orientation_error.pdf'))

  # ------------------------------------------------------------------------------
  # scale error
  if plot_scale_drift:
    fig_scale = plt.figure(figsize=(8,2.5))
    ax_scale = fig_scale.add_subplot(111, xlabel='time [s]', ylabel='scale change [\%]')

    for exp in experiments:

      # load dataset parameters
      params = yaml.load(open(os.path.join(results_dir, exp, 'params.yaml')))

      # plot translation error
      scale_drift = open(os.path.join(results_dir, exp, 'scale_drift.txt'))
      scale_drift[:,0] = scale_drift[:,0]-scale_drift[0,0]
      ax_scale.plot(scale_drift[:,0], scale_drift[:,1], label=params['experiment_label'])

    ax_scale.set_xlim([0, rot_error[-1,0]+4])
    ax_scale.legend(bbox_to_anchor=[0, 1], loc='upper left', ncol=3)
    ax_scale.grid()
    fig_scale.tight_layout()
    fig_scale.savefig(os.path.join(comparison_dir, 'scale_drift.pdf'))

  # ------------------------------------------------------------------------------
  # trajectory

#   fig_traj = plt.figure(figsize=(8,4.8))
#   ax_traj = fig_traj.add_subplot(111, xlabel='x [m]', ylabel='y [m]', aspect='equal', xlim=[-3.1, 4], ylim=[-1.5, 2.6])
#
#   plotTrajectory(ax_traj, '/home/cforster/Datasets/asl_vicon_d2/groundtruth_filtered.txt', 'Groundtruth', 'k', 1.5)
#   plotTrajectory(ax_traj, results_dir+'/20130911_2229_nslam_i7_asl2_fast/traj_estimate_rotated.txt', 'Fast', 'g', 1)
#   plotTrajectory(ax_traj, results_dir+'/20130906_2149_ptam_i7_asl2/traj_estimate_rotated.txt', 'PTAM', 'r', 1)
#
#   mark_inset(ax_traj, axins, loc1=2, loc2=4, fc="none", ec='b')
#   plt.draw()
#   plt.show()
#   ax_traj.legend(bbox_to_anchor=[1, 0], loc='lower right', ncol=3)
#   ax_traj.grid()
#   fig_traj.tight_layout()
#   fig_traj.savefig('../results/trajectory_asl.pdf')

if __name__ == '__main__':

  default_name = time.strftime("%Y%m%d_%H%M", time.localtime())+'_comparison'
  parser = argparse.ArgumentParser(description='Compare results.')
  parser.add_argument('result_directories', nargs='+', help='list of result directories to compare')
  parser.add_argument('--name', help='name of the comparison', default=default_name)
  args = parser.parse_args()

  # create folder for comparison results
  results_dir = os.path.join(rospkg.RosPack().get_path('svo_analysis'), 'results')
  comparison_dir = os.path.join(results_dir, args.name)
  if not os.path.exists(comparison_dir):
    os.makedirs(comparison_dir)

  # run comparison
  compare_results(args.result_directories, results_dir, comparison_dir)


02 - Programs being tested/RV_function.py

# -*- coding: utf-8 -*-
"""
Created on Fri Feb  3 11:36:58 2017

@author: camacho
"""
import  numpy as np
import matplotlib.pyplot as pl
pl.close("all")

##### RV FUNCTION 1 - circular orbit
def RV_circular(P=365,K=0.1,T=0,gamma=0,time=100,space=20):
    #parameters
    #P = period in days
    #K =  semi-amplitude of the signal
    #T = velocity at zero phase
    #gamma = average velocity of the star
    #time  = time of the simulation
    #space => I want an observation every time/space days

    t=np.linspace(0,time,space)
    RV=[K*np.sin(2*np.pi*x/P - T) + gamma for x in t]
    RV=[x for x in RV] #m/s
    return [t,RV]

##### RV FUNCTION 2 - keplerian orbit
def RV_kepler(P=365,e=0,K=0.1,T=0,gamma=0,w=np.pi,time=100,space=1000):
    #parameters
    #P = period in days
    #e = eccentricity
    #K = RV amplitude
    #gamma = constant system RV
    #T = zero phase
    #w = longitude of the periastron
    #time = time of the simulation
    #space => I want an observation every time/space days

    t=np.linspace(0,time,space)

    #mean anomaly
    Mean_anom=[2*np.pi*(x1-T)/P  for x1 in t]

    #eccentric anomaly -> E0=M + e*sin(M) + 0.5*(e**2)*sin(2*M)
    E0=[x + e*np.sin(x)  + 0.5*(e**2)*np.sin(2*x) for x in Mean_anom]
    #mean anomaly -> M0=E0 - e*sin(E0)
    M0=[x - e*np.sin(x) for x in E0]

    i=0
    while i<100:
        #[x + y for x, y in zip(first, second)]
        calc_aux=[x2-y for x2,y in zip(Mean_anom,M0)]
        E1=[x3 + y/(1-e*np.cos(x3)) for x3,y in zip(E0,calc_aux)]
        M1=[x4 - e*np.sin(x4) for x4 in E0]
        i+=1
        E0=E1
        M0=M1

    nu=[2*np.arctan(np.sqrt((1+e)/(1-e))*np.tan(x5/2)) for x5 in E0]
    RV=[ gamma + K*(e*np.cos(w)+np.cos(w+x6)) for x6 in nu]
    RV=[x for x in RV] #m/s
    return t,RV

#Examples
#a=RV_circular()
#pl.figure('RV_circular with P=365')
#pl.plot(a[0],a[1],':',)
#pl.title('planet of 365 days orbit')
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')

#b=RV_circular(P=100)
#pl.figure('RV_circular with P=100')
#pl.title('planet of 100 days orbit')
#pl.plot(b[0],b[1],':',)
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')

#c=RV_kepler(P=100,e=0,w=np.pi,time=100)
#pl.figure()
#pl.plot(c[0],c[1],':',)
#pl.title('P=100, e=0, w=pi, time=100')
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')

#d1=RV_kepler(P=100,e=0, w=0,time=500)
#pl.figure()
#pl.title('P=100, e=0, w=pi, time=25')
#pl.plot(d[0],d[1],'-',)
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')
#d2=RV_kepler(P=100,e=0, w=np.pi,time=500)
#pl.figure()
#pl.title('P=100, e=0, w=pi, time=25')
#pl.plot(d[0],d[1],'-',)
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')
#d3=RV_kepler(P=100,e=0.5, w=np.pi,time=500)
#pl.figure()
#pl.title('P=100, e=0, w=pi, time=25')
#pl.plot(d[0],d[1],'-',)
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')
#d4=RV_kepler(P=100,e=0.5, w=np.pi/2,time=500)
#pl.figure()
#pl.title('P=100, e=0, w=pi, time=25')
#pl.plot(d[0],d[1],'-',)
#pl.xlabel('time')
#pl.ylabel('RV (Km/s)')

d1=RV_kepler(P=100,e=0, w=0,time=500)
d2=RV_kepler(P=100,e=0.5, w=0,time=500)
d3=RV_kepler(P=100,e=0.5, w=np.pi,time=500)
d4=RV_kepler(P=100,e=0.5, w=np.pi/2,time=500)

# Four axes, returned as a 2-d array
f, axarr = pl.subplots(2, 2)

axarr[0, 0].plot(d1[0],d1[1])
axarr[0, 0].set_title('e=0 and w=0')

axarr[0, 1].plot(d2[0],d2[1])
axarr[0, 1].set_title('e=0.5, w=0')

axarr[1, 0].plot(d3[0],d3[1])
axarr[1, 0].set_title('e=0.5, w=pi')

axarr[1, 1].plot(d4[0],d4[1])
axarr[1, 1].set_title('e=0.5, w=pi/2')
#pl.setp(pl.xticks(fontsize = 18) for a in axarr[0,:])#pl.yticks(fontsize=18))
pl.setp([a.get_xticklabels() for a in axarr[0, :]], visible=False)


from __future__ import print_function

import warnings

import numpy as np
from six import next
from six.moves import xrange
from shapely.geometry import Polygon

def plot_polygon(ax, poly, facecolor='red', edgecolor='black', alpha=0.5, linewidth=1.0, **kwargs):
    """ Plot a single Polygon geometry """
    from descartes.patch import PolygonPatch
    a = np.asarray(poly.exterior)
    if poly.has_z:
        poly = Polygon(zip(*poly.exterior.xy))

    # without Descartes, we could make a Patch of exterior
    ax.add_patch(PolygonPatch(poly, facecolor=facecolor, linewidth=0, alpha=alpha))  # linewidth=0 because boundaries are drawn separately
    ax.plot(a[:, 0], a[:, 1], color=edgecolor, linewidth=linewidth, **kwargs)
    for p in poly.interiors:
        x, y = zip(*p.coords)
        ax.plot(x, y, color=edgecolor, linewidth=linewidth)

def plot_multipolygon(ax, geom, facecolor='red', edgecolor='black', alpha=0.5, linewidth=1.0, **kwargs):
    """ Can safely call with either Polygon or Multipolygon geometry
    """
    if geom.type == 'Polygon':
        plot_polygon(ax, geom, facecolor=facecolor, edgecolor=edgecolor, alpha=alpha, linewidth=linewidth, **kwargs)
    elif geom.type == 'MultiPolygon':
        for poly in geom.geoms:
            plot_polygon(ax, poly, facecolor=facecolor, edgecolor=edgecolor, alpha=alpha, linewidth=linewidth, **kwargs)

def plot_linestring(ax, geom, color='black', linewidth=1.0, **kwargs):
    """ Plot a single LineString geometry """
    a = np.array(geom)
    ax.plot(a[:, 0], a[:, 1], color=color, linewidth=linewidth, **kwargs)

def plot_multilinestring(ax, geom, color='red', linewidth=1.0, **kwargs):
    """ Can safely call with either LineString or MultiLineString geometry
    """
    if geom.type == 'LineString':
        plot_linestring(ax, geom, color=color, linewidth=linewidth, **kwargs)
    elif geom.type == 'MultiLineString':
        for line in geom.geoms:
            plot_linestring(ax, line, color=color, linewidth=linewidth, **kwargs)

def plot_point(ax, pt, marker='o', markersize=2, color='black', **kwargs):
    """ Plot a single Point geometry """
    ax.plot(pt.x, pt.y, marker=marker, markersize=markersize, color=color, **kwargs)

def gencolor(N, colormap='Set1'):
    """
    Color generator intended to work with one of the ColorBrewer
    qualitative color scales.

    Suggested values of colormap are the following:

        Accent, Dark2, Paired, Pastel1, Pastel2, Set1, Set2, Set3

    (although any matplotlib colormap will work).
    """
    from matplotlib import cm
    # don't use more than 9 discrete colors
    n_colors = min(N, 9)
    cmap = cm.get_cmap(colormap, n_colors)
    colors = cmap(range(n_colors))
    for i in xrange(N):
        yield colors[i % n_colors]

def plot_series(s, cmap='Set1', color=None, ax=None, linewidth=1.0,
                figsize=None, **color_kwds):
    """ Plot a GeoSeries

        Generate a plot of a GeoSeries geometry with matplotlib.

        Parameters
        ----------

        Series
            The GeoSeries to be plotted.  Currently Polygon,
            MultiPolygon, LineString, MultiLineString and Point
            geometries can be plotted.

        cmap : str (default 'Set1')
            The name of a colormap recognized by matplotlib.  Any
            colormap will work, but categorical colormaps are
            generally recommended.  Examples of useful discrete
            colormaps include:

                Accent, Dark2, Paired, Pastel1, Pastel2, Set1, Set2, Set3

        color : str (default None)
            If specified, all objects will be colored uniformly.

        ax : matplotlib.pyplot.Artist (default None)
            axes on which to draw the plot

        linewidth : float (default 1.0)
            Line width for geometries.

        figsize : pair of floats (default None)
            Size of the resulting matplotlib.figure.Figure. If the argument
            ax is given explicitly, figsize is ignored.

        **color_kwds : dict
            Color options to be passed on to the actual plot function

        Returns
        -------

        matplotlib axes instance
    """
    if 'colormap' in color_kwds:
        warnings.warn("'colormap' is deprecated, please use 'cmap' instead "
                      "(for consistency with matplotlib)", FutureWarning)
        cmap = color_kwds.pop('colormap')
    if 'axes' in color_kwds:
        warnings.warn("'axes' is deprecated, please use 'ax' instead "
                      "(for consistency with pandas)", FutureWarning)
        ax = color_kwds.pop('axes')

    import matplotlib.pyplot as plt
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
        ax.set_aspect('equal')
    color_generator = gencolor(len(s), colormap=cmap)
    for geom in s:
        if color is None:
            col = next(color_generator)
        else:
            col = color
        if geom.type == 'Polygon' or geom.type == 'MultiPolygon':
            if 'facecolor' in color_kwds:
                plot_multipolygon(ax, geom, linewidth=linewidth, **color_kwds)
            else:
                plot_multipolygon(ax, geom, facecolor=col, linewidth=linewidth, **color_kwds)
        elif geom.type == 'LineString' or geom.type == 'MultiLineString':
            plot_multilinestring(ax, geom, color=col, linewidth=linewidth, **color_kwds)
        elif geom.type == 'Point':
            plot_point(ax, geom, color=col, **color_kwds)
    plt.draw()
    return ax

def plot_dataframe(s, column=None, cmap=None, color=None, linewidth=1.0,
                   categorical=False, legend=False, ax=None,
                   scheme=None, k=5, vmin=None, vmax=None, figsize=None,
                   **color_kwds):
    """ Plot a GeoDataFrame

        Generate a plot of a GeoDataFrame with matplotlib.  If a
        column is specified, the plot coloring will be based on values
        in that column.  Otherwise, a categorical plot of the
        geometries in the `geometry` column will be generated.

        Parameters
        ----------

        GeoDataFrame
            The GeoDataFrame to be plotted.  Currently Polygon,
            MultiPolygon, LineString, MultiLineString and Point
            geometries can be plotted.

        column : str (default None)
            The name of the column to be plotted.

        categorical : bool (default False)
            If False, cmap will reflect numerical values of the
            column being plotted.  For non-numerical columns (or if
            column=None), this will be set to True.

        cmap : str (default 'Set1')
            The name of a colormap recognized by matplotlib.

        color : str (default None)
            If specified, all objects will be colored uniformly.

        linewidth : float (default 1.0)
            Line width for geometries.

        legend : bool (default False)
            Plot a legend (Experimental; currently for categorical
            plots only)

        ax : matplotlib.pyplot.Artist (default None)
            axes on which to draw the plot

        scheme : pysal.esda.mapclassify.Map_Classifier
            Choropleth classification schemes (requires PySAL)

        k   : int (default 5)
            Number of classes (ignored if scheme is None)

        vmin : None or float (default None)

            Minimum value of cmap. If None, the minimum data value
            in the column to be plotted is used.

        vmax : None or float (default None)

            Maximum value of cmap. If None, the maximum data value
            in the column to be plotted is used.

        figsize
            Size of the resulting matplotlib.figure.Figure. If the argument
            axes is given explicitly, figsize is ignored.

        **color_kwds : dict
            Color options to be passed on to the actual plot function

        Returns
        -------

        matplotlib axes instance
    """
    if 'colormap' in color_kwds:
        warnings.warn("'colormap' is deprecated, please use 'cmap' instead "
                      "(for consistency with matplotlib)", FutureWarning)
        cmap = color_kwds.pop('colormap')
    if 'axes' in color_kwds:
        warnings.warn("'axes' is deprecated, please use 'ax' instead "
                      "(for consistency with pandas)", FutureWarning)
        ax = color_kwds.pop('axes')

    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D
    from matplotlib.colors import Normalize
    from matplotlib import cm

    if column is None:
        return plot_series(s.geometry, cmap=cmap, color=color,
                           ax=ax, linewidth=linewidth, figsize=figsize,
                           **color_kwds)
    else:
        if s[column].dtype is np.dtype('O'):
            categorical = True
        if categorical:
            if cmap is None:
                cmap = 'Set1'
            categories = list(set(s[column].values))
            categories.sort()
            valuemap = dict([(k, v) for (v, k) in enumerate(categories)])
            values = [valuemap[k] for k in s[column]]
        else:
            values = s[column]
        if scheme is not None:
            binning = __pysal_choro(values, scheme, k=k)
            values = binning.yb
            # set categorical to True for creating the legend
            categorical = True
            binedges = [binning.yb.min()] + binning.bins.tolist()
            categories = ['{0:.2f} - {1:.2f}'.format(binedges[i], binedges[i+1])
                          for i in range(len(binedges)-1)]
        cmap = norm_cmap(values, cmap, Normalize, cm, vmin=vmin, vmax=vmax)
        if ax is None:
            fig, ax = plt.subplots(figsize=figsize)
            ax.set_aspect('equal')
        for geom, value in zip(s.geometry, values):
            if color is None:
                col = cmap.to_rgba(value)
            else:
                col = color
            if geom.type == 'Polygon' or geom.type == 'MultiPolygon':
                plot_multipolygon(ax, geom, facecolor=col, linewidth=linewidth, **color_kwds)
            elif geom.type == 'LineString' or geom.type == 'MultiLineString':
                plot_multilinestring(ax, geom, color=col, linewidth=linewidth, **color_kwds)
            elif geom.type == 'Point':
                plot_point(ax, geom, color=col, **color_kwds)
        if legend:
            if categorical:
                patches = []
                for value, cat in enumerate(categories):
                    patches.append(Line2D([0], [0], linestyle="none",
                                          marker="o", alpha=color_kwds.get('alpha', 0.5),
                                          markersize=10, markerfacecolor=cmap.to_rgba(value)))
                ax.legend(patches, categories, numpoints=1, loc='best')
            else:
                # TODO: show a colorbar
                raise NotImplementedError
    plt.draw()
    return ax

def __pysal_choro(values, scheme, k=5):
    """ Wrapper for choropleth schemes from PySAL for use with plot_dataframe

        Parameters
        ----------

        values
            Series to be plotted

        scheme
            pysal.esda.mapclassify classificatin scheme
            ['Equal_interval'|'Quantiles'|'Fisher_Jenks']

        k
            number of classes (2 <= k <=9)

        Returns
        -------

        binning
            Binning objects that holds the Series with values replaced with
            class identifier and the bins.
    """

    try:
        from pysal.esda.mapclassify import Quantiles, Equal_Interval, Fisher_Jenks
        schemes = {}
        schemes['equal_interval'] = Equal_Interval
        schemes['quantiles'] = Quantiles
        schemes['fisher_jenks'] = Fisher_Jenks
        s0 = scheme
        scheme = scheme.lower()
        if scheme not in schemes:
            scheme = 'quantiles'
            warnings.warn('Unrecognized scheme "{0}". Using "Quantiles" '
                          'instead'.format(s0), UserWarning, stacklevel=3)
        if k < 2 or k > 9:
            warnings.warn('Invalid k: {0} (2 <= k <= 9), setting k=5 '
                          '(default)'.format(k), UserWarning, stacklevel=3)
            k = 5
        binning = schemes[scheme](values, k)
        return binning
    except ImportError:
        raise ImportError("PySAL is required to use the 'scheme' keyword")

def norm_cmap(values, cmap, normalize, cm, vmin=None, vmax=None):

    """ Normalize and set colormap

        Parameters
        ----------

        values
            Series or array to be normalized

        cmap
            matplotlib Colormap

        normalize
            matplotlib.colors.Normalize

        cm
            matplotlib.cm

        vmin
            Minimum value of colormap. If None, uses min(values).

        vmax
            Maximum value of colormap. If None, uses max(values).

        Returns
        -------
        n_cmap
            mapping of normalized values to colormap (cmap)

    """

    mn = min(values) if vmin is None else vmin
    mx = max(values) if vmax is None else vmax
    norm = normalize(vmin=mn, vmax=mx)
    n_cmap = cm.ScalarMappable(norm=norm, cmap=cmap)
    return n_cmap


import numpy as _np

from matplotlib.widgets import AxesWidget as _AxesWidget
from matplotlib.colors import is_color_like as _is_color_like
from matplotlib.axes import Axes as _mplAxes
from matplotlib.figure import Figure as _mplFigure
from IPython.display import display as _ipydisplay

from pyemma.util.types import is_int as _is_int
from scipy.spatial import cKDTree as _cKDTree

from ._bmutils import get_ascending_coord_idx

from mdtraj import Trajectory as _mdTrajectory
from nglview import NGLWidget as _NGLwdg

from ipywidgets import HBox as _HBox, VBox as _VBox

def pts_per_axis_unit(mplax, pt_per_inch=72):
    r"""
    Return how many pt per axis unit of a given maptplotlib axis a figure has

    Parameters
    ----------

    mplax : :obj:`matplotlib.axes._subplots.AxesSubplot`

    pt_per_inch : how many points are in an inch (this number should not change)

    Returns
    --------

    pt_per_xunit, pt_per_yunit

    """

    # matplotlib voodoo
    # Get bounding box
    bbox = mplax.get_window_extent().transformed(mplax.get_figure().dpi_scale_trans.inverted())

    span_inch = _np.array([bbox.width, bbox.height], ndmin=2).T

    span_units = [mplax.get_xlim(), mplax.get_ylim()]
    span_units = _np.diff(span_units, axis=1)

    inch_per_unit = span_inch / span_units
    return inch_per_unit * pt_per_inch

def update2Dlines(iline, x, y):
    """
    provide a common interface to update objects on the plot to a new position (x,y) depending
    on whether they are hlines, vlines, dots etc

    Parameters
    ----------

    iline: :obj:`matplotlib.lines.Line2D` object

    x : float with new position

    y : float with new position

    """
    # TODO FIND OUT A CLEANER WAY TO DO THIS (dict or class)

    if not hasattr(iline,'whatisthis'):
        raise AttributeError("This method will only work if iline has the attribute 'whatsthis'")
    else:
        # TODO find cleaner way of distinguishing these 2Dlines
        if iline.whatisthis in ['dot']:
            iline.set_xdata((x))
            iline.set_ydata((y))
        elif iline.whatisthis in ['lineh']:
            iline.set_ydata((y,y))
        elif iline.whatisthis in ['linev']:
            iline.set_xdata((x,x))
        else:
            # TODO: FIND OUT WNY EXCEPTIONS ARE NOT BEING RAISED
            raise TypeError("what is this type of 2Dline?")

class ClickOnAxisListener(object):
    def __init__(self, ngl_wdg, crosshairs, showclick_objs, ax, pos,
                 list_mpl_objects_to_update):
        self.ngl_wdg = ngl_wdg
        self.crosshairs = crosshairs
        self.showclick_objs = showclick_objs
        self.ax = ax
        self.pos = pos
        self.list_mpl_objects_to_update = list_mpl_objects_to_update
        self.list_of_dots = [None]*self.pos.shape[0]
        self.fig_size = self.ax.figure.get_size_inches()
        self.kdtree = None

    def build_tree(self):
        # Use ax.transData to compute distance in pixels
        # regardelss of the axes units (http://matplotlib.org/users/transforms_tutorial.html)
        # Corresponds to the visual distance between clicked point and target point
        self.kdtree = _cKDTree(self.ax.transData.transform(self.pos))

    @property
    def figure_changed_size(self):
        return not _np.allclose(self.fig_size, self.ax.figure.get_size_inches())

    def __call__(self, event):
        # Wait for the first click or a a figsize change
        # to build the kdtree
        if self.figure_changed_size or self.kdtree is None:
            self.build_tree()
            self.fig_size = self.ax.figure.get_size_inches()

        # Was the click inside the bounding box?
        if self.ax.get_window_extent().contains(event.x, event.y):
            if self.crosshairs:
                for iline in self.showclick_objs:
                    update2Dlines(iline, event.xdata, event.ydata)

            _, index = self.kdtree.query(x=[event.x, event.y], k=1)
            for idot in self.list_mpl_objects_to_update:
                update2Dlines(idot, self.pos[index, 0], self.pos[index, 1])

            self.ngl_wdg.isClick = True
            if hasattr(self.ngl_wdg, '_GeomsInWid'):
                # We're in a sticky situation
                if event.button == 1:
                    # Pressed left
                    self.ngl_wdg._GeomsInWid[index].show()
                    if self.list_of_dots[index] is None:
                        # Plot and store the dot in case there wasn't
                        self.list_of_dots[index] = self.ax.plot(self.pos[index, 0], self.pos[index, 1], 'o',
                                                                c=self.ngl_wdg._GeomsInWid[index].color_dot, ms=7)[0]
                elif event.button in [2, 3]:
                    #  Pressed right or middle
                    self.ngl_wdg._GeomsInWid[index].hide()
                    # Delete dot if the geom is not visible anymore
                    if not self.ngl_wdg._GeomsInWid[index].is_visible() and self.list_of_dots[index] is not None:
                        self.list_of_dots[index].remove()
                        self.list_of_dots[index] = None
            else:
                # We're not sticky, just go to the frame
                self.ngl_wdg.frame = index

class MolPXBox(object):
    r"""
    Class created to be the parent class of MolPXHBox and MolPXVBox, which inherit from
    MolPXBox and the ipywidget classes HBox and VBox (*args and **kwargs are for these)

    The sole purpose of this class is to  avoid monkey-patching elsewhere in the code,
     this class creates them as empty lists on instantiation.

    It also implements two methods:

    * self.display (=IPython.display(self)

    * append_if_existing
    """
    def __init__(self, *args, **kwargs):
        self.linked_axes = []
        self.linked_mdgeoms = []
        self.linked_ngl_wdgs = []
        self.linked_data_arrays = []
        self.linked_ax_wdgs = []
        self.linked_figs = []

    def display(self):
        _ipydisplay(self)

    def append_if_existing(self, args0, startswith_arg="linked_"):
        r"""
        args0  is the tuple containing all widgets to be included in the MolPXBox
        this tuple can contain itself other MolPXWidget
        so we iterate through them and appending linked stuff
        """
        for iarg in args0:
            for attrname in dir(iarg):
                if attrname.startswith(startswith_arg) and len(iarg.__dict__[attrname]) != 0:
                    self.__dict__[attrname] += iarg.__dict__[attrname]

def auto_append_these_mpx_attrs(iobj, *attrs):
    r""" The attribute s name is automatically derived
    from the attribute s type via a type:name dictionary

    *attrs : any number of unnamed objects of the types in type2attrname.
             If the object type is a list, it will be flattened prior to attempting
    """

    attrs_flat_list = []
    for sublist in attrs:
        if isinstance(sublist, list):
            for item in sublist:
                attrs_flat_list.append(item)
        else:
            attrs_flat_list.append(sublist)

    # Go through the arguments and assign them an attrname according to their types
    for iattr in attrs_flat_list:
        for attrname, itype in type2attrname.items():
            if isinstance(iattr, itype):
                iobj.__dict__[attrname].append(iattr)
                break

class MolPXHBox(_HBox, MolPXBox):
    def __init__(self, *args, **kwargs):
        super(MolPXHBox, self).__init__(*args, **kwargs)
        self.append_if_existing(args[0])

class MolPXVBox(_VBox, MolPXBox):
    def __init__(self, *args, **kwargs):
        super(MolPXVBox, self).__init__(*args, **kwargs)
        self.append_if_existing(args[0])

type2attrname = {"linked_axes": _mplAxes,
                 "linked_mdgeoms": _mdTrajectory,
                 "linked_ngl_wdgs": _NGLwdg,
                 "linked_data_arrays": _np.ndarray,
                 "linked_ax_wdgs": _AxesWidget,
                 "linked_figs": _mplFigure,
                 }

class ChangeInNGLWidgetListener(object):

    def __init__(self, ngl_wdg, list_mpl_objects_to_update, pos):
        self.ngl_wdg = ngl_wdg
        self.list_mpl_objects_to_update = list_mpl_objects_to_update
        self.pos = pos
    def __call__(self, change):
        self.ngl_wdg.isClick = False
        _idx = change["new"]
        try:
            for idot in self.list_mpl_objects_to_update:
                update2Dlines(idot, self.pos[_idx, 0], self.pos[_idx, 1])
            #print("caught index error with index %s (new=%s, old=%s)" % (_idx, change["new"], change["old"]))
        except IndexError as e:
            for idot in self.list_mpl_objects_to_update:
                update2Dlines(idot, self.pos[0, 0], self.pos[0, 1])
            print("caught index error with index %s (new=%s, old=%s)" % (_idx, change["new"], change["old"]))
        #print("set xy = (%s, %s)" % (x[_idx], y[_idx]))

class GeometryInNGLWidget(object):
    r"""
    returns an object that is aware of where its geometries are located in the NGLWidget their representation status

    The object exposes two methods, show and hide, to automagically know what to do
    """

    def __init__(self, geom, ngl_wdg, list_of_repr_dicts=None,
                 color_molecule_hex='Element', n_small=10):
        self.lives_at_components = []
        self.geom = geom
        self.ngl_wdg = ngl_wdg
        self.have_repr = []

        sticky_rep = 'cartoon'
        if self.geom[0].top.n_residues < n_small:
            sticky_rep = 'ball+stick'
        if list_of_repr_dicts is None:
            list_of_repr_dicts = [{'repr_type': sticky_rep, 'selection': 'all'}]

        self.list_of_repr_dicts = list_of_repr_dicts
        self.color_molecule_hex = color_molecule_hex
        self.color_dot = color_molecule_hex
        if isinstance(self.color_molecule_hex, str) and color_molecule_hex == 'Element':
            self.color_dot = 'red'

    def show(self):
        # Show can mean either
        #    - add a whole new component (case 1)
        #    - add the representation again to a representation-less component (case 2)

        # CASE 1
        if self.is_empty() or self.all_reps_are_on():
            if len(self.have_repr) == self.geom.n_frames:
                print("arrived at the end")
                component = None
            else:
                idx = len(self.have_repr)

                self.ngl_wdg.add_trajectory(self.geom[idx])
                self.lives_at_components.append(len(self.ngl_wdg._ngl_component_ids) - 1)
                self.ngl_wdg.clear_representations(component=self.lives_at_components[-1])
                self.have_repr.append(True)
                component = self.lives_at_components[-1]

        # CASE 2
        elif self.any_rep_is_off():  # Some are living in the widget already but have no rep
            idx = _np.argwhere(~_np.array(self.have_repr))[0].squeeze()
            component = self.lives_at_components[idx]
            self.have_repr[idx] = True
        else:
            raise Exception("This situation should not arise. This is a bug")

        if component is not None:
            for irepr in self.list_of_repr_dicts:
                self.ngl_wdg.add_representation(irepr['repr_type'],
                                                selection=irepr['selection'],
                                                component=component,
                                                color=self.color_molecule_hex)

    def hide(self):
        if self.is_empty() or self.all_reps_are_off():
            print("nothing to hide")
            pass
        elif self.any_rep_is_on():  # There's represented components already in the widget
            idx = _np.argwhere(self.have_repr)[-1].squeeze()
            self.ngl_wdg.clear_representations(component=self.lives_at_components[idx])
            self.have_repr[idx] = False
        else:
            raise Exception("This situation should not arise. This is a bug")

    # Quickhand methods for knowing what's up
    def is_empty(self):
        if len(self.have_repr) == 0:
            return True
        else:
            return False

    def all_reps_are_off(self):
        if len(self.have_repr) == 0:
            return True
        else:
            return _np.all(~_np.array(self.have_repr))

    def all_reps_are_on(self):
        if len(self.have_repr) == 0:
            return False
        else:
            return _np.all(self.have_repr)

    def any_rep_is_off(self):
        return _np.any(~_np.array(self.have_repr))

    def any_rep_is_on(self):
        return _np.any(self.have_repr)

    def is_visible(self):
        if self.is_empty() or self.all_reps_are_off():
            return False
        else:
            return True

def link_ax_w_pos_2_nglwidget(ax, pos, ngl_wdg,
                              crosshairs=True,
                              dot_color='red',
                              band_width=None,
                              radius=False,
                              directionality=None,
                              exclude_coord=None,
                              ):
    r"""
    Initial idea for this function comes from @arose, the rest is @gph82

    Parameters
    ----------
    ax : matplotlib axis object to be linked

    pos : ndarray of shape (N,2) with the positions of the geoms in the ngl_wdg

    crosshairs : Boolean or str
        If True, a crosshair will show where the mouse-click ocurred. If 'h' or 'v', only the horizontal or
        vertical line of the crosshair will be shown, respectively. If False, no crosshair will appear

    dot_color : Anything that yields matplotlib.colors.is_color_like(dot_color)==True
        Default is 'red'. dot_color='None' yields no dot

    band_width : None or iterable of len = 2
        If band_width is not None, the method tries to figure out on its own if
        there is an ascending coordinate and will include a moving band on :obj:ax
        of this width (in units of the axis along which the band is plotted)

        If the method cannot find an ascending coordinate, an exception is thrown

    directionality : str or None, default is None
        If not None, directionality can be either 'a2w' or 'w2a', meaning that connectivity
         between axis and widget will be only established as
         * 'a2w' : action in axis   triggers action in widget, but not the other way around
         * 'w2a' : action in widget triggers action in axis, but not the other way around

    exclude_coord : None or int , default is None
        The excluded coordinate will not be considered when computing the nearest-point-to-click.
        Typical use case is for visualize.traj to only compute distances horizontally along the time axis

    Returns
    -------

    axes_widget : :obj:`matplotlib.Axes.Axeswidget` that has been linked to the NGLWidget
    """

    assert directionality in [None, 'a2w', 'w2a'], "The directionality parameter has to be in [None, 'a2w', 'w2a'] " \
                                                   "not %s"%directionality

    assert crosshairs in [True, False, 'h', 'v'], "The crosshairs parameter has to be in [True, False, 'h','v'], " \
                                                   "not %s" % crosshairs
    ipos = _np.copy(pos)
    if _is_int(exclude_coord):
        ipos[:,exclude_coord] = 0

    # Are we in a sticky situation?
    if hasattr(ngl_wdg, '_GeomsInWid'):
        sticky = True
    else:
        assert ngl_wdg.trajectory_0.n_frames == pos.shape[0], \
            ("Mismatching frame numbers %u vs %u" % (ngl_wdg.trajectory_0.n_frames, pos.shape[0]))
        sticky = False

    # Basic interactive objects
    showclick_objs = []
    if crosshairs in [True, 'h']:
        lineh = ax.axhline(ax.get_ybound()[0], c="black", ls='--')
        setattr(lineh, 'whatisthis', 'lineh')
        showclick_objs.append(lineh)
    if crosshairs in [True, 'v']:
        linev = ax.axvline(ax.get_xbound()[0], c="black", ls='--')
        setattr(linev, 'whatisthis', 'linev')
        showclick_objs.append(linev)

    if _is_color_like(dot_color):
        pass
    else:
        raise TypeError('dot_color should be a matplotlib color')

    dot = ax.plot(pos[0,0],pos[0,1], 'o', c=dot_color, ms=7, zorder=100)[0]
    setattr(dot,'whatisthis','dot')
    list_mpl_objects_to_update = [dot]

    # Other objects, related to smoothing options
    if band_width is not None:
        if radius:
            band_width_in_pts = int(_np.round(pts_per_axis_unit(ax).mean() * _np.mean(band_width)))
            rad = ax.plot(pos[0, 0], pos[0, 1], 'o',
                          ms=_np.round(band_width_in_pts),
                          c='green', alpha=.25, markeredgecolor='None')[0]
            setattr(rad, 'whatisthis', 'dot')
            if not sticky:
                list_mpl_objects_to_update.append(rad)
        else:
            # print("Band_width(x,y) is %s" % (band_width))
            coord_idx = get_ascending_coord_idx(pos)
            if _np.ndim(coord_idx)>0 and len(coord_idx)==0:
                raise ValueError("Must have an ascending coordinate for band_width usage")
            band_width_in_pts = int(_np.round(pts_per_axis_unit(ax)[coord_idx] * band_width[coord_idx]))
            # print("Band_width in %s is %s pts"%('xy'[coord_idx], band_width_in_pts))

            band_call = [ax.axvline, ax.axhline][coord_idx]
            band_init = [ax.get_xbound, ax.get_ybound][coord_idx]
            band_type = ['linev',  'lineh'][coord_idx]
            band = band_call(band_init()[0],
                             lw=band_width_in_pts,
                             c="green", ls='-',
                             alpha=.25)
            setattr(band, 'whatisthis', band_type)
            list_mpl_objects_to_update.append(band)

    ngl_wdg.isClick = False

    CLA_listener = ClickOnAxisListener(ngl_wdg, crosshairs, showclick_objs, ax, pos,
                                       list_mpl_objects_to_update)

    NGL_listener = ChangeInNGLWidgetListener(ngl_wdg, list_mpl_objects_to_update, pos)
    # Connect axes to widget
    axes_widget = _AxesWidget(ax)
    if directionality in [None, 'a2w']:
        axes_widget.connect_event('button_release_event', CLA_listener)

    # Connect widget to axes
    if directionality in [None, 'w2a']:
        ngl_wdg.observe(NGL_listener, "frame", "change")

    ngl_wdg.center()
    return axes_widget