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_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_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_))


#import pdb # pause code for debugging at pdb.set_trace()
import numpy as np
import toolbox as tool
import slab_functions as sf
from pysac.plot.mayavi_seed_streamlines import SeedStreamline
import matplotlib.pyplot as plt
from mayavi import mlab
import gc
#import move_seed_points as msp
import mayavi_plotting_functions as mpf
import dispersion_diagram
import img2vid as i2v
from functools import partial
import os

# ================================
# Preamble: set mode options and view parameters
# ================================

# What mode do you want? OPTIONS:
mode_options = ['slow-kink-surf', 'slow-saus-surf', 'slow-saus-body-3',
                'slow-kink-body-3', 'slow-saus-body-2', 'slow-kink-body-2',
                'slow-saus-body-1', 'slow-kink-body-1', 'fast-saus-body-1',
                'fast-kink-body-1', 'fast-saus-body-2', 'fast-kink-body-2',
                'fast-saus-body-3', 'fast-kink-body-3', 'fast-kink-surf',
                'fast-saus-surf', 'shear-alfven', 'shear-alfven-broadband']

# Which angle shall we view from? OPTIONS:
view_options = ['front', 'front-parallel', 'top', 'top-parallel', 'front-top',
                'front-side', 'front-top-side']

# Uniform lighting?
#uniform_light = True
uniform_light = False

show_density = False
show_density_pert = False
show_mag = False
show_mag_scale = False
show_mag_fade = False
show_mag_vec = False
show_vel_front = False
show_vel_front_pert = False
show_vel_top = False
show_vel_top_pert = False
show_disp_top = False
show_disp_front = False
show_axes = False
show_axis_labels = False
show_mini_axis = False
show_boundary = False

# Uncomment the parametrer you would like to see
# No density perturbations or vel/disp pert for alfven modes.

#show_density = True
#show_density_pert = True
show_mag = True
#show_mag_scale = True #must also have show_mag = True
#show_mag_fade = True
#show_mag_vec = True
#show_vel_front = True
#show_vel_front_pert = True
#show_vel_top = True
#show_vel_top_pert = True
#show_disp_top = True
#show_disp_front = True
show_axes = True
#show_axis_labels = True
show_mini_axis = True
show_boundary = True

# Visualisation modules in string form for file-names
vis_modules = [show_density, show_density_pert, show_mag, show_mag_scale,
               show_mag_fade, show_mag_vec, show_vel_front, show_vel_front_pert,
               show_vel_top, show_vel_top_pert, show_disp_top, show_disp_front]
vis_modules_strings = ['show_density', 'show_density_pert', 'show_mag', 'show_mag_scale',
                       'show_mag_fade', 'show_mag_vec', 'show_vel_front', 'show_vel_front_pert',
                       'show_vel_top', 'show_vel_top_pert', 'show_disp_top', 'show_disp_front']
vis_mod_string = ''
for i, j in enumerate(vis_modules):
    if vis_modules[i]:
        vis_mod_string = vis_mod_string + vis_modules_strings[i][5:] + '_'

# Set to True if you would like the dispersion diagram with chosen mode highlighted.
show_dispersion = False
#show_dispersion = True

# Wanna see the animation? Of course you do
#show_animation = False
show_animation = True

# Basic plot to see which eigensolutions have been found.
show_quick_plot = False
#show_quick_plot = True

# Video resolution
#res = (1920,1080) # There is a problem with this resolution- height must be odd number - Mayavi bug apparently
res = tuple(101 * np.array((16,9)))
#res = tuple(51 * np.array((16,9)))
#res = tuple(21 * np.array((16,9)))

number_of_frames = 1

# Frames per second of output video
fps = 20

#save_images = False
save_images = True

make_video = False
#make_video = True

# Where should I save the animation images/videos?
os.path.abspath(os.curdir)
os.chdir('..')
save_directory = os.path.join(os.path.abspath(os.curdir), '3D_vis_animations')

# Where should I save the dispersion diagrams?
save_dispersion_diagram_directory = os.path.join(os.path.abspath(os.curdir), '3D_vis_dispersion_diagrams')

# ================================
# Visualisation set-up
# ================================

# Variable definitions (for reference):
# x = k*x
# y = k*y
# z = k*z
# W = omega/k
# K = k*x_0
# t = omega*t

# Loop through selected modes
for mode_ind in [0]:#range(8,14): # for all others. REMEMBER SBB pparameters
#for mode_ind in [14,15]: #for fast body surf. REMEMBER SBS parameters
#for mode_ind in [16, 17]:
#for mode_ind in [13]: #for an individual mode
#for mode_ind in range(2,14):
    if mode_ind not in range(len(mode_options)):
        raise NameError('Mode not in mode_options')

    # (note that fast surface modes, i.e. 14 and 15, can only be
    # found with SBS parameters in slab_functions...)
    mode = mode_options[mode_ind]

    # Specify oscillation parameters
    if 'slow' in mode and 'surf' in mode or 'alfven' in mode:
        K = 2.
    elif 'slow' in mode and 'body' in mode:
        K = 8.
    elif 'fast' in mode and 'body-1' in mode:
        K = 8.
    elif 'fast' in mode and 'body-2' in mode:
        K = 15.
    elif 'fast' in mode and 'body-3' in mode:
        K = 22.
    elif 'fast' in mode and 'surf' in mode:
        K = 8.
    else:
        raise NameError('Mode not found')

# Specify density ratio R1 := rho_1 / rho_0
#    R1 = 1.5 # Higher denisty on left than right
#    R1 = 1.8
#    R1 = 1.9 # Disp_diagram will only work for R1=1.5, 1.8, 2.0
    R1 = 2. # Symmetric slab

    # Reduce number of variables in dispersion relation
    disp_rel_partial = partial(sf.disp_rel_asym, R1=R1)

    # find eigenfrequencies W (= omega/k) within the range Wrange for the given parameters.
    Wrange1 = np.linspace(0., sf.cT, 11)
    Wrange2 = np.linspace(sf.cT, sf.c0, 401)
    Wrange3 = np.linspace(sf.c0, sf.c2, 11)

    Woptions_slow_surf = np.real(tool.point_find(disp_rel_partial, np.array(K), Wrange1, args=None).transpose())
    Woptions_slow_body = np.real(tool.point_find(disp_rel_partial, np.array(K), Wrange2, args=None).transpose())
    Woptions_fast = np.real(tool.point_find(disp_rel_partial, np.array(K), Wrange3, args=None).transpose())

    # Remove W values that are very close to characteristic speeds - these are spurious solutions
    tol = 1e-2
    indices_to_rm = []
    for i, w in enumerate(Woptions_slow_surf):
        spurious_roots_diff = abs(np.array([w, w - sf.c0, w - sf.c1(R1), w - sf.c2, w - sf.vA]))
        if min(spurious_roots_diff) < tol or w < 0 or w > sf.cT:
            indices_to_rm.append(i)
    Woptions_slow_surf = np.delete(Woptions_slow_surf, indices_to_rm)
    Woptions_slow_surf.sort()

    indices_to_rm = []
    for i, w in enumerate(Woptions_slow_body):
        spurious_roots_diff = abs(np.array([w, w - sf.c0, w - sf.c1(R1), w - sf.c2, w - sf.vA]))
        if min(spurious_roots_diff) < tol or w < sf.cT or w > sf.c0:
            indices_to_rm.append(i)
    Woptions_slow_body = np.delete(Woptions_slow_body, indices_to_rm)
    Woptions_slow_body.sort()

    indices_to_rm = []
    for i, w in enumerate(Woptions_fast):
        spurious_roots_diff = abs(np.array([w, w - sf.c0, w - sf.c1(R1), w - sf.c2, w - sf.vA]))
        if min(spurious_roots_diff) < tol or w < sf.c0 or w > min(sf.c1, sf.c2):
            indices_to_rm.append(i)
    Woptions_fast = np.delete(Woptions_fast, indices_to_rm)
    Woptions_fast.sort()

    # remove any higher order slow body modes - we only want to do the first 3 saus/kink
    if len(Woptions_slow_body) > 6:
        Woptions_slow_body = np.delete(Woptions_slow_body, range(len(Woptions_slow_body) - 6))
    Woptions = np.concatenate((Woptions_slow_surf, Woptions_slow_body, Woptions_fast))

    # set W to be the eigenfrequency for the requested mode
    if 'fast-saus-body' in mode or 'fast-kink-surf' in mode:
        W = Woptions_fast[-2]
    elif 'fast-kink-body' in mode or 'fast-saus-surf' in mode:
        W = Woptions_fast[-1]
    elif 'slow' in mode and 'surf' in mode:
        W = Woptions_slow_surf[mode_ind]
    elif 'slow' in mode and 'body' in mode:
        W = Woptions_slow_body[mode_ind-2]

    if 'alfven' in mode:
        W = sf.vA
    else:
        W = np.real(W)

    # Quick plot to see if we are hitting correct mode
    if show_quick_plot:
        plt.plot([K] * len(Woptions), Woptions, '.')
        plt.plot(K+0.5, W, 'go')
        plt.xlim([0,23])
        plt.show()

# ================================
# Dispersion diagram
# ================================

    if show_dispersion:
        if 'alfven' in mode:
            raise NameError('Disperion plot requested for an alfven mode. Cant do that.')

        dispersion_diagram.dispersion_diagram(mode_options, mode,
                                              disp_rel_partial, K, W, R1)
#        plt.tight_layout() # seems to make it chop the sides off with this
        plt.savefig(os.path.join(save_dispersion_diagram_directory, 'R1_' + str(R1) + '_' + mode + '.png') )
        plt.close()

# ================================
# Animation
# ================================

    if show_animation:
        print('Starting ' + mode)

        # set grid parameters
        xmin = -2.*K
        xmax = 2.*K
        ymin = 0.
        ymax = 4.
        zmin = 0.
        zmax = 2*np.pi

        # You can change ny but be careful changing nx, nz.
        nx = 300#100 #100 #300 gives us reduced bouncing of field lines for the same video size, but there is significant computational cost.
        ny = 300#100 #100 #100#20 #100
        nz = 300#100 #100
        nt = number_of_frames

        if nz % nt != 0:
            print("nt doesnt divide nz so there may be a problem with chopping in z direction for each time step")

        t_start = 0.
        t_end = zmax

        t = t_start

        xvals = np.linspace(xmin, xmax, nx)
        yvals = np.linspace(ymin, ymax, ny)
        zvals = np.linspace(zmin, zmax, nz, endpoint=False) # A fudge to give the height as exactly one wavelength

        x_spacing = max(nx, ny, nz) / nx
        y_spacing = max(nx, ny, nz) / ny
        z_spacing = max(nx, ny, nz) / nz

        # For masking points for plotting vector fields- have to do it manually due to Mayavi bug
        mod = int(4 * nx / 100)
        mod_y = int(np.ceil(mod / y_spacing))

        # Get the data xi=displacement, v=velocity, b=mag field
        if show_disp_top or show_disp_front:
            xixvals = np.real(np.repeat(sf.xix(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
            xizvals = np.real(np.repeat(sf.xiz(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
            xiyvals = np.real(np.repeat(sf.xiy(mode, xvals, zvals, t, W, K)[:, :, np.newaxis], ny, axis=2))

        if show_vel_front or show_vel_top:
            vxvals = np.real(np.repeat(sf.vx(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
            vzvals = np.real(np.repeat(sf.vz(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
            vyvals = np.real(np.repeat(sf.vy(mode, xvals, zvals, t, K)[:, :, np.newaxis], ny, axis=2))

        if show_vel_front_pert or show_vel_top_pert:
            vxvals = np.real(np.repeat(sf.vx_pert(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
            vzvals = np.real(np.repeat(sf.vz_pert(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
            vyvals = np.zeros_like(vxvals)

        # Axis is defined on the mag field so we have to set up this data
        bxvals = np.real(np.repeat(sf.bx(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
        byvals = np.real(np.repeat(sf.by(mode, xvals, zvals, t, K)[:, :, np.newaxis], ny, axis=2))
        bz_eq3d = np.repeat(sf.bz_eq(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2)
        bzvals = np.real(np.repeat(-sf.bz(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2) +
                         bz_eq3d)

        # displacement at the right and left boundaries
        if show_boundary:
            xix_boundary_r_vals = np.real(np.repeat(K + sf.xix_boundary(mode, zvals, t, W, K, R1, boundary='r')[:, np.newaxis], ny, axis=1))
            xix_boundary_l_vals = np.real(np.repeat(-K + sf.xix_boundary(mode, zvals, t, W, K, R1, boundary='l')[:, np.newaxis], ny, axis=1))

        if show_density:
            rho_vals = np.real(np.repeat(sf.rho(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))

        if show_density_pert:
            rho_vals = np.real(np.repeat(sf.rho_pert(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))

        bxvals_t = bxvals
        byvals_t = byvals
        bzvals_t = bzvals

        if show_disp_top or show_disp_front:
            xixvals_t = xixvals
            xiyvals_t = xiyvals
            xizvals_t = xizvals

        if show_vel_top or show_vel_top_pert or show_vel_front or show_vel_front_pert:
            vxvals_t = vxvals
            vyvals_t = vyvals
            vzvals_t = vzvals

        if show_boundary:
            xix_boundary_r_vals_t = xix_boundary_r_vals
            xix_boundary_l_vals_t = xix_boundary_l_vals

        if show_density or show_density_pert:
            rho_vals_t = rho_vals

# ================================
# Starting figure and visualisation modules
# ================================

        zgrid_zy, ygrid_zy = np.mgrid[0:nz:(nz)*1j,
                          0:ny:(ny)*1j]

        fig = mlab.figure(size=res) # (1920, 1080) for 1080p , tuple(101 * np.array((16,9))) #16:9 aspect ratio for video upload

        # Spacing of grid so that we can display a visualisation cube without having the same number of grid points in each dimension
        spacing =  np.array([x_spacing, z_spacing, y_spacing])

        if show_density or show_density_pert:
            # Scalar field density
            rho = mlab.pipeline.scalar_field(rho_vals_t, name="density", figure=fig)
            rho.spacing = spacing
            mpf.volume_red_blue(rho, rho_vals_t)

        #Masking points
        if show_mag_vec:
            bxvals_mask_front_t, byvals_mask_front_t, bzvals_mask_front_t = mpf.mask_points(bxvals_t, byvals_t, bzvals_t,
                                                                                'front', mod, mod_y)
        if show_disp_top:
            xixvals_mask_top_t, xiyvals_mask_top_t, xizvals_mask_top_t = mpf.mask_points(xixvals_t, xiyvals_t, xizvals_t,
                                                                                'top', mod, mod_y)
        if show_disp_front:
            xixvals_mask_front_t, xiyvals_mask_front_t, xizvals_mask_front_t = mpf.mask_points(xixvals_t, xiyvals_t, xizvals_t,
                                                                                'front', mod, mod_y)
        if show_vel_top or show_vel_top_pert:
            vxvals_mask_top_t, vyvals_mask_top_t, vzvals_mask_top_t = mpf.mask_points(vxvals_t, vyvals_t, vzvals_t,
                                                                                'top', mod, mod_y)
        if show_vel_front or show_vel_front_pert:
            vxvals_mask_front_t, vyvals_mask_front_t, vzvals_mask_front_t = mpf.mask_points(vxvals_t, vyvals_t, vzvals_t,
                                                                                'front', mod, mod_y)

        xgrid, zgrid, ygrid = np.mgrid[0:nx:(nx)*1j,
                                       0:nz:(nz)*1j,
                                       0:ny:(ny)*1j]

        field = mlab.pipeline.vector_field(bxvals_t, bzvals_t, byvals_t, name="B field",
                                           figure=fig, scalars=zgrid)
        field.spacing = spacing

        if show_axes:
            mpf.axes_no_label(field)

        if show_mini_axis:
            mpf.mini_axes()

        if uniform_light:
            #uniform lighting, but if we turn shading of volumes off, we are ok without
            mpf.uniform_lighting(fig)

        #Black background
        mpf.background_colour(fig, (0., 0., 0.))

        scalefactor = 8. * nx / 100. # scale factor for direction field vectors

        # Set up visualisation modules
        if show_mag_vec:
            bdirfield_front = mlab.pipeline.vector_field(bxvals_mask_front_t, bzvals_mask_front_t,
                                                         byvals_mask_front_t, name="B field front",
                                                         figure=fig)
            bdirfield_front.spacing = spacing
            mpf.vector_cut_plane(bdirfield_front, 'front', nx, ny, nz,
                                 y_spacing, scale_factor=scalefactor)

        if show_vel_top or show_vel_top_pert:
            vdirfield_top = mlab.pipeline.vector_field(vxvals_mask_top_t, np.zeros_like(vxvals_mask_top_t),
                                                       vyvals_mask_top_t, name="V field top",
                                                       figure=fig)
            vdirfield_top.spacing = spacing
            mpf.vector_cut_plane(vdirfield_top, 'top', nx, ny, nz,
                                 y_spacing, scale_factor=scalefactor)

        if show_vel_front or show_vel_front_pert:
            vdirfield_front = mlab.pipeline.vector_field(vxvals_mask_front_t, vzvals_mask_front_t,
                                                         vyvals_mask_front_t, name="V field front",
                                                         figure=fig)
            vdirfield_front.spacing = spacing
            mpf.vector_cut_plane(vdirfield_front,'front', nx, ny, nz,
                                 y_spacing, scale_factor=scalefactor)

        if show_disp_top:
            xidirfield_top = mlab.pipeline.vector_field(xixvals_mask_top_t, np.zeros_like(xixvals_mask_top_t),
                                                        xiyvals_mask_top_t, name="Xi field top",
                                                        figure=fig)
            xidirfield_top.spacing = spacing
            mpf.vector_cut_plane(xidirfield_top, 'top', nx, ny, nz,
                                 y_spacing, scale_factor=scalefactor)

        if show_disp_front:
            xidirfield_front = mlab.pipeline.vector_field(xixvals_mask_front_t, xizvals_mask_front_t,
                                                         xiyvals_mask_front_t, name="Xi field front",
                                                         figure=fig)
            xidirfield_front.spacing = spacing
            mpf.vector_cut_plane(xidirfield_front, 'front', nx, ny, nz,
                                 y_spacing, scale_factor=scalefactor)

        # Loop through time
        for t_ind in range(nt):
            if t_ind == 0:
                bxvals_t = bxvals
                byvals_t = byvals
                bzvals_t = bzvals

                if show_disp_top or show_disp_front:
                    xixvals_t = xixvals
                    xiyvals_t = xiyvals
                    xizvals_t = xizvals

                if show_vel_top or show_vel_top_pert or show_vel_front or show_vel_front_pert:
                    vxvals_t = vxvals
                    vyvals_t = vyvals
                    vzvals_t = vzvals

                if show_boundary:
                    xix_boundary_r_vals_t = xix_boundary_r_vals
                    xix_boundary_l_vals_t = xix_boundary_l_vals

                if show_density or show_density_pert:
                    rho_vals_t = rho_vals

            else:

                bxvals = np.real(np.repeat(sf.bx(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2))
                byvals = np.real(np.repeat(sf.by(mode, xvals, zvals, t, K)[:, :, np.newaxis], ny, axis=2))
                bz_eq3d = np.repeat(sf.bz_eq(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2)
                bzvals = np.real(np.repeat(-sf.bz(mode, xvals, zvals, t, W, K, R1)[:, :, np.newaxis], ny, axis=2) +
                                 bz_eq3d)

                bxvals_t = bxvals
                byvals_t = byvals
                bzvals_t = bzvals

                # Update mag field data
                field.mlab_source.set(u=bxvals_t, v=bzvals_t, w=byvals_t)

                # Update mag field visualisation module
                if show_mag_vec:
                    bxvals_mask_front_t, byvals_mask_front_t, bzvals_mask_front_t = mpf.mask_points(bxvals_t, byvals_t, bzvals_t,
                                                                                'front', mod, mod_y)
                    bdirfield_front.mlab_source.set(u=bxvals_mask_front_t, v=bzvals_mask_front_t, w=byvals_mask_front_t)

                # Update displacement field data
                if show_disp_top or show_disp_front:
                    xixvals_split = np.split(xixvals, [nz - (nz / nt) * t_ind], axis=1)
                    xiyvals_split = np.split(xiyvals, [nz - (nz / nt) * t_ind], axis=1)
                    xizvals_split = np.split(xizvals, [nz - (nz / nt) * t_ind], axis=1)

                    xixvals_t = np.concatenate((xixvals_split[1], xixvals_split[0]), axis=1)
                    xiyvals_t = np.concatenate((xiyvals_split[1], xiyvals_split[0]), axis=1)
                    xizvals_t = np.concatenate((xizvals_split[1], xizvals_split[0]), axis=1)

                    # Update displacement field visualisation module
                    if show_disp_top:
                        xixvals_mask_top_t, xiyvals_mask_top_t, xizvals_mask_top_t = mpf.mask_points(xixvals_t, xiyvals_t, xizvals_t,
                                                                                'top', mod, mod_y)
                        xidirfield_top.mlab_source.set(u=xixvals_mask_top_t, v=np.zeros_like(xixvals_mask_top_t), w=xiyvals_mask_top_t)
                    if show_disp_front:
                        xixvals_mask_front_t, xiyvals_mask_front_t, xizvals_mask_front_t = mpf.mask_points(xixvals_t, xiyvals_t, xizvals_t,
                                                                                'front', mod, mod_y)
                        xidirfield_front.mlab_source.set(u=xixvals_mask_front_t, v=xizvals_mask_front_t, w=xiyvals_mask_front_t)

                # Update velocity field data
                if show_vel_top or show_vel_top_pert or show_vel_front or show_vel_front_pert:
                    vxvals_split = np.split(vxvals, [nz - (nz / nt) * t_ind], axis=1)
                    vyvals_split = np.split(vyvals, [nz - (nz / nt) * t_ind], axis=1)
                    vzvals_split = np.split(vzvals, [nz - (nz / nt) * t_ind], axis=1)

                    vxvals_t = np.concatenate((vxvals_split[1], vxvals_split[0]), axis=1)
                    vyvals_t = np.concatenate((vyvals_split[1], vyvals_split[0]), axis=1)
                    vzvals_t = np.concatenate((vzvals_split[1], vzvals_split[0]), axis=1)

                    # Update velocity field visualisation module
                    if show_vel_top or show_vel_top_pert:
                        vxvals_mask_top_t, vyvals_mask_top_t, vzvals_mask_top_t = mpf.mask_points(vxvals_t, vyvals_t, vzvals_t,
                                                                                'top', mod, mod_y)
                        vdirfield_top.mlab_source.set(u=vxvals_mask_top_t, v=np.zeros_like(vxvals_mask_top_t), w=vyvals_mask_top_t)
                    if show_vel_front or show_vel_front_pert:
                        vxvals_mask_front_t, vyvals_mask_front_t, vzvals_mask_front_t = mpf.mask_points(vxvals_t, vyvals_t, vzvals_t,
                                                                                'front', mod, mod_y)
                        vdirfield_front.mlab_source.set(u=vxvals_mask_front_t, v=vzvals_mask_front_t, w=vyvals_mask_front_t)

                # Update boundary displacement data
                if show_boundary:
                    xix_boundary_r_vals_split = np.split(xix_boundary_r_vals, [nz - (nz / nt) * t_ind], axis=0)
                    xix_boundary_l_vals_split = np.split(xix_boundary_l_vals, [nz - (nz / nt) * t_ind], axis=0)

                    xix_boundary_r_vals_t = np.concatenate((xix_boundary_r_vals_split[1], xix_boundary_r_vals_split[0]), axis=0)
                    xix_boundary_l_vals_t = np.concatenate((xix_boundary_l_vals_split[1], xix_boundary_l_vals_split[0]), axis=0)

                # Update density data
                if show_density or show_density_pert:
                    rho_vals_split = np.split(rho_vals, [nz - (nz / nt) * t_ind], axis=1)

                    rho_vals_t = np.concatenate((rho_vals_split[1], rho_vals_split[0]), axis=1)

                    rho.mlab_source.set(scalars=rho_vals_t)

            # Boundary data - Letting mayavi know where to plot the boundary
            if show_boundary:
                ext_min_r = ((nx) * (xix_boundary_r_vals_t.min() - xmin) / (xmax - xmin)) * x_spacing
                ext_max_r = ((nx) * (xix_boundary_r_vals_t.max() - xmin) / (xmax - xmin)) * x_spacing

                ext_min_l = ((nx) * (xix_boundary_l_vals_t.min() - xmin) / (xmax - xmin)) * x_spacing
                ext_max_l = ((nx) * (xix_boundary_l_vals_t.max() - xmin) / (xmax - xmin)) * x_spacing

            #Make field lines
            if show_mag:
                # move seed points up with phase speed. - Bit of a fudge.
                # Create an array of points for which we want mag field seeds
                nx_seed = 9
                ny_seed = 13
                start_x = 30. * nx / 100.
                end_x = nx+1 - start_x
                start_y = 1.
                if ny == 20: # so that the lines dont go right up to the edge of the box
                    end_y = ny - 1.
                elif ny == 100:
                    end_y = ny - 2.
                elif ny == 300:
                    end_y = ny - 6.
                else:
                    end_y = ny - 1
                seeds=[]
                dx_res = (end_x - start_x) / (nx_seed-1)
                dy_res = (end_y - start_y) / (ny_seed-1)
                for j in range(ny_seed):
                    for i in range(nx_seed):
                        x = start_x + (i * dx_res) * x_spacing
                        y = start_y + (j * dy_res) * y_spacing
                        z = 1. + (t_start + t_ind*(t_end - t_start)/nt)/zmax * nz
                        seeds.append((x,z,y))

                if 'alfven' in mode:
                    for i in range(nx_seed):
                        del seeds[0]
                        del seeds[-1]

                # Remove previous field lines - field lines cannot be updated, just the data that they are built from
                if t_ind != 0:
                    field_lines.remove() # field_lines is defined in first go through loop
                field_lines = SeedStreamline(seed_points=seeds)

                # Field line visualisation tinkering
                field_lines.stream_tracer.integration_direction='both'
                field_lines.streamline_type = 'tube'
                field_lines.stream_tracer.maximum_propagation = nz * 2
                field_lines.tube_filter.number_of_sides = 20
                field_lines.tube_filter.radius = 0.7 * max(nx, ny, nz) / 100.
                field_lines.tube_filter.capping = True
                field_lines.actor.property.opacity = 1.0

                field.add_child(field_lines)
                module_manager = field_lines.parent

                # Colormap of magnetic field strength plotted on the field lines
                if show_mag_scale:
                    module_manager.scalar_lut_manager.lut_mode = 'coolwarm'
                    module_manager.scalar_lut_manager.data_range=[7,18]
                else:
                    mag_lut = module_manager.scalar_lut_manager.lut.table.to_array()
                    mag_lut[:,0] = [220]*256
                    mag_lut[:,1] = [20]*256
                    mag_lut[:,2] = [20]*256
                    module_manager.scalar_lut_manager.lut.table = mag_lut
                if show_mag_fade:
                    mpf.colormap_fade(module_manager, fade_value=20)

            # Which views do you want to show? Options are defined at the start
            views_selected = [0]#[0,1,4,5,6] #range(7) #[2,3]
            for view_ind, view_selected in enumerate(views_selected):
                view = view_options[view_selected]

                # Display boundary - cannot be updated each time
                if show_boundary:
                    # Boundaries should look different depending on view
                    if view == 'front-parallel':
                        #remove previous boundaries
                        if t != 0 or view_ind != 0:
                            boundary_r.remove()
                            boundary_l.remove()

                        # Make a fading colormap by changing opacity at ends
                        lut = np.reshape(np.array([150, 150, 150, 255]*256), (256,4))
                        fade_value = 125
                        lut[:fade_value,-1] = np.linspace(0, 255, fade_value)
                        lut[-fade_value:,-1] = np.linspace(255, 0, fade_value)

                        # Set up boundary visualisation
                        boundary_r = mlab.mesh(xix_boundary_r_vals_t, zgrid_zy, ygrid_zy,
                                                     extent=[ext_min_r, ext_max_r, 1, nz, 0, (ny-1) * y_spacing],
                                                     opacity=1., representation='wireframe',
                                                     line_width=12., scalars=zgrid_zy)
                        boundary_l = mlab.mesh(xix_boundary_l_vals_t, zgrid_zy, ygrid_zy,
                                                     extent=[ext_min_l, ext_max_l, 1, nz, 0, (ny-1) * y_spacing],
                                                     opacity=1., representation='wireframe',
                                                     line_width=12., scalars=zgrid_zy)

                        # Boundary color and other options
                        boundary_r.module_manager.scalar_lut_manager.lut.table = lut
                        boundary_l.module_manager.scalar_lut_manager.lut.table = lut
                        boundary_r.actor.property.lighting = False
                        boundary_r.actor.property.shading = False
                        boundary_l.actor.property.lighting = False
                        boundary_l.actor.property.shading = False

                    else:
                        #remove previous boundaries
                        if t != 0 or view_ind != 0:
                            boundary_r.remove()
                            boundary_l.remove()

                        # Make a fading colormap by changing opacity at ends
                        lut = np.reshape(np.array([150, 150, 150, 255]*256), (256,4))
                        fade_value = 20
                        lut[:fade_value,-1] = np.linspace(0, 255, fade_value)
                        lut[-fade_value:,-1] = np.linspace(255, 0, fade_value)

                        # Set up boundary visualisation
                        boundary_r = mlab.mesh(xix_boundary_r_vals_t, zgrid_zy, ygrid_zy,
                                               extent=[ext_min_r, ext_max_r, 1, nz, 0, (ny-1) * y_spacing],
                                               opacity=0.7, scalars=zgrid_zy)

                        boundary_l = mlab.mesh(xix_boundary_l_vals_t, zgrid_zy, ygrid_zy,
                                               extent=[ext_min_l, ext_max_l, 1, nz, 0, (ny-1) * y_spacing],
                                               opacity=0.7, scalars=zgrid_zy)

                        # Boundary color and other options
                        boundary_r.module_manager.scalar_lut_manager.lut.table = lut
                        boundary_l.module_manager.scalar_lut_manager.lut.table = lut
                        boundary_r.actor.property.lighting = False
                        boundary_r.actor.property.shading = False
                        boundary_l.actor.property.lighting = False
                        boundary_l.actor.property.shading = False

                # Set viewing angle - For some unknown reason we must redefine the camera position each time.
                # This is something to do with the boundaries being replaced each time.
                mpf.view_position(fig, view, nx, ny, nz)

                if save_images:
                    prefix = 'R1_'+str(R1) + '_' + mode + '_' + vis_mod_string + view + '_'# + '_norho_'
                    mlab.savefig(os.path.join(save_directory, prefix + str(t_ind+1) + '.png'))
                if t_ind == nt - 1:
                    if make_video:
                        i2v.image2video(filepath=save_directory, prefix=prefix,
                                        output_name=prefix+'video', out_extension='mp4',
                                        fps=fps, n_loops=4, delete_images=True,
                                        delete_old_videos=True, res=res[1])

            # Log: to keep us updated with progress
            if t_ind % 5 == 4:
                print('Finished frame number ' + str(t_ind + 1) + ' out of ' + str(number_of_frames))

            #Release some memory after each time step
            gc.collect()

            #step t forward
            t = t + (t_end - t_start) / nt

        # Close Mayavi window each time if we cant to make a video
        if make_video:
            mlab.close(fig)
        print('Finished ' + mode)


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.optim.lr_scheduler as lr_scheduler
from torch.autograd.function import Function
import torchvision

import os
import matplotlib.pyplot as plt
import argparse
from tqdm import trange
import numpy as np
from sklearn.metrics import classification_report

from losses import CenterLoss
from mnist_net import Net
import mnist_loader

# cf. https://cpp-learning.com/center-loss/

def main():
	args = parse_args()

	# Dataset
	train_loader, test_loader, classes = mnist_loader.load_dataset(args.dataset_dir, img_show=True)

	# Device
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	# Model
	model = Net().to(device)
	print(model)

	# Loss
	nllloss = nn.NLLLoss().to(device)  # CrossEntropyLoss = log_softmax + NLLLoss
	loss_weight = 1
	centerloss = CenterLoss(10, 2).to(device)

	# Optimizer
	dnn_optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)
	sheduler = lr_scheduler.StepLR(dnn_optimizer, 20, gamma=0.8)
	center_optimizer = optim.SGD(centerloss.parameters(), lr =0.5)

	print('Start training...')
	for epoch in range(100):
		# Update parameters.
		epoch += 1
		sheduler.step()

		# Train and test a model.
		train_acc, train_loss, feat, labels = train(device, train_loader, model, nllloss, loss_weight, centerloss, dnn_optimizer, center_optimizer)
		test_acc, test_loss = test(device, test_loader, model, nllloss, loss_weight, centerloss)
		stdout_temp = 'Epoch: {:>3}, train acc: {:<8}, train loss: {:<8}, test acc: {:<8}, test loss: {:<8}'
		print(stdout_temp.format(epoch, train_acc, train_loss, test_acc, test_loss))

		# Visualize features of each class.
		vis_img_path = args.vis_img_path_temp.format(str(epoch).zfill(3))
		visualize(feat.data.cpu().numpy(), labels.data.cpu().numpy(), epoch, vis_img_path)

		# Save a trained model.
		model_path = args.model_path_temp.format(str(epoch).zfill(3))
		torch.save(model.state_dict(), model_path)

def train(device, train_loader, model, nllloss, loss_weight, centerloss, dnn_optimizer, center_optimizer):
	running_loss = 0.0
	pred_list = []
	label_list = []
	ip1_loader = []
	idx_loader = []

	model.train()
	for i,(imgs, labels) in enumerate(train_loader):
		# Set batch data.
		imgs, labels = imgs.to(device), labels.to(device)
		# Predict labels.
		ip1, pred = model(imgs)
		# Calculate loss.
		loss = nllloss(pred, labels) + loss_weight * centerloss(labels, ip1)
		# Initilize gradient.
		dnn_optimizer.zero_grad()
		center_optimizer.zero_grad()
		# Calculate gradient.
		loss.backward()
		# Update parameters.
		dnn_optimizer.step()
		center_optimizer.step()
		# For calculation.
		running_loss += loss.item()
		pred_list += [int(p.argmax()) for p in pred]
		label_list += [int(l) for l in labels]
		# For visualization.
		ip1_loader.append(ip1)
		idx_loader.append((labels))

	# Calculate training accurary and loss.
	result = classification_report(pred_list, label_list, output_dict=True)
	train_acc = round(result['weighted avg']['f1-score'], 6)
	train_loss = round(running_loss / len(train_loader.dataset), 6)

	# Concatinate features and labels.
	feat = torch.cat(ip1_loader, 0)
	labels = torch.cat(idx_loader, 0)

	return train_acc, train_loss, feat, labels

def test(device, test_loader, model, nllloss, loss_weight, centerloss):
	model = model.eval()

	# Prediciton
	running_loss = 0.0
	pred_list = []
	label_list = []
	for i,(imgs, labels) in enumerate(test_loader):
		with torch.no_grad():
			# Set batch data.
			imgs, labels = imgs.to(device), labels.to(device)
			# Predict labels.
			ip1, pred = model(imgs)
			# Calculate loss.
			loss = nllloss(pred, labels) + loss_weight * centerloss(labels, ip1)
			# Append predictions and labels.
			running_loss += loss.item()
			pred_list += [int(p.argmax()) for p in pred]
			label_list += [int(l) for l in labels]

	# Calculate accuracy.
	result = classification_report(pred_list, label_list, output_dict=True)
	test_acc = round(result['weighted avg']['f1-score'], 6)
	test_loss = round(running_loss / len(test_loader.dataset), 6)

	return test_acc, test_loss

def visualize(feat, labels, epoch, vis_img_path):
	colors = ['#ff0000', '#ffff00', '#00ff00', '#00ffff', '#0000ff',
			  '#ff00ff', '#990000', '#999900', '#009900', '#009999']
	plt.figure()
	for i in range(10):
		plt.plot(feat[labels==i, 0], feat[labels==i, 1], '.', color=colors[i])
	plt.legend(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'], loc='best')
	plt.xlim(left=-8, right=8)
	plt.ylim(bottom=-8, top=8)
	plt.text(-7.8, 7.3, "epoch=%d" % epoch)
	plt.savefig(vis_img_path)
	plt.clf()

def parse_args():
	arg_parser = argparse.ArgumentParser(description="parser for focus one")

	arg_parser.add_argument("--dataset_dir", type=str, default='../inputs/')
	arg_parser.add_argument("--model_dir", type=str, default='../outputs/models/checkpoints/')
	arg_parser.add_argument("--model_path_temp", type=str, default='../outputs/models/checkpoints/mnist_original_softmax_center_epoch_{}.pth')
	arg_parser.add_argument("--vis_img_dir", type=str, default='../outputs/visual/')
	arg_parser.add_argument("--vis_img_path_temp", type=str, default='../outputs/visual/epoch_{}.png')

	args = arg_parser.parse_args()

	os.makedirs(args.dataset_dir, exist_ok=True)
	os.makedirs(args.model_dir, exist_ok=True)
	os.makedirs(args.vis_img_dir, exist_ok=True)

	return args

if __name__ == "__main__":
	main()


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

import os
import matplotlib.pyplot as plt

from framework.mongo import database
from website import settings

from utils import plot_dates, mkdirp

user_collection = database['user']

FIG_PATH = os.path.join(settings.ANALYTICS_PATH, 'figs', 'features')
mkdirp(FIG_PATH)

def analyze_email_invites():
    invited = user_collection.find({'unclaimed_records': {'$ne': {}}})
    dates_invited = [
        user['date_registered']
        for user in invited
    ]
    if not dates_invited:
        return
    fig = plot_dates(dates_invited)
    plt.title('email invitations ({}) total)'.format(len(dates_invited)))
    plt.savefig(os.path.join(FIG_PATH, 'email-invites.png'))
    plt.close()

def analyze_email_confirmations():
    confirmed = user_collection.find({
        'unclaimed_records': {'$ne': {}},
        'is_claimed': True,
    })
    dates_confirmed = [
        user['date_confirmed']
        for user in confirmed
    ]
    if not dates_confirmed:
        return
    fig = plot_dates(dates_confirmed)
    plt.title('confirmed email invitations ({}) total)'.format(len(dates_confirmed)))
    plt.savefig(os.path.join(FIG_PATH, 'email-invite-confirmations.png'))
    plt.close()

def main():
    analyze_email_invites()
    analyze_email_confirmations()

if __name__ == '__main__':
    main()


import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import PYTHON_PATH
if PYTHON_PATH is not None:
	for i in PYTHON_PATH: sys.path.insert(0,i)
import argparse
import h5py
import pandas as pd
import numpy as np
from hdgwas.tools import Timer
import tables
import glob

def probes_VCF2hdf5(data_path, save_path,study_name, chunk_size=1000000):

	if os.path.isfile(os.path.join(save_path,'probes',study_name+'.h5')):
		os.remove(os.path.join(save_path,'probes',study_name+'.h5'))

	hash_table={'keys':np.array([],dtype=np.int),'allele':np.array([])}

	df=pd.read_csv(data_path,sep='\t',chunksize=chunk_size, header=None,index_col=None)
	for i,chunk in enumerate(df):
		print 'add chunk {}'.format(i)
		print chunk.head()
		chunk.columns=[ "CHR","bp" ,"ID",'allele1','allele2','QUAL','FILTER','INFO'] #TODO (high) parse INFO
		hash_1=chunk.allele1.apply(hash)
		hash_2=chunk.allele2.apply(hash)
		k,indices=np.unique(np.append(hash_1,hash_2),return_index=True)
		s=np.append(chunk.allele1,chunk.allele2)[indices]
		ind=np.invert(np.in1d(k,hash_table['keys']))
		hash_table['keys']=np.append(hash_table['keys'],k[ind])
		hash_table['allele']=np.append(hash_table['allele'],s[ind])
		chunk.allele1=hash_1
		chunk.allele2=hash_2
		chunk.to_hdf(os.path.join(save_path,'probes',study_name+'.h5'),data_columns=["CHR","bp" ,"ID",'allele1','allele2'], key='probes',format='table',append=True,
			 min_itemsize = 25, complib='zlib',complevel=9 )
	pd.DataFrame.from_dict(hash_table).to_csv(os.path.join(save_path,'probes',study_name+'_hash_table.csv.gz'),index=False,compression='gzip', sep='\t')

def ind_VCF2hdf5(data_path, save_path,study_name):

	if os.path.isfile(os.path.join(save_path,'individuals',study_name+'.h5')):
		os.remove(os.path.join(save_path,'individuals',study_name+'.h5'))
	n=[]
	f=open(data_path,'r')
	for i,j in enumerate(f):
		n.append((j[:-1]))
	f.close()
	n=np.array(n)
	chunk=pd.DataFrame.from_dict({"individual":n})
	chunk.to_hdf(os.path.join(save_path,'individuals',study_name+'.h5'), key='individuals',format='table',
				 min_itemsize = 25, complib='zlib',complevel=9 )

def genotype_VCF2hdf5(data_path,id, save_path, study_name):

	df=pd.read_csv(data_path, header=None, index_col=None,sep='\t', dtype=np.float16)
	data=df.as_matrix()
	print data.shape
	print 'Saving chunk...{}'.format(os.path.join(save_path,'genotype',str(id)+'_'+study_name+'.h5'))
	h5_gen_file = tables.open_file(
		os.path.join(save_path,'genotype',str(id)+'_'+study_name+'.h5'), 'w', title=study_name)

	atom = tables.Float16Atom()
	genotype = h5_gen_file.create_carray(h5_gen_file.root, 'genotype', atom,
										(data.shape),
										title='Genotype',
										filters=tables.Filters(complevel=9, complib='zlib'))
	genotype[:] = data
	h5_gen_file.close()
	os.remove(data_path)

if __name__=="__main__":

	parser = argparse.ArgumentParser(description='Script to convert VCF data')
	parser.add_argument("-study_name", required=True, type=str, help="Study specific name")
	parser.add_argument("-id", type=str, help="subject id")
	parser.add_argument("-data",required=True, type=str, help="path to file")
	parser.add_argument("-out",required=True, type=str, help="path to results save folder")
	parser.add_argument("-flag",required=True,type=str,choices=['individuals','probes','chunk'], help="path to file with SNPs info")

	args = parser.parse_args()

	print args
	try:
		print ('Creating directories...')
		os.mkdir(os.path.join(args.out,'genotype') )
		os.mkdir(os.path.join(args.out,'individuals') )
		os.mkdir(os.path.join(args.out,'probes') )
		os.mkdir(os.path.join(args.out,'tmp_files'))
	except:
		print('Directories "genotype","probes","individuals" are already exist in {}...'.format(args.out))

	if args.flag=='probes':
		probes_VCF2hdf5(args.data, args.out, args.study_name)
	elif args.flag=='individuals':
		ind_VCF2hdf5(args.data, args.out,args.study_name)
	elif args.flag=='chunk':
		genotype_VCF2hdf5(args.data,args.id, args.out,args.study_name)