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.utils.testing import ignore_warnings

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)

    @raises(ValueError)
    def test_sgd_bad_alpha_for_optimal_learning_rate(self):
        # Check whether expected ValueError on bad alpha, i.e. 0
        # since alpha is used to compute the optimal learning rate
        self.factory(alpha=0, learning_rate="optimal")

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_balanced(self):
        # partial_fit with class_weight='balanced' not supported"""
        assert_raises_regexp(ValueError,
                             "class_weight 'balanced' is not supported for "
                             "partial_fit. In order to use 'balanced' weights, "
                             "use compute_class_weight\('balanced', 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='balanced').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_balanced_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 balanced class_weight
        clf_balanced = self.factory(alpha=0.0001, n_iter=1000,
                                    class_weight="balanced",
                                    shuffle=False).fit(X, y)
        assert_almost_equal(metrics.f1_score(y, clf_balanced.predict(X), average='weighted'), 0.96,
                            decimal=1)

        # Make sure that in the balanced case it does not change anything
        # to use "balanced"
        assert_array_almost_equal(clf.coef_, clf_balanced.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 balanced class_weight enabled
        clf = self.factory(n_iter=1000, class_weight="balanced", 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="balanced", 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_partial_fit_multiclass_average(self):
        third = X2.shape[0] // 3
        clf = self.factory(alpha=0.01, average=X2.shape[0])
        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,))

        clf.partial_fit(X2[third:], Y2[third:])
        assert_equal(clf.coef_.shape, (3, X2.shape[1]))
        assert_equal(clf.intercept_.shape, (3,))

    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)

    @ignore_warnings
    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.predict([[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_))


# Authors: Vlad Niculae, Mathieu Blondel
# License: BSD 3 clause
"""
=========================
Multilabel classification
=========================

This example simulates a multi-label document classification problem. The
dataset is generated randomly based on the following process:

    - pick the number of labels: n ~ Poisson(n_labels)
    - n times, choose a class c: c ~ Multinomial(theta)
    - pick the document length: k ~ Poisson(length)
    - k times, choose a word: w ~ Multinomial(theta_c)

In the above process, rejection sampling is used to make sure that n is more
than 2, and that the document length is never zero. Likewise, we reject classes
which have already been chosen.  The documents that are assigned to both
classes are plotted surrounded by two colored circles.

The classification is performed by projecting to the first two principal
components found by PCA and CCA for visualisation purposes, followed by using
the :class:`sklearn.multiclass.OneVsRestClassifier` metaclassifier using two
SVCs with linear kernels to learn a discriminative model for each class.
Note that PCA is used to perform an unsupervised dimensionality reduction,
while CCA is used to perform a supervised one.

Note: in the plot, "unlabeled samples" does not mean that we don't know the
labels (as in semi-supervised learning) but that the samples simply do *not*
have a label.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt

from sklearn.datasets import make_multilabel_classification
from sklearn.multiclass import OneVsRestClassifier
from sklearn.svm import SVC
from sklearn.preprocessing import LabelBinarizer
from sklearn.decomposition import PCA
from sklearn.cross_decomposition import CCA

def plot_hyperplane(clf, min_x, max_x, linestyle, label):
    # get the separating hyperplane
    w = clf.coef_[0]
    a = -w[0] / w[1]
    xx = np.linspace(min_x - 5, max_x + 5)  # make sure the line is long enough
    yy = a * xx - (clf.intercept_[0]) / w[1]
    plt.plot(xx, yy, linestyle, label=label)

def plot_subfigure(X, Y, subplot, title, transform):
    if transform == "pca":
        X = PCA(n_components=2).fit_transform(X)
    elif transform == "cca":
        X = CCA(n_components=2).fit(X, Y).transform(X)
    else:
        raise ValueError

    min_x = np.min(X[:, 0])
    max_x = np.max(X[:, 0])

    min_y = np.min(X[:, 1])
    max_y = np.max(X[:, 1])

    classif = OneVsRestClassifier(SVC(kernel='linear'))
    classif.fit(X, Y)

    plt.subplot(2, 2, subplot)
    plt.title(title)

    zero_class = np.where(Y[:, 0])
    one_class = np.where(Y[:, 1])
    plt.scatter(X[:, 0], X[:, 1], s=40, c='gray')
    plt.scatter(X[zero_class, 0], X[zero_class, 1], s=160, edgecolors='b',
               facecolors='none', linewidths=2, label='Class 1')
    plt.scatter(X[one_class, 0], X[one_class, 1], s=80, edgecolors='orange',
               facecolors='none', linewidths=2, label='Class 2')

    plot_hyperplane(classif.estimators_[0], min_x, max_x, 'k--',
                    'Boundary\nfor class 1')
    plot_hyperplane(classif.estimators_[1], min_x, max_x, 'k-.',
                    'Boundary\nfor class 2')
    plt.xticks(())
    plt.yticks(())

    plt.xlim(min_x - .5 * max_x, max_x + .5 * max_x)
    plt.ylim(min_y - .5 * max_y, max_y + .5 * max_y)
    if subplot == 2:
        plt.xlabel('First principal component')
        plt.ylabel('Second principal component')
        plt.legend(loc="upper left")

plt.figure(figsize=(8, 6))

X, Y = make_multilabel_classification(n_classes=2, n_labels=1,
                                      allow_unlabeled=True,
                                      random_state=1)

plot_subfigure(X, Y, 1, "With unlabeled samples + CCA", "cca")
plot_subfigure(X, Y, 2, "With unlabeled samples + PCA", "pca")

X, Y = make_multilabel_classification(n_classes=2, n_labels=1,
                                      allow_unlabeled=False,
                                      random_state=1)

plot_subfigure(X, Y, 3, "Without unlabeled samples + CCA", "cca")
plot_subfigure(X, Y, 4, "Without unlabeled samples + PCA", "pca")

plt.subplots_adjust(.04, .02, .97, .94, .09, .2)
plt.show()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This file is part of OpenDrift.
#
# OpenDrift is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 2
#
# OpenDrift is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with OpenDrift.  If not, see <https://www.gnu.org/licenses/>.
#
# Copyright 2015, Knut-Frode Dagestad, MET Norway

import unittest
from datetime import datetime, timedelta

import numpy as np

from opendrift.models.oceandrift import OceanDrift
from opendrift.models.leeway import Leeway
from opendrift.models.openoil import OpenOil
from opendrift.readers import reader_netCDF_CF_generic
from opendrift.readers import reader_ROMS_native
from opendrift.readers import reader_global_landmask
from opendrift.readers import reader_constant
from opendrift.readers import reader_lazy
from opendrift.readers import reader_from_url
from opendrift.models.pelagicegg import PelagicEggDrift
from opendrift.readers import reader_current_from_track

o = OceanDrift(loglevel=20)

reader_list = [
    'www.nonexistingurl.com',
    o.test_data_folder() +
        '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc',
    '/nonexistingdisk/nonexistingfile.ext',
    o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/AROME_MetCoOp_00_DEF_20160202_subset.nc']

class TestReaders(unittest.TestCase):
    """Tests for readers"""

    def test_adding_readers(self):
        o = OceanDrift()
        landmask = reader_global_landmask.Reader(
            extent=[-1.5, 7, 59, 64])
        r = reader_ROMS_native.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        o.add_reader([r, landmask])
        self.assertEqual(o.priority_list['land_binary_mask'],
                         ['roms native', 'global_landmask'])
        self.assertEqual(o.priority_list['x_sea_water_velocity'],
                         ['roms native'])
        # Switch order
        o = OceanDrift()
        o.add_reader([landmask, r])
        self.assertEqual(o.priority_list['land_binary_mask'],
                         ['global_landmask', 'roms native'])
        self.assertEqual(o.priority_list['x_sea_water_velocity'],
                         ['roms native'])

        # Test add_readers_from_list
        o = OceanDrift()
        o.add_readers_from_list(reader_list, lazy=False)
        self.assertEqual(o.priority_list['x_sea_water_velocity'],
                         ['roms native'])
        self.assertEqual(o.priority_list['x_wind'],
                         [o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/AROME_MetCoOp_00_DEF_20160202_subset.nc'])

    def test_repeated_run(self):
        # NOTE: this test fails if outfile is not None
        #outfile = 'leeway_test.nc'
        outfile = None
        o = OceanDrift(loglevel=50)
        o.set_config('drift:vertical_mixing', False)
        o.add_readers_from_list(reader_list)
        o.seed_elements(lon=14, lat=67.85,
                        time=datetime(2016, 2, 2, 12))
        o.run(steps=5, outfile=outfile)
        lon1 = o.get_property('lon')[0]
        # Repeated run with same object
        o.seed_elements(lon=14, lat=67.85,
                        time=datetime(2016, 2, 2, 12))
        o.run(steps=5, outfile=outfile)
        lon2 = o.get_property('lon')[0]
        # Third run, with different config
        o.seed_elements(lon=14, lat=67.85,
                        time=datetime(2016, 2, 2, 12),
                        wind_drift_factor=.1)
        o.run(steps=5)
        lon3 = o.get_property('lon')[0]
        # Fourth run, with different time
        o.reset()  # Reset is needed due to new start_time
        o.seed_elements(lon=14, lat=67.85,
                        time=datetime(2016, 2, 2, 13),
                        wind_drift_factor=.1)
        o.run(steps=5, outfile=outfile)
        lon4 = o.get_property('lon')[0]

        # Check results
        self.assertEqual(lon1[-1][0], lon2[-1][0])
        self.assertNotEqual(lon3[-1][0], lon2[-1][0])

        #os.remove(outfile)

    def test_reader_from_url(self):
        readers = reader_from_url(reader_list)
        self.assertIsNone(readers[0])
        self.assertTrue(isinstance(readers[1],
                                   reader_ROMS_native.Reader))
        self.assertIsNone(readers[2])
        self.assertTrue(isinstance(readers[3],
                                   reader_netCDF_CF_generic.Reader))

    def test_lazy_reader(self):
        o = OceanDrift(loglevel=20)
        lr = reader_lazy.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        self.assertFalse(lr.initialised)
        self.assertEqual(len(lr.covers_positions([15], [69])[0]), 1)
        self.assertEqual(len(lr.covers_positions([0], [0])[0]), 0)
        self.assertTrue(lr.initialised)

        # Make a corresponding, unlazy reader
        rr = reader_ROMS_native.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        self.assertEqual(len(rr.covers_positions([15], [69])[0]), 1)
        self.assertEqual(len(rr.covers_positions([0], [0])[0]), 0)

        # Check that both readers provide the same attributes
        for att in rr.__dict__:
            self.assertEqual(type(lr.__getattr__(att)),
                             type(getattr(rr, att)))
            if type(getattr(rr, att)) in [float, int, dict, str, list,
                                          datetime, timedelta, bool,
                                          np.float64]:
                self.assertEqual(lr.__getattr__(att),
                                 getattr(rr, att))
            elif type(getattr(rr, att)) in [np.ndarray]:
                self.assertIsNone(np.testing.assert_array_equal(
                                  lr.__getattr__(att),
                                  getattr(rr, att)))
            else:
                print('Skipping: ' + att + ' ' +
                      str(type(getattr(rr, att))))

    def test_lazy_reader_oildrift(self):
        o = OpenOil(loglevel=0)
        reader_constant_wind = \
            reader_constant.Reader({'x_wind':5, 'y_wind': 6,
                                    'sea_ice_area_fraction': 0})
            # Added ice area to prevent problems with masking
            # with older versions of netCDF library
        o.add_reader(reader_constant_wind)

        o.add_readers_from_list(reader_list, lazy=True)

        self.assertEqual(len(o._lazy_readers()), 4)
        o.seed_elements(lon=14, lat=67.85,
                        time=datetime(2016, 2, 2, 12))
        o.run(steps=5)
        print(o)  # Debug, this fails for old libraries
        self.assertEqual(len(o._lazy_readers()), 2)
        self.assertEqual(len(o.discarded_readers), 1)

    def test_ROMS_native_stranding(self):
        o = OceanDrift(loglevel=0)
        r = reader_ROMS_native.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        o.add_reader(r)
        o.set_config('general:use_auto_landmask', False)
        o.set_config('drift:vertical_mixing', False)
        o.set_config('environment:fallback:x_wind', 0)
        o.set_config('environment:fallback:y_wind', 10)
        o.seed_elements(lon=15.2, lat=68.3, time=r.start_time,
                        wind_drift_factor=.02,
                        number=10, radius=1000)
        o.run(steps=8)
        self.assertEqual(o.num_elements_deactivated(), 2)

    #def test_lazy_readers_and_corrupt_data(self):
    #    o = OceanDrift(loglevel=0)

    #    o.add_readers_from_list([o.test_data_folder() +
    #        '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc'])

    #    reader_constant_current_corrupt = \
    #        reader_constant.Reader({'x_sea_water_velocity': np.nan,
    #                                'y_sea_water_velocity': np.nan})
    #    o.add_reader(reader_constant_current_corrupt)
    #    o.add_readers_from_list([o.test_data_folder() +
    #        '2Feb2016_Nordic_sigma_3d/Arctic20_1to5Feb_2016.nc'])

    #    print o
    #    o.seed_elements(lon=14.5, lat=68, time=datetime(2016,2,4))
    #    o.set_config('environment:fallback:'x_wind', 0)
    #    o.set_config('environment:fallback:'y_wind', 0)
    #    o.set_config('environment:fallback:'x_sea_water_velocity', None)
    #    o.set_config('environment:fallback:'y_sea_water_velocity', None)
    #    o.set_config('environment:fallback:'land_binary_mask', 0)
    #    print o
    #    o.run(steps=1)

    #def test_oildrift_backwards(self):
    #    o = OpenOil(loglevel=20)
    #    reader_constant_wind = \
    #        reader_constant.Reader({'x_wind':5, 'y_wind': 6})
    #    o.add_reader(reader_constant_wind)

    #    o.add_readers_from_list(reader_list, lazy=True)

    #    self.assertEqual(len(o._lazy_readers()), 4)
    #    o.seed_elements(lon=14, lat=67.85,
    #                    time=datetime(2016, 2, 2, 12))
    #    o.set_config()
    #    o.run(steps=5)
    #    self.assertEqual(len(o._lazy_readers()), 2)
    #    self.assertEqual(len(o.discarded_readers), 1)

    #def test_lazy_reader_oildrift_real(self):
    #    o = OpenOil(loglevel=0)
    #    o.add_readers_from_file(o.test_data_folder() +
    #        '../../opendrift/scripts/data_sources.txt')

    #    o.seed_elements(lon=4, lat=60.0,
    #                    time=datetime(2018, 7, 2, 12))
    #    o.run(steps=5)
    #    print o

    def test_lazy_reader_leeway_compare(self):

        o1 = Leeway(loglevel=0)
        #o1.set_config('environment:fallback:land_binary_mask', 0)
        o1.required_variables = [r for r in o1.required_variables
                                 if r != 'land_binary_mask']
        o1.add_readers_from_list(reader_list, lazy=False)
        time = o1.readers['roms native'].start_time
        o1.seed_elements(lat=67.85, lon=14, time=time)
        o1.run(steps=5)

        o2 = Leeway(loglevel=20)
        #o2.set_config('environment:fallback:land_binary_mask', 0)
        o2.required_variables = [r for r in o1.required_variables
                                 if r != 'land_binary_mask']
        o2.add_readers_from_list(reader_list, lazy=True)
        o2.seed_elements(lat=67.85, lon=14, time=time)
        o2.run(steps=5)

        # Some differences in wind and current components
        # due to different coordinate system
        for var in o1.history.dtype.names:
            if var in ['x_wind', 'y_wind', 'x_sea_water_velocity',
                       'y_sea_water_velocity']:
                tolerance = 1
            else:
                tolerance = 5
            self.assertIsNone(np.testing.assert_array_almost_equal(
                o1.history[var], o2.history[var], tolerance))

    def test_constant_and_lazy_reader_leeway(self):
        cw = reader_constant.Reader({'x_wind':5, 'y_wind': 6})
        cc = reader_constant.Reader({'x_sea_water_velocity':0,
                                     'y_sea_water_velocity': .2})

        o = Leeway(loglevel=20)
        o.add_reader([cw, cc])
        o.add_readers_from_list(reader_list)
        o.set_config('environment:fallback:x_sea_water_velocity', 0.0)
        o.set_config('environment:fallback:y_sea_water_velocity', 0.1)
        time = datetime(2016,2,2,12)
        o.seed_elements(lat=67.85, lon=14, time=time)
        o.run(steps=2)
        self.assertAlmostEqual(o.elements.lat[0], 67.8548, 3)

    def test_automatic_landmask(self):
        o = OceanDrift(loglevel=20)
        self.assertRaises(ValueError, o.run)
        o.seed_elements(lon=4, lat=60, time=datetime(2016,9,1))
        o.run(steps=2)

    def test_reader_coverage(self):
        r = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
            '16Nov2015_NorKyst_z_surface/norkyst800_subset_16Nov2015.nc')
        # Element outside reader domain
        self.assertEqual(len(r.covers_positions(5, 80)[0]), 0)
        x, y = r.lonlat2xy(5, 80)
        self.assertRaises(ValueError, r.check_arguments,
                          'y_sea_water_velocity', r.start_time, x, y, 0)
        # Element inside reader domain
        self.assertEqual(len(r.covers_positions(5, 60)[0]), 1)
        x, y = r.lonlat2xy(5, 60)
        var, time, x2, y2, z2, outside = \
            r.check_arguments('y_sea_water_velocity', r.start_time, x, y, 0)
        self.assertEqual(var, ['y_sea_water_velocity'])
        self.assertEqual(time, r.start_time)
        self.assertEqual(x, x2)
        self.assertEqual(y, y2)
        self.assertEqual(0, z2)
        self.assertEqual(len(outside), 0)

    def test_outside_reader_time_coverage(self):
        o = PelagicEggDrift()
        reader = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
            '16Nov2015_NorKyst_z_surface/norkyst800_subset_16Nov2015.nc')
        o.add_reader(reader)
        o.set_config('environment:fallback:x_sea_water_velocity', 1)
        o.set_config('environment:fallback:land_binary_mask', 0)
        o.set_config('drift:vertical_mixing', False)
        o.seed_elements(lon=4.8, lat=60, number=1, time=reader.end_time)
        o.run(steps=2)
        # Check that fallback value is used when outside time coverage
        self.assertEqual(o.history['x_sea_water_velocity'][0][-1], 1.0)

    def test_reader_netcdf(self):
        """Check reader functionality."""

        reader1 = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
            '16Nov2015_NorKyst_z_surface/norkyst800_subset_16Nov2015.nc')
        reader2 = reader_ROMS_native.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        readers = [reader1, reader2]

        for r in readers:
            print(r)
            # Make four points:
            #  1) outside lower left, 2) lower left,  3) center of domain
            #  4) outside upper right
            # and assure that only 2) and 3) are marked as covered
            # Upper right is skipped, as lonlat2xy may lie slightly outside
            x = np.array([r.xmin - r.delta_x, r.xmin, (r.xmin + r.xmax)/2,
                          r.xmax + r.delta_x])
            y = np.array([r.ymin - r.delta_y, r.ymin, (r.ymin + r.ymax)/2,
                          r.ymax + r.delta_y])
            lons, lats = r.xy2lonlat(x,  y)
            covered = r.covers_positions(lons, lats, 0)[0]
            if len(covered) != 1:
                self.assertEqual(covered.tolist(), [1, 2])
            else:
                if covered == [2]:
                    print('#'*60)
                    print('#'*60)
                    print('WARNING: A point on the boundary is considered ' \
                          'outside after conversion x,y -> lon,lat -> x,y. ' \
                          'This is different from "standard", but is due to ' \
                          'rounding differences and not considered to be an ' \
                          'error. Numpy version is %s' % (np.__version__))
                    print('#'*60)
                    print('#'*60)
                else:
                    self.assertTrue(False)  # Should never happen!

            self.assertTrue(r.covers_time(r.start_time))
            self.assertFalse(r.covers_time(r.start_time - r.time_step))
            self.assertFalse(r.proj.crs.is_geographic)

    def test_vertical_profiles(self):

        norkyst3d = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
            '14Jan2016_NorKyst_z_3d/NorKyst-800m_ZDEPTHS_his_00_3Dsubset.nc')
        lon = np.array([4.73])
        lat = np.array([62.35])
        variables = ['x_sea_water_velocity', 'x_sea_water_velocity',
                     'sea_water_temperature']
        x,y = norkyst3d.lonlat2xy(lon, lat)
        data = norkyst3d.get_variables(variables,
                                       time=norkyst3d.start_time,
                                       x=x, y=y, z=[0, -100])
        self.assertEqual(data['z'][4], -25)
        self.assertEqual(data['z'][4], -25)
        self.assertAlmostEqual(data['sea_water_temperature'][:,0,0][7],
                         9.220000267028809)

    def test_vertical_interpolation(self):
        norkyst3d = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
            '14Jan2016_NorKyst_z_3d/NorKyst-800m_ZDEPTHS_his_00_3Dsubset.nc')
        lon = np.array([4.73, 4.75])
        lat = np.array([62.35, 62.30])
        z = np.array([0, -33])
        variables = ['x_sea_water_velocity', 'x_sea_water_velocity',
                     'sea_water_temperature']
        # Call get_variables_interpolated which interpolates both in
        # space (horizontally, vertically) and then in time
        data, profiles = norkyst3d.get_variables_interpolated(
                variables, profiles=['sea_water_temperature'],
                profiles_depth = [-100, 0],
                time = norkyst3d.start_time + timedelta(seconds=900),
                lon=lon, lat=lat, z=z)
        # Check surface value
        self.assertEqual(data['sea_water_temperature'][0],
                         profiles['sea_water_temperature'][0,0])
        # Check interpolated temperature at 33 m depth
        self.assertAlmostEqual(data['sea_water_temperature'][1],
                               8.32, 2)
        #import matplotlib.pyplot as plt
        #plt.plot(profiles['sea_water_temperature'][:,0])
        #plt.plot(profiles['sea_water_temperature'][:,1], 'r')
        #plt.show()

    def test_vertical_interpolation_sigma(self):
        nordic3d = reader_ROMS_native.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        lon = np.array([12.46, 12.46, 12.46])
        lat = np.array([68.21, 69.31, 69.31])
        z = np.array([-33, 0, -2500])
        x, y = nordic3d.lonlat2xy(lon, lat)
        variables = ['x_sea_water_velocity', 'y_sea_water_velocity',
                     'sea_water_temperature']
        # Call get_variables_interpolated which interpolates both in
        data = nordic3d.get_variables(variables,
                time = nordic3d.start_time + timedelta(seconds=900),
                x=x, y=y, z=z)
        self.assertAlmostEqual(data['sea_water_temperature'][0,60, 60],
                               3.447, 2)
                               #3.59, 2)
        self.assertAlmostEqual(data['sea_water_temperature'][-1,60, 60],
                               -0.783, 2)
                               #-0.803, 2)

    def test_get_environment(self):
        o = PelagicEggDrift(loglevel=0)
        reader_nordic = reader_ROMS_native.Reader(o.test_data_folder() + '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc', name='Nordic')
        reader_arctic = reader_netCDF_CF_generic.Reader(o.test_data_folder() + '2Feb2016_Nordic_sigma_3d/Arctic20_1to5Feb_2016.nc', name='Arctic')
        ######################################################
        # Vertical interpolation is another issue to be fixed:
        reader_nordic.zlevels = reader_arctic.z
        ######################################################
        o.add_reader([reader_nordic, reader_arctic])
        # One point covered only by Nordic, two points coverd
        # by both readers, and two points covered by none of the readers
        testlon = np.array((14.0, 20.0, 20.1, 4, 5))
        testlat = np.array((70.1, 76.0, 76.1, 60, 60))
        testz = np.random.uniform(0, 0, len(testlon))
        self.assertIsNone(np.testing.assert_array_almost_equal(
            [0], reader_nordic.covers_positions(testlon, testlat, testz)[0]))
        self.assertIsNone(np.testing.assert_array_almost_equal(
            [0, 1, 2],
            reader_arctic.covers_positions(testlon, testlat, testz)[0]))
        o.seed_elements(testlon, testlat, z=testz, time=reader_nordic.start_time)
        o.set_config('environment:fallback:land_binary_mask', 0)
        env, env_profiles, missing = \
            o.get_environment(list(o.required_variables),
                              reader_nordic.start_time,
                              testlon, testlat, testz,
                              o.required_profiles)
        self.assertAlmostEqual(env['sea_water_temperature'][0], 4.251, 2)
        self.assertAlmostEqual(env['sea_water_temperature'][1], 0.122, 3)
        self.assertAlmostEqual(env['sea_water_temperature'][4], 10.0)
        self.assertIsNone(np.testing.assert_array_almost_equal(
            missing, [False,False,False,False,False]))
        self.assertAlmostEqual(env_profiles['sea_water_temperature'][0,0],
                               4.251, 2)
        self.assertAlmostEqual(env_profiles['sea_water_temperature'][0,4], 10)
        #self.assertAlmostEqual(env_profiles['sea_water_temperature'][8,2], 10)
        self.assertAlmostEqual(env_profiles['sea_water_temperature'][7,2],
                               2.159, 3)
        # Get separate data
        env2, env_profiles2, missing2 = \
            o.get_environment(['x_sea_water_velocity', 'y_sea_water_velocity',
                               'sea_water_temperature'],
                              reader_nordic.start_time,
                              testlon, testlat, testz,
                              ['sea_water_temperature'])
        self.assertTrue(env_profiles2 is not None)
        self.assertEqual(set(env_profiles2.keys()),
            set(['z', 'sea_water_temperature']))
        # Get separate data, without profile
        env3, env_profiles3, missing3 = \
            o.get_environment(['x_sea_water_velocity', 'y_sea_water_velocity',
                               'sea_water_temperature'],
                              reader_nordic.start_time,
                              testlon, testlat, testz,
                              profiles=None)
        self.assertTrue(env_profiles3 is None)
        # Get separate data
        env4, env_profiles4, missing4 = \
            o.get_environment(['x_sea_water_velocity', 'y_sea_water_velocity',
                               'sea_water_temperature'],
                              reader_nordic.start_time,
                              testlon, testlat, testz,
                              ['sea_water_temperature'])

        self.assertIsNone(np.testing.assert_array_almost_equal(
            env['x_sea_water_velocity'],
            env2['x_sea_water_velocity']))
        self.assertIsNone(np.testing.assert_array_almost_equal(
            env_profiles2['sea_water_temperature'].ravel(),
            env_profiles4['sea_water_temperature'].ravel()))

    def test_constant_reader(self):
        o = OpenOil(loglevel=0)
        cw = reader_constant.Reader({'x_wind':5, 'y_wind': 6})
        cc = reader_constant.Reader({'x_sea_water_velocity':0, 'y_sea_water_velocity': .2})
        cs = reader_constant.Reader({'sea_water_temperature': 278})
        r = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
            '16Nov2015_NorKyst_z_surface/norkyst800_subset_16Nov2015.nc')
        o.add_reader([cw, cc, r])
        # TODO: should check why adding constant reader with
        #   sea_water_temperature gives Deprecated warning
        #o.add_reader([cw, cc, cs, r])
        o.seed_elements(lon=4, lat=60, time=r.start_time, number=5)
        o.run(steps=3)

    def test_clip_domain(self):
        o = OceanDrift(loglevel=50)
        r1 = reader_ROMS_native.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        r1.clip_boundary_pixels(20)
        r2 = reader_ROMS_native.Reader(o.test_data_folder() +
            '2Feb2016_Nordic_sigma_3d/Nordic-4km_SLEVELS_avg_00_subset2Feb2016.nc')
        self.assertEqual(r2.shape, (151, 81))
        self.assertEqual(r1.shape, (111, 41))
        self.assertEqual(r1.xmin, 20)

        o1 = OceanDrift(loglevel=50)
        o1.set_config('environment:fallback:x_sea_water_velocity', None)
        o1.add_reader(r1)
        o1.seed_elements(lon=15, lat=70.1, time=r1.start_time)
        o1.set_config('environment:fallback:land_binary_mask', 0)
        o1.run(time_step=3600*3, duration=timedelta(hours=48))
        o2 = OceanDrift(loglevel=50)
        o2.set_config('environment:fallback:x_sea_water_velocity', None)
        o2.add_reader(r2)
        o2.seed_elements(lon=15, lat=70.1, time=r1.start_time)
        o2.set_config('environment:fallback:land_binary_mask', 0)
        o2.run(time_step=3600*3, duration=timedelta(hours=48))
        # Compare
        lat1 = o1.get_property('lat')[0]
        lat2 = o2.get_property('lat')[0]
        self.assertEqual(len(lat1), 13)
        self.assertEqual(len(lat2), 17)
        self.assertIsNone(np.testing.assert_allclose(
                            lat1[0:12], lat2[0:12]))
        # Test reader netCDF_CF_generic
        r = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
            '16Nov2015_NorKyst_z_surface/norkyst800_subset_16Nov2015.nc')
        self.assertEqual(r.shape, (301, 201))
        o3 = OceanDrift(loglevel=50)
        o3.set_config('environment:fallback:x_sea_water_velocity', None)
        o3.set_config('environment:fallback:land_binary_mask', 0)
        o3.add_reader(r)
        o3.seed_elements(lon=4.36, lat=61.7, time=r.start_time)
        o3.run(steps=24)
        r.clip_boundary_pixels(10)
        self.assertEqual(r.shape, (281, 181))
        o4 = OceanDrift(loglevel=50)
        o4.set_config('environment:fallback:x_sea_water_velocity', None)
        o4.set_config('environment:fallback:land_binary_mask', 0)
        o4.add_reader(r)
        o4.seed_elements(lon=4.36, lat=61.7, time=r.start_time)
        o4.run(steps=24)
        # Compare
        lat3 = o3.get_property('lat')[0]
        lat4 = o4.get_property('lat')[0]
        self.assertEqual(len(lat3), 25)
        self.assertEqual(len(lat4), 13)
        self.assertIsNone(np.testing.assert_allclose(
                            lat3[0:12], lat4[0:12]))

    def test_reader_current_from_track(self):
        """Check if extrapolated currents are of expected value"""
        obslon = [3.1, 3.123456]
        obslat = [61.1, 61.132198]
        obstime = [datetime(2015, 11, 16, 0), datetime(2015, 11, 16, 6)]

        o = OceanDrift(loglevel=20)
        reader_wind = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
                    '16Nov2015_NorKyst_z_surface/arome_subset_16Nov2015.nc')

        reader_current = reader_current_from_track.Reader(obslon, obslat, obstime,
                    wind_east=0, wind_north=0, windreader=reader_wind, wind_factor=0.018)
        self.assertAlmostEqual(reader_current.x_sea_water_velocity.data[0],0.2236, 4)

    def test_valid_minmax(self):
        """Check that invalid values are replaced with fallback."""
        o = OceanDrift(loglevel=20)
        from opendrift.readers.basereader import variables
        minval = variables.standard_names['x_wind']['valid_min']
        # Setting valid_min to -5, to check that replacement works
        variables.standard_names['x_wind']['valid_min'] = -5
        reader_wind = reader_netCDF_CF_generic.Reader(o.test_data_folder() +
                    '16Nov2015_NorKyst_z_surface/arome_subset_16Nov2015.nc')
        o.add_reader(reader_wind)
        o.set_config('environment:fallback:x_sea_water_velocity', 0)
        o.set_config('environment:fallback:x_wind', 2.0)
        o.set_config('environment:fallback:y_sea_water_velocity', 0)
        o.set_config('environment:fallback:land_binary_mask', 0)
        o.seed_elements(lon=4, lat=60, time=reader_wind.start_time)
        o.run(steps=1)
        variables.standard_names['x_wind']['valid_min'] = minval  # reset
        w = o.get_property('x_wind')[0][0]
        self.assertAlmostEqual(w, 2.0, 1)

    def test_valid_minmax_nanvalues(self):
        from opendrift.readers.basereader import variables
        # Reducing max current speed to test masking
        maxval = variables.standard_names['x_sea_water_velocity']['valid_max']
        variables.standard_names['x_sea_water_velocity']['valid_max'] = .1
        o = OceanDrift(loglevel=20)
        o.set_config('environment:fallback:land_binary_mask', 0)
        norkyst = reader_netCDF_CF_generic.Reader(o.test_data_folder() + '14Jan2016_NorKyst_z_3d/NorKyst-800m_ZDEPTHS_his_00_3Dsubset.nc')
        o.add_reader(norkyst)

        o.seed_elements(lon=4.95, lat=62, number=10, time=norkyst.start_time)
        o.run(steps=2)
        variables.standard_names['x_sea_water_velocity']['valid_max'] = maxval  # reset
        u = o.get_property('x_sea_water_velocity')[0]
        self.assertAlmostEqual(u.max(), -.069, 3)  # Some numerical error allowed

if __name__ == '__main__':
    unittest.main()


"""
===========================================
Comparison of F-test and mutual information
===========================================

This example illustrates the differences between univariate F-test statistics
and mutual information.

We consider 3 features x_1, x_2, x_3 distributed uniformly over [0, 1], the
target depends on them as follows:

y = x_1 + sin(6 * pi * x_2) + 0.1 * N(0, 1), that is the third features is completely irrelevant.

The code below plots the dependency of y against individual x_i and normalized
values of univariate F-tests statistics and mutual information.

As F-test captures only linear dependency, it rates x_1 as the most
discriminative feature. On the other hand, mutual information can capture any
kind of dependency between variables and it rates x_2 as the most
discriminative feature, which probably agrees better with our intuitive
perception for this example. Both methods correctly marks x_3 as irrelevant.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
from sklearn.feature_selection import f_regression, mutual_info_regression

np.random.seed(0)
X = np.random.rand(1000, 3)
y = X[:, 0] + np.sin(6 * np.pi * X[:, 1]) + 0.1 * np.random.randn(1000)

f_test, _ = f_regression(X, y)
f_test /= np.max(f_test)

mi = mutual_info_regression(X, y)
mi /= np.max(mi)

plt.figure(figsize=(15, 5))
for i in range(3):
    plt.subplot(1, 3, i + 1)
    plt.scatter(X[:, i], y)
    plt.xlabel("$x_{}$".format(i + 1), fontsize=14)
    if i == 0:
        plt.ylabel("$y$", fontsize=14)
    plt.title("F-test={:.2f}, MI={:.2f}".format(f_test[i], mi[i]),
              fontsize=16)
plt.show()


'''
Code for error analysis

Copyright (c) 2014, Helen Ramsden
All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
disclaimer in the documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

import Image, ImageChops
import numpy as np
from scipy import ndimage
from GenericFunctions import checkOSpath, adjust_spines,st
import matplotlib.pyplot as plt
plt.rc('ytick', labelsize=12)
plt.rc('xtick', labelsize=12)
plt.rc('axes', labelsize=12)
plt.rc('axes', titlesize=20)

def checksegmented(segmaskfilepath,filedict,resultsfilepath):
	'''
	FUNCTION runs through all segmented masks and checks location of centre of mass and size of mask
	input SegmentedMask/
	output is a list containing name of file, size of mask,
	'''
	newfile = open(resultsfilepath + 'maskstatssize.txt','w')
	for f in filedict:
		# print f
		newfile.write(f )
		for filepath in [segmaskfilepath]:#, segmask2filepath]:
			maskim =  Image.open(filepath+ f).convert('L') # need to convert to 8 bit (not rgb)
			maskim = ImageChops.invert(maskim)
			maskarray = np.array(maskim)
			# print maskarray.shape
			com = ndimage.measurements.center_of_mass(maskarray)
			blackpixels = np.nonzero(maskarray==0)
			whitepixels = np.nonzero(maskarray>0)
			# print len(blackpixels[0]),len(whitepixels[0])
			masksize = len(blackpixels[0])
			newfile.write('\t' + '\t'.join([str(com[0]),str(com[1]),str(masksize)]))
		newfile.write('\n')

def plotmi():
	'''
	Plot the distribution of MI scores from the registration output
	'''
	milog = np.loadtxt('alllogdata.txt',delimiter = '\t',dtype = float,usecols=[2,3])
	diffs = milog[:,0] - milog[:,1]
	milognew = np.ma.masked_array(milog, np.isnan(milog))
	diffsnew = np.ma.masked_array(diffs, np.isnan(diffs))

	# Get rid of nans
	milogmaskpre = np.ma.masked_array(milog[:,0],np.isnan(milog[:,0]))
	milogmaskpost = np.ma.masked_array(milog[:,1],np.isnan(milog[:,1]))
	milogmaskpre = milogmaskpre[milogmaskpre>-1000]
	milogmaskpost = milogmaskpost[milogmaskpost>-1000]

	fig = plt.figure(figsize = (8,8))
	fig.subplots_adjust(bottom=0.2)
	fig.subplots_adjust(left=0.2)
	ax = fig.add_subplot(1,1,1)
	adjust_spines(ax, ['left','bottom'])
	cols = ['r','b','k','g','m','y']
#	histpre, binspre = np.histogram(milogmaskpre, bins=20)
#	histpost, binspre = np.histogram(milogmaskpre, bins=20)
	ax.hist(milogmaskpre, bins=20,histtype='step',color='b', range = [-600,0])
	ax.hist(milogmaskpost,bins=20,histtype='step',color='g', range = [-600,0]) # normed=True,
	[xmin, xmax, ymin, ymax] = ax.axis()
	ax.set_yticks([ymin,ymax])
	ax.set_yticklabels([int(ymin),int(ymax)], fontsize = 25)
	ax.xaxis.set_label_coords(0.5, -0.15)
	ax.set_xticks([xmin,xmax])
	ax.set_xticklabels([xmin,xmax], fontsize = 25)
	ax.set_xlabel('Joint Entropy', fontsize = 25)
	ax.set_ylabel('Frequency', fontsize = 25)
	ax.yaxis.set_label_coords( -0.15, 0.5)

	fig.savefig('MIlogdata.png', transparent = True)