from __future__ import division
import numpy as np
import scipy.sparse as sp
from itertools import product

from sklearn.externals.six.moves import xrange
from sklearn.externals.six import iteritems

from scipy.sparse import issparse
from scipy.sparse import csc_matrix
from scipy.sparse import csr_matrix
from scipy.sparse import coo_matrix
from scipy.sparse import dok_matrix
from scipy.sparse import lil_matrix

from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_false
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_raises_regex

from sklearn.utils.multiclass import unique_labels
from sklearn.utils.multiclass import is_multilabel
from sklearn.utils.multiclass import type_of_target
from sklearn.utils.multiclass import class_distribution

class NotAnArray(object):
    """An object that is convertable to an array. This is useful to
    simulate a Pandas timeseries."""

    def __init__(self, data):
        self.data = data

    def __array__(self):
        return self.data

EXAMPLES = {
    'multilabel-indicator': [
        # valid when the data is formated as sparse or dense, identified
        # by CSR format when the testing takes place
        csr_matrix(np.random.RandomState(42).randint(2, size=(10, 10))),
        csr_matrix(np.array([[0, 1], [1, 0]])),
        csr_matrix(np.array([[0, 1], [1, 0]], dtype=np.bool)),
        csr_matrix(np.array([[0, 1], [1, 0]], dtype=np.int8)),
        csr_matrix(np.array([[0, 1], [1, 0]], dtype=np.uint8)),
        csr_matrix(np.array([[0, 1], [1, 0]], dtype=np.float)),
        csr_matrix(np.array([[0, 1], [1, 0]], dtype=np.float32)),
        csr_matrix(np.array([[0, 0], [0, 0]])),
        csr_matrix(np.array([[0, 1]])),
        # Only valid when data is dense
        np.array([[-1, 1], [1, -1]]),
        np.array([[-3, 3], [3, -3]]),
        NotAnArray(np.array([[-3, 3], [3, -3]])),
    ],
    'multiclass': [
        [1, 0, 2, 2, 1, 4, 2, 4, 4, 4],
        np.array([1, 0, 2]),
        np.array([1, 0, 2], dtype=np.int8),
        np.array([1, 0, 2], dtype=np.uint8),
        np.array([1, 0, 2], dtype=np.float),
        np.array([1, 0, 2], dtype=np.float32),
        np.array([[1], [0], [2]]),
        NotAnArray(np.array([1, 0, 2])),
        [0, 1, 2],
        ['a', 'b', 'c'],
        np.array([u'a', u'b', u'c']),
        np.array([u'a', u'b', u'c'], dtype=object),
        np.array(['a', 'b', 'c'], dtype=object),
    ],
    'multiclass-multioutput': [
        np.array([[1, 0, 2, 2], [1, 4, 2, 4]]),
        np.array([[1, 0, 2, 2], [1, 4, 2, 4]], dtype=np.int8),
        np.array([[1, 0, 2, 2], [1, 4, 2, 4]], dtype=np.uint8),
        np.array([[1, 0, 2, 2], [1, 4, 2, 4]], dtype=np.float),
        np.array([[1, 0, 2, 2], [1, 4, 2, 4]], dtype=np.float32),
        np.array([['a', 'b'], ['c', 'd']]),
        np.array([[u'a', u'b'], [u'c', u'd']]),
        np.array([[u'a', u'b'], [u'c', u'd']], dtype=object),
        np.array([[1, 0, 2]]),
        NotAnArray(np.array([[1, 0, 2]])),
    ],
    'binary': [
        [0, 1],
        [1, 1],
        [],
        [0],
        np.array([0, 1, 1, 1, 0, 0, 0, 1, 1, 1]),
        np.array([0, 1, 1, 1, 0, 0, 0, 1, 1, 1], dtype=np.bool),
        np.array([0, 1, 1, 1, 0, 0, 0, 1, 1, 1], dtype=np.int8),
        np.array([0, 1, 1, 1, 0, 0, 0, 1, 1, 1], dtype=np.uint8),
        np.array([0, 1, 1, 1, 0, 0, 0, 1, 1, 1], dtype=np.float),
        np.array([0, 1, 1, 1, 0, 0, 0, 1, 1, 1], dtype=np.float32),
        np.array([[0], [1]]),
        NotAnArray(np.array([[0], [1]])),
        [1, -1],
        [3, 5],
        ['a'],
        ['a', 'b'],
        ['abc', 'def'],
        np.array(['abc', 'def']),
        [u'a', u'b'],
        np.array(['abc', 'def'], dtype=object),
    ],
    'continuous': [
        [1e-5],
        [0, .5],
        np.array([[0], [.5]]),
        np.array([[0], [.5]], dtype=np.float32),
    ],
    'continuous-multioutput': [
        np.array([[0, .5], [.5, 0]]),
        np.array([[0, .5], [.5, 0]], dtype=np.float32),
        np.array([[0, .5]]),
    ],
    'unknown': [
        [[]],
        [()],
        # sequence of sequences that were'nt supported even before deprecation
        np.array([np.array([]), np.array([1, 2, 3])], dtype=object),
        [np.array([]), np.array([1, 2, 3])],
        [set([1, 2, 3]), set([1, 2])],
        [frozenset([1, 2, 3]), frozenset([1, 2])],

        # and also confusable as sequences of sequences
        [{0: 'a', 1: 'b'}, {0: 'a'}],

        # empty second dimension
        np.array([[], []]),

        # 3d
        np.array([[[0, 1], [2, 3]], [[4, 5], [6, 7]]]),
    ]
}

NON_ARRAY_LIKE_EXAMPLES = [
    set([1, 2, 3]),
    {0: 'a', 1: 'b'},
    {0: [5], 1: [5]},
    'abc',
    frozenset([1, 2, 3]),
    None,
]

MULTILABEL_SEQUENCES = [
    [[1], [2], [0, 1]],
    [(), (2), (0, 1)],
    np.array([[], [1, 2]], dtype='object'),
    NotAnArray(np.array([[], [1, 2]], dtype='object'))
]

def test_unique_labels():
    # Empty iterable
    assert_raises(ValueError, unique_labels)

    # Multiclass problem
    assert_array_equal(unique_labels(xrange(10)), np.arange(10))
    assert_array_equal(unique_labels(np.arange(10)), np.arange(10))
    assert_array_equal(unique_labels([4, 0, 2]), np.array([0, 2, 4]))

    # Multilabel indicator
    assert_array_equal(unique_labels(np.array([[0, 0, 1],
                                               [1, 0, 1],
                                               [0, 0, 0]])),
                       np.arange(3))

    assert_array_equal(unique_labels(np.array([[0, 0, 1],
                                               [0, 0, 0]])),
                       np.arange(3))

    # Several arrays passed
    assert_array_equal(unique_labels([4, 0, 2], xrange(5)),
                       np.arange(5))
    assert_array_equal(unique_labels((0, 1, 2), (0,), (2, 1)),
                       np.arange(3))

    # Border line case with binary indicator matrix
    assert_raises(ValueError, unique_labels, [4, 0, 2], np.ones((5, 5)))
    assert_raises(ValueError, unique_labels, np.ones((5, 4)), np.ones((5, 5)))
    assert_array_equal(unique_labels(np.ones((4, 5)), np.ones((5, 5))),
                       np.arange(5))

def test_unique_labels_non_specific():
    # Test unique_labels with a variety of collected examples

    # Smoke test for all supported format
    for format in ["binary", "multiclass", "multilabel-indicator"]:
        for y in EXAMPLES[format]:
            unique_labels(y)

    # We don't support those format at the moment
    for example in NON_ARRAY_LIKE_EXAMPLES:
        assert_raises(ValueError, unique_labels, example)

    for y_type in ["unknown", "continuous", 'continuous-multioutput',
                   'multiclass-multioutput']:
        for example in EXAMPLES[y_type]:
            assert_raises(ValueError, unique_labels, example)

def test_unique_labels_mixed_types():
    # Mix with binary or multiclass and multilabel
    mix_clf_format = product(EXAMPLES["multilabel-indicator"],
                             EXAMPLES["multiclass"] +
                             EXAMPLES["binary"])

    for y_multilabel, y_multiclass in mix_clf_format:
        assert_raises(ValueError, unique_labels, y_multiclass, y_multilabel)
        assert_raises(ValueError, unique_labels, y_multilabel, y_multiclass)

    assert_raises(ValueError, unique_labels, [[1, 2]], [["a", "d"]])
    assert_raises(ValueError, unique_labels, ["1", 2])
    assert_raises(ValueError, unique_labels, [["1", 2], [1, 3]])
    assert_raises(ValueError, unique_labels, [["1", "2"], [2, 3]])

def test_is_multilabel():
    for group, group_examples in iteritems(EXAMPLES):
        if group in ['multilabel-indicator']:
            dense_assert_, dense_exp = assert_true, 'True'
        else:
            dense_assert_, dense_exp = assert_false, 'False'

        for example in group_examples:
            # Only mark explicitly defined sparse examples as valid sparse
            # multilabel-indicators
            if group == 'multilabel-indicator' and issparse(example):
                sparse_assert_, sparse_exp = assert_true, 'True'
            else:
                sparse_assert_, sparse_exp = assert_false, 'False'

            if (issparse(example) or
                (hasattr(example, '__array__') and
                 np.asarray(example).ndim == 2 and
                 np.asarray(example).dtype.kind in 'biuf' and
                 np.asarray(example).shape[1] > 0)):
                examples_sparse = [sparse_matrix(example)
                                   for sparse_matrix in [coo_matrix,
                                                         csc_matrix,
                                                         csr_matrix,
                                                         dok_matrix,
                                                         lil_matrix]]
                for exmpl_sparse in examples_sparse:
                    sparse_assert_(is_multilabel(exmpl_sparse),
                                   msg=('is_multilabel(%r)'
                                   ' should be %s')
                                   % (exmpl_sparse, sparse_exp))

            # Densify sparse examples before testing
            if issparse(example):
                example = example.toarray()

            dense_assert_(is_multilabel(example),
                          msg='is_multilabel(%r) should be %s'
                          % (example, dense_exp))

def test_type_of_target():
    for group, group_examples in iteritems(EXAMPLES):
        for example in group_examples:
            assert_equal(type_of_target(example), group,
                         msg=('type_of_target(%r) should be %r, got %r'
                              % (example, group, type_of_target(example))))

    for example in NON_ARRAY_LIKE_EXAMPLES:
        msg_regex = 'Expected array-like \(array or non-string sequence\).*'
        assert_raises_regex(ValueError, msg_regex, type_of_target, example)

    for example in MULTILABEL_SEQUENCES:
        msg = ('You appear to be using a legacy multi-label data '
               'representation. Sequence of sequences are no longer supported;'
               ' use a binary array or sparse matrix instead.')
        assert_raises_regex(ValueError, msg, type_of_target, example)

def test_class_distribution():
    y = np.array([[1, 0, 0, 1],
                  [2, 2, 0, 1],
                  [1, 3, 0, 1],
                  [4, 2, 0, 1],
                  [2, 0, 0, 1],
                  [1, 3, 0, 1]])
    # Define the sparse matrix with a mix of implicit and explicit zeros
    data = np.array([1, 2, 1, 4, 2, 1, 0, 2, 3, 2, 3, 1, 1, 1, 1, 1, 1])
    indices = np.array([0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 5, 0, 1, 2, 3, 4, 5])
    indptr = np.array([0, 6, 11, 11, 17])
    y_sp = sp.csc_matrix((data, indices, indptr), shape=(6, 4))

    classes, n_classes, class_prior = class_distribution(y)
    classes_sp, n_classes_sp, class_prior_sp = class_distribution(y_sp)
    classes_expected = [[1, 2, 4],
                        [0, 2, 3],
                        [0],
                        [1]]
    n_classes_expected = [3, 3, 1, 1]
    class_prior_expected = [[3/6, 2/6, 1/6],
                            [1/3, 1/3, 1/3],
                            [1.0],
                            [1.0]]

    for k in range(y.shape[1]):
        assert_array_almost_equal(classes[k], classes_expected[k])
        assert_array_almost_equal(n_classes[k], n_classes_expected[k])
        assert_array_almost_equal(class_prior[k], class_prior_expected[k])

        assert_array_almost_equal(classes_sp[k], classes_expected[k])
        assert_array_almost_equal(n_classes_sp[k], n_classes_expected[k])
        assert_array_almost_equal(class_prior_sp[k], class_prior_expected[k])

    # Test again with explicit sample weights
    (classes,
     n_classes,
     class_prior) = class_distribution(y, [1.0, 2.0, 1.0, 2.0, 1.0, 2.0])
    (classes_sp,
     n_classes_sp,
     class_prior_sp) = class_distribution(y, [1.0, 2.0, 1.0, 2.0, 1.0, 2.0])
    class_prior_expected = [[4/9, 3/9, 2/9],
                            [2/9, 4/9, 3/9],
                            [1.0],
                            [1.0]]

    for k in range(y.shape[1]):
        assert_array_almost_equal(classes[k], classes_expected[k])
        assert_array_almost_equal(n_classes[k], n_classes_expected[k])
        assert_array_almost_equal(class_prior[k], class_prior_expected[k])

        assert_array_almost_equal(classes_sp[k], classes_expected[k])
        assert_array_almost_equal(n_classes_sp[k], n_classes_expected[k])
        assert_array_almost_equal(class_prior_sp[k], class_prior_expected[k])


import pickle
from io import BytesIO
import numpy as np
import scipy.sparse

from sklearn.datasets import load_digits, load_iris

from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_val_score

from sklearn.externals.six.moves import zip
from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_greater

from sklearn.naive_bayes import GaussianNB, BernoulliNB, MultinomialNB

# Data is just 6 separable points in the plane
X = np.array([[-2, -1], [-1, -1], [-1, -2], [1, 1], [1, 2], [2, 1]])
y = np.array([1, 1, 1, 2, 2, 2])

# A bit more random tests
rng = np.random.RandomState(0)
X1 = rng.normal(size=(10, 3))
y1 = (rng.normal(size=(10)) > 0).astype(np.int)

# Data is 6 random integer points in a 100 dimensional space classified to
# three classes.
X2 = rng.randint(5, size=(6, 100))
y2 = np.array([1, 1, 2, 2, 3, 3])

def test_gnb():
    # Gaussian Naive Bayes classification.
    # This checks that GaussianNB implements fit and predict and returns
    # correct values for a simple toy dataset.

    clf = GaussianNB()
    y_pred = clf.fit(X, y).predict(X)
    assert_array_equal(y_pred, y)

    y_pred_proba = clf.predict_proba(X)
    y_pred_log_proba = clf.predict_log_proba(X)
    assert_array_almost_equal(np.log(y_pred_proba), y_pred_log_proba, 8)

    # Test whether label mismatch between target y and classes raises
    # an Error
    # FIXME Remove this test once the more general partial_fit tests are merged
    assert_raises(ValueError, GaussianNB().partial_fit, X, y, classes=[0, 1])

def test_gnb_prior():
    # Test whether class priors are properly set.
    clf = GaussianNB().fit(X, y)
    assert_array_almost_equal(np.array([3, 3]) / 6.0,
                              clf.class_prior_, 8)
    clf.fit(X1, y1)
    # Check that the class priors sum to 1
    assert_array_almost_equal(clf.class_prior_.sum(), 1)

def test_gnb_sample_weight():
    """Test whether sample weights are properly used in GNB. """
    # Sample weights all being 1 should not change results
    sw = np.ones(6)
    clf = GaussianNB().fit(X, y)
    clf_sw = GaussianNB().fit(X, y, sw)

    assert_array_almost_equal(clf.theta_, clf_sw.theta_)
    assert_array_almost_equal(clf.sigma_, clf_sw.sigma_)

    # Fitting twice with half sample-weights should result
    # in same result as fitting once with full weights
    sw = rng.rand(y.shape[0])
    clf1 = GaussianNB().fit(X, y, sample_weight=sw)
    clf2 = GaussianNB().partial_fit(X, y, classes=[1, 2], sample_weight=sw / 2)
    clf2.partial_fit(X, y, sample_weight=sw / 2)

    assert_array_almost_equal(clf1.theta_, clf2.theta_)
    assert_array_almost_equal(clf1.sigma_, clf2.sigma_)

    # Check that duplicate entries and correspondingly increased sample
    # weights yield the same result
    ind = rng.randint(0, X.shape[0], 20)
    sample_weight = np.bincount(ind, minlength=X.shape[0])

    clf_dupl = GaussianNB().fit(X[ind], y[ind])
    clf_sw = GaussianNB().fit(X, y, sample_weight)

    assert_array_almost_equal(clf_dupl.theta_, clf_sw.theta_)
    assert_array_almost_equal(clf_dupl.sigma_, clf_sw.sigma_)

def test_discrete_prior():
    # Test whether class priors are properly set.
    for cls in [BernoulliNB, MultinomialNB]:
        clf = cls().fit(X2, y2)
        assert_array_almost_equal(np.log(np.array([2, 2, 2]) / 6.0),
                                  clf.class_log_prior_, 8)

def test_mnnb():
    # Test Multinomial Naive Bayes classification.
    # This checks that MultinomialNB implements fit and predict and returns
    # correct values for a simple toy dataset.

    for X in [X2, scipy.sparse.csr_matrix(X2)]:
        # Check the ability to predict the learning set.
        clf = MultinomialNB()
        assert_raises(ValueError, clf.fit, -X, y2)
        y_pred = clf.fit(X, y2).predict(X)

        assert_array_equal(y_pred, y2)

        # Verify that np.log(clf.predict_proba(X)) gives the same results as
        # clf.predict_log_proba(X)
        y_pred_proba = clf.predict_proba(X)
        y_pred_log_proba = clf.predict_log_proba(X)
        assert_array_almost_equal(np.log(y_pred_proba), y_pred_log_proba, 8)

        # Check that incremental fitting yields the same results
        clf2 = MultinomialNB()
        clf2.partial_fit(X[:2], y2[:2], classes=np.unique(y2))
        clf2.partial_fit(X[2:5], y2[2:5])
        clf2.partial_fit(X[5:], y2[5:])

        y_pred2 = clf2.predict(X)
        assert_array_equal(y_pred2, y2)

        y_pred_proba2 = clf2.predict_proba(X)
        y_pred_log_proba2 = clf2.predict_log_proba(X)
        assert_array_almost_equal(np.log(y_pred_proba2), y_pred_log_proba2, 8)
        assert_array_almost_equal(y_pred_proba2, y_pred_proba)
        assert_array_almost_equal(y_pred_log_proba2, y_pred_log_proba)

        # Partial fit on the whole data at once should be the same as fit too
        clf3 = MultinomialNB()
        clf3.partial_fit(X, y2, classes=np.unique(y2))

        y_pred3 = clf3.predict(X)
        assert_array_equal(y_pred3, y2)
        y_pred_proba3 = clf3.predict_proba(X)
        y_pred_log_proba3 = clf3.predict_log_proba(X)
        assert_array_almost_equal(np.log(y_pred_proba3), y_pred_log_proba3, 8)
        assert_array_almost_equal(y_pred_proba3, y_pred_proba)
        assert_array_almost_equal(y_pred_log_proba3, y_pred_log_proba)

def check_partial_fit(cls):
    clf1 = cls()
    clf1.fit([[0, 1], [1, 0]], [0, 1])

    clf2 = cls()
    clf2.partial_fit([[0, 1], [1, 0]], [0, 1], classes=[0, 1])
    assert_array_equal(clf1.class_count_, clf2.class_count_)
    assert_array_equal(clf1.feature_count_, clf2.feature_count_)

    clf3 = cls()
    clf3.partial_fit([[0, 1]], [0], classes=[0, 1])
    clf3.partial_fit([[1, 0]], [1])
    assert_array_equal(clf1.class_count_, clf3.class_count_)
    assert_array_equal(clf1.feature_count_, clf3.feature_count_)

def test_discretenb_partial_fit():
    for cls in [MultinomialNB, BernoulliNB]:
        yield check_partial_fit, cls

def test_gnb_partial_fit():
    clf = GaussianNB().fit(X, y)
    clf_pf = GaussianNB().partial_fit(X, y, np.unique(y))
    assert_array_almost_equal(clf.theta_, clf_pf.theta_)
    assert_array_almost_equal(clf.sigma_, clf_pf.sigma_)
    assert_array_almost_equal(clf.class_prior_, clf_pf.class_prior_)

    clf_pf2 = GaussianNB().partial_fit(X[0::2, :], y[0::2], np.unique(y))
    clf_pf2.partial_fit(X[1::2], y[1::2])
    assert_array_almost_equal(clf.theta_, clf_pf2.theta_)
    assert_array_almost_equal(clf.sigma_, clf_pf2.sigma_)
    assert_array_almost_equal(clf.class_prior_, clf_pf2.class_prior_)

def test_discretenb_pickle():
    # Test picklability of discrete naive Bayes classifiers

    for cls in [BernoulliNB, MultinomialNB, GaussianNB]:
        clf = cls().fit(X2, y2)
        y_pred = clf.predict(X2)

        store = BytesIO()
        pickle.dump(clf, store)
        clf = pickle.load(BytesIO(store.getvalue()))

        assert_array_equal(y_pred, clf.predict(X2))

        if cls is not GaussianNB:
            # TODO re-enable me when partial_fit is implemented for GaussianNB

            # Test pickling of estimator trained with partial_fit
            clf2 = cls().partial_fit(X2[:3], y2[:3], classes=np.unique(y2))
            clf2.partial_fit(X2[3:], y2[3:])
            store = BytesIO()
            pickle.dump(clf2, store)
            clf2 = pickle.load(BytesIO(store.getvalue()))
            assert_array_equal(y_pred, clf2.predict(X2))

def test_input_check_fit():
    # Test input checks for the fit method
    for cls in [BernoulliNB, MultinomialNB, GaussianNB]:
        # check shape consistency for number of samples at fit time
        assert_raises(ValueError, cls().fit, X2, y2[:-1])

        # check shape consistency for number of input features at predict time
        clf = cls().fit(X2, y2)
        assert_raises(ValueError, clf.predict, X2[:, :-1])

def test_input_check_partial_fit():
    for cls in [BernoulliNB, MultinomialNB]:
        # check shape consistency
        assert_raises(ValueError, cls().partial_fit, X2, y2[:-1],
                      classes=np.unique(y2))

        # classes is required for first call to partial fit
        assert_raises(ValueError, cls().partial_fit, X2, y2)

        # check consistency of consecutive classes values
        clf = cls()
        clf.partial_fit(X2, y2, classes=np.unique(y2))
        assert_raises(ValueError, clf.partial_fit, X2, y2,
                      classes=np.arange(42))

        # check consistency of input shape for partial_fit
        assert_raises(ValueError, clf.partial_fit, X2[:, :-1], y2)

        # check consistency of input shape for predict
        assert_raises(ValueError, clf.predict, X2[:, :-1])

def test_discretenb_predict_proba():
    # Test discrete NB classes' probability scores

    # The 100s below distinguish Bernoulli from multinomial.
    # FIXME: write a test to show this.
    X_bernoulli = [[1, 100, 0], [0, 1, 0], [0, 100, 1]]
    X_multinomial = [[0, 1], [1, 3], [4, 0]]

    # test binary case (1-d output)
    y = [0, 0, 2]   # 2 is regression test for binary case, 02e673
    for cls, X in zip([BernoulliNB, MultinomialNB],
                      [X_bernoulli, X_multinomial]):
        clf = cls().fit(X, y)
        assert_equal(clf.predict(X[-1:]), 2)
        assert_equal(clf.predict_proba([X[0]]).shape, (1, 2))
        assert_array_almost_equal(clf.predict_proba(X[:2]).sum(axis=1),
                                  np.array([1., 1.]), 6)

    # test multiclass case (2-d output, must sum to one)
    y = [0, 1, 2]
    for cls, X in zip([BernoulliNB, MultinomialNB],
                      [X_bernoulli, X_multinomial]):
        clf = cls().fit(X, y)
        assert_equal(clf.predict_proba(X[0:1]).shape, (1, 3))
        assert_equal(clf.predict_proba(X[:2]).shape, (2, 3))
        assert_almost_equal(np.sum(clf.predict_proba([X[1]])), 1)
        assert_almost_equal(np.sum(clf.predict_proba([X[-1]])), 1)
        assert_almost_equal(np.sum(np.exp(clf.class_log_prior_)), 1)
        assert_almost_equal(np.sum(np.exp(clf.intercept_)), 1)

def test_discretenb_uniform_prior():
    # Test whether discrete NB classes fit a uniform prior
    # when fit_prior=False and class_prior=None

    for cls in [BernoulliNB, MultinomialNB]:
        clf = cls()
        clf.set_params(fit_prior=False)
        clf.fit([[0], [0], [1]], [0, 0, 1])
        prior = np.exp(clf.class_log_prior_)
        assert_array_equal(prior, np.array([.5, .5]))

def test_discretenb_provide_prior():
    # Test whether discrete NB classes use provided prior

    for cls in [BernoulliNB, MultinomialNB]:
        clf = cls(class_prior=[0.5, 0.5])
        clf.fit([[0], [0], [1]], [0, 0, 1])
        prior = np.exp(clf.class_log_prior_)
        assert_array_equal(prior, np.array([.5, .5]))

        # Inconsistent number of classes with prior
        assert_raises(ValueError, clf.fit, [[0], [1], [2]], [0, 1, 2])
        assert_raises(ValueError, clf.partial_fit, [[0], [1]], [0, 1],
                      classes=[0, 1, 1])

def test_discretenb_provide_prior_with_partial_fit():
    # Test whether discrete NB classes use provided prior
    # when using partial_fit

    iris = load_iris()
    iris_data1, iris_data2, iris_target1, iris_target2 = train_test_split(
        iris.data, iris.target, test_size=0.4, random_state=415)

    for cls in [BernoulliNB, MultinomialNB]:
        for prior in [None, [0.3, 0.3, 0.4]]:
            clf_full = cls(class_prior=prior)
            clf_full.fit(iris.data, iris.target)
            clf_partial = cls(class_prior=prior)
            clf_partial.partial_fit(iris_data1, iris_target1,
                                    classes=[0, 1, 2])
            clf_partial.partial_fit(iris_data2, iris_target2)
            assert_array_almost_equal(clf_full.class_log_prior_,
                                      clf_partial.class_log_prior_)

def test_sample_weight_multiclass():
    for cls in [BernoulliNB, MultinomialNB]:
        # check shape consistency for number of samples at fit time
        yield check_sample_weight_multiclass, cls

def check_sample_weight_multiclass(cls):
    X = [
        [0, 0, 1],
        [0, 1, 1],
        [0, 1, 1],
        [1, 0, 0],
    ]
    y = [0, 0, 1, 2]
    sample_weight = np.array([1, 1, 2, 2], dtype=np.float64)
    sample_weight /= sample_weight.sum()
    clf = cls().fit(X, y, sample_weight=sample_weight)
    assert_array_equal(clf.predict(X), [0, 1, 1, 2])

    # Check sample weight using the partial_fit method
    clf = cls()
    clf.partial_fit(X[:2], y[:2], classes=[0, 1, 2],
                    sample_weight=sample_weight[:2])
    clf.partial_fit(X[2:3], y[2:3], sample_weight=sample_weight[2:3])
    clf.partial_fit(X[3:], y[3:], sample_weight=sample_weight[3:])
    assert_array_equal(clf.predict(X), [0, 1, 1, 2])

def test_sample_weight_mnb():
    clf = MultinomialNB()
    clf.fit([[1, 2], [1, 2], [1, 0]],
            [0, 0, 1],
            sample_weight=[1, 1, 4])
    assert_array_equal(clf.predict([[1, 0]]), [1])
    positive_prior = np.exp(clf.intercept_[0])
    assert_array_almost_equal([1 - positive_prior, positive_prior],
                              [1 / 3., 2 / 3.])

def test_coef_intercept_shape():
    # coef_ and intercept_ should have shapes as in other linear models.
    # Non-regression test for issue #2127.
    X = [[1, 0, 0], [1, 1, 1]]
    y = [1, 2]  # binary classification

    for clf in [MultinomialNB(), BernoulliNB()]:
        clf.fit(X, y)
        assert_equal(clf.coef_.shape, (1, 3))
        assert_equal(clf.intercept_.shape, (1,))

def test_check_accuracy_on_digits():
    # Non regression test to make sure that any further refactoring / optim
    # of the NB models do not harm the performance on a slightly non-linearly
    # separable dataset
    digits = load_digits()
    X, y = digits.data, digits.target
    binary_3v8 = np.logical_or(digits.target == 3, digits.target == 8)
    X_3v8, y_3v8 = X[binary_3v8], y[binary_3v8]

    # Multinomial NB
    scores = cross_val_score(MultinomialNB(alpha=10), X, y, cv=10)
    assert_greater(scores.mean(), 0.86)

    scores = cross_val_score(MultinomialNB(alpha=10), X_3v8, y_3v8, cv=10)
    assert_greater(scores.mean(), 0.94)

    # Bernoulli NB
    scores = cross_val_score(BernoulliNB(alpha=10), X > 4, y, cv=10)
    assert_greater(scores.mean(), 0.83)

    scores = cross_val_score(BernoulliNB(alpha=10), X_3v8 > 4, y_3v8, cv=10)
    assert_greater(scores.mean(), 0.92)

    # Gaussian NB
    scores = cross_val_score(GaussianNB(), X, y, cv=10)
    assert_greater(scores.mean(), 0.77)

    scores = cross_val_score(GaussianNB(), X_3v8, y_3v8, cv=10)
    assert_greater(scores.mean(), 0.86)

def test_feature_log_prob_bnb():
    # Test for issue #4268.
    # Tests that the feature log prob value computed by BernoulliNB when
    # alpha=1.0 is equal to the expression given in Manning, Raghavan,
    # and Schuetze's "Introduction to Information Retrieval" book:
    # http://nlp.stanford.edu/IR-book/html/htmledition/the-bernoulli-model-1.html

    X = np.array([[0, 0, 0], [1, 1, 0], [0, 1, 0], [1, 0, 1], [0, 1, 0]])
    Y = np.array([0, 0, 1, 2, 2])

    # Fit Bernoulli NB w/ alpha = 1.0
    clf = BernoulliNB(alpha=1.0)
    clf.fit(X, Y)

    # Manually form the (log) numerator and denominator that
    # constitute P(feature presence | class)
    num = np.log(clf.feature_count_ + 1.0)
    denom = np.tile(np.log(clf.class_count_ + 2.0), (X.shape[1], 1)).T

    # Check manual estimate matches
    assert_array_equal(clf.feature_log_prob_, (num - denom))

def test_bnb():
    # Tests that BernoulliNB when alpha=1.0 gives the same values as
    # those given for the toy example in Manning, Raghavan, and
    # Schuetze's "Introduction to Information Retrieval" book:
    # http://nlp.stanford.edu/IR-book/html/htmledition/the-bernoulli-model-1.html

    # Training data points are:
    # Chinese Beijing Chinese (class: China)
    # Chinese Chinese Shanghai (class: China)
    # Chinese Macao (class: China)
    # Tokyo Japan Chinese (class: Japan)

    # Features are Beijing, Chinese, Japan, Macao, Shanghai, and Tokyo
    X = np.array([[1, 1, 0, 0, 0, 0],
                  [0, 1, 0, 0, 1, 0],
                  [0, 1, 0, 1, 0, 0],
                  [0, 1, 1, 0, 0, 1]])

    # Classes are China (0), Japan (1)
    Y = np.array([0, 0, 0, 1])

    # Fit BernoulliBN w/ alpha = 1.0
    clf = BernoulliNB(alpha=1.0)
    clf.fit(X, Y)

    # Check the class prior is correct
    class_prior = np.array([0.75, 0.25])
    assert_array_almost_equal(np.exp(clf.class_log_prior_), class_prior)

    # Check the feature probabilities are correct
    feature_prob = np.array([[0.4, 0.8, 0.2, 0.4, 0.4, 0.2],
                             [1/3.0, 2/3.0, 2/3.0, 1/3.0, 1/3.0, 2/3.0]])
    assert_array_almost_equal(np.exp(clf.feature_log_prob_), feature_prob)

    # Testing data point is:
    # Chinese Chinese Chinese Tokyo Japan
    X_test = np.array([[0, 1, 1, 0, 0, 1]])

    # Check the predictive probabilities are correct
    unnorm_predict_proba = np.array([[0.005183999999999999,
                                      0.02194787379972565]])
    predict_proba = unnorm_predict_proba / np.sum(unnorm_predict_proba)
    assert_array_almost_equal(clf.predict_proba(X_test), predict_proba)

def test_naive_bayes_scale_invariance():
    # Scaling the data should not change the prediction results
    iris = load_iris()
    X, y = iris.data, iris.target
    labels = [GaussianNB().fit(f * X, y).predict(f * X)
              for f in [1E-10, 1, 1E10]]
    assert_array_equal(labels[0], labels[1])
    assert_array_equal(labels[1], labels[2])


from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six

from matplotlib.patches import _Style, FancyArrowPatch
from matplotlib.transforms import IdentityTransform
from matplotlib.path import Path
import numpy as np

class _FancyAxislineStyle:
    class SimpleArrow(FancyArrowPatch):
        """
        The artist class that will be returned for SimpleArrow style.
        """
        _ARROW_STYLE = "->"

        def __init__(self, axis_artist, line_path, transform,
                     line_mutation_scale):
            self._axis_artist = axis_artist
            self._line_transform = transform
            self._line_path = line_path
            self._line_mutation_scale = line_mutation_scale

            FancyArrowPatch.__init__(self,
                                     path=self._line_path,
                                     arrowstyle=self._ARROW_STYLE,
                                     arrow_transmuter=None,
                                     patchA=None,
                                     patchB=None,
                                     shrinkA=0.,
                                     shrinkB=0.,
                                     mutation_scale=line_mutation_scale,
                                     mutation_aspect=None,
                                     transform=IdentityTransform(),
                                     )

        def set_line_mutation_scale(self, scale):
            self.set_mutation_scale(scale*self._line_mutation_scale)

        def _extend_path(self, path, mutation_size=10):
            """
            Extend the path to make a room for drawing arrow.
            """
            from matplotlib.bezier import get_cos_sin

            x0, y0 = path.vertices[-2]
            x1, y1 = path.vertices[-1]
            cost, sint = get_cos_sin(x0, y0, x1, y1)

            d = mutation_size * 1.
            x2, y2 = x1 + cost*d, y1+sint*d

            if path.codes is None:
                _path = Path(np.concatenate([path.vertices, [[x2, y2]]]))
            else:
                _path = Path(np.concatenate([path.vertices, [[x2, y2]]]),
                             np.concatenate([path.codes, [Path.LINETO]]))

            return _path

        def set_path(self, path):
            self._line_path = path

        def draw(self, renderer):
            """
            Draw the axis line.
             1) transform the path to the display coordinate.
             2) extend the path to make a room for arrow
             3) update the path of the FancyArrowPatch.
             4) draw
            """
            path_in_disp = self._line_transform.transform_path(self._line_path)
            mutation_size = self.get_mutation_scale() #line_mutation_scale()
            extented_path = self._extend_path(path_in_disp,
                                              mutation_size=mutation_size)

            self._path_original = extented_path
            FancyArrowPatch.draw(self, renderer)

    class FilledArrow(SimpleArrow):
        """
        The artist class that will be returned for SimpleArrow style.
        """
        _ARROW_STYLE = "-|>"

class AxislineStyle(_Style):
    """
    :class:`AxislineStyle` is a container class which defines style classes
    for AxisArtists.

    An instance of any axisline style class is an callable object,
    whose call signature is ::

       __call__(self, axis_artist, path, transform)

    When called, this should return a mpl artist with following
    methods implemented. ::

      def set_path(self, path):
          # set the path for axisline.

      def set_line_mutation_scale(self, scale):
          # set the scale

      def draw(self, renderer):
          # draw

    """

    _style_list = {}

    class _Base(object):
        # The derived classes are required to be able to be initialized
        # w/o arguments, i.e., all its argument (except self) must have
        # the default values.

        def __init__(self):
            """
            initialization.
            """
            super(AxislineStyle._Base, self).__init__()

        def __call__(self, axis_artist, transform):
            """
            Given the AxisArtist instance, and transform for the path
            (set_path method), return the mpl artist for drawing the axis line.
            """

            return self.new_line(axis_artist, transform)

    class SimpleArrow(_Base):
        """
        A simple arrow.
        """

        ArrowAxisClass = _FancyAxislineStyle.SimpleArrow

        def __init__(self, size=1):
            """
             *size*
                size of the arrow as a fraction of the ticklabel size.
            """

            self.size = size
            super(AxislineStyle.SimpleArrow, self).__init__()

        def new_line(self, axis_artist, transform):

            linepath = Path([(0,0), (0, 1)])
            axisline = self.ArrowAxisClass(axis_artist, linepath, transform,
                                           line_mutation_scale=self.size)
            return axisline

    _style_list["->"] = SimpleArrow

    class FilledArrow(SimpleArrow):
        ArrowAxisClass = _FancyAxislineStyle.FilledArrow

    _style_list["-|>"] = FilledArrow


"""
===================
Isotonic Regression
===================

An illustration of the isotonic regression on generated data. The
isotonic regression finds a non-decreasing approximation of a function
while minimizing the mean squared error on the training data. The benefit
of such a model is that it does not assume any form for the target
function such as linearity. For comparison a linear regression is also
presented.

"""
print(__doc__)

# Author: Nelle Varoquaux <nelle.varoquaux@gmail.com>
#         Alexandre Gramfort <alexandre.gramfort@inria.fr>
# Licence: BSD

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

from sklearn.linear_model import LinearRegression
from sklearn.isotonic import IsotonicRegression
from sklearn.utils import check_random_state

n = 100
x = np.arange(n)
rs = check_random_state(0)
y = rs.randint(-50, 50, size=(n,)) + 50. * np.log(1 + np.arange(n))

###############################################################################
# Fit IsotonicRegression and LinearRegression models

ir = IsotonicRegression()

y_ = ir.fit_transform(x, y)

lr = LinearRegression()
lr.fit(x[:, np.newaxis], y)  # x needs to be 2d for LinearRegression

###############################################################################
# plot result

segments = [[[i, y[i]], [i, y_[i]]] for i in range(n)]
lc = LineCollection(segments, zorder=0)
lc.set_array(np.ones(len(y)))
lc.set_linewidths(0.5 * np.ones(n))

fig = plt.figure()
plt.plot(x, y, 'r.', markersize=12)
plt.plot(x, y_, 'g.-', markersize=12)
plt.plot(x, lr.predict(x[:, np.newaxis]), 'b-')
plt.gca().add_collection(lc)
plt.legend(('Data', 'Isotonic Fit', 'Linear Fit'), loc='lower right')
plt.title('Isotonic regression')
plt.show()


"""Module :mod:`sklearn.kernel_ridge` implements kernel ridge regression."""

# Authors: Mathieu Blondel <mathieu@mblondel.org>
#          Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
# License: BSD 3 clause

import numpy as np

from .base import BaseEstimator, RegressorMixin
from .metrics.pairwise import pairwise_kernels
from .linear_model.ridge import _solve_cholesky_kernel
from .utils import check_array, check_X_y
from .utils.validation import check_is_fitted

class KernelRidge(BaseEstimator, RegressorMixin):
    """Kernel ridge regression.

    Kernel ridge regression (KRR) combines ridge regression (linear least
    squares with l2-norm regularization) with the kernel trick. It thus
    learns a linear function in the space induced by the respective kernel and
    the data. For non-linear kernels, this corresponds to a non-linear
    function in the original space.

    The form of the model learned by KRR is identical to support vector
    regression (SVR). However, different loss functions are used: KRR uses
    squared error loss while support vector regression uses epsilon-insensitive
    loss, both combined with l2 regularization. In contrast to SVR, fitting a
    KRR model can be done in closed-form and is typically faster for
    medium-sized datasets. On the other  hand, the learned model is non-sparse
    and thus slower than SVR, which learns a sparse model for epsilon > 0, at
    prediction-time.

    This estimator has built-in support for multi-variate regression
    (i.e., when y is a 2d-array of shape [n_samples, n_targets]).

    Read more in the :ref:`User Guide <kernel_ridge>`.

    Parameters
    ----------
    alpha : {float, array-like}, shape = [n_targets]
        Small positive values of alpha improve the conditioning of the problem
        and reduce the variance of the estimates.  Alpha corresponds to
        ``(2*C)^-1`` in other linear models such as LogisticRegression or
        LinearSVC. If an array is passed, penalties are assumed to be specific
        to the targets. Hence they must correspond in number.

    kernel : string or callable, default="linear"
        Kernel mapping used internally. A callable should accept two arguments
        and the keyword arguments passed to this object as kernel_params, and
        should return a floating point number.

    gamma : float, default=None
        Gamma parameter for the RBF, laplacian, polynomial, exponential chi2
        and sigmoid kernels. Interpretation of the default value is left to
        the kernel; see the documentation for sklearn.metrics.pairwise.
        Ignored by other kernels.

    degree : float, default=3
        Degree of the polynomial kernel. Ignored by other kernels.

    coef0 : float, default=1
        Zero coefficient for polynomial and sigmoid kernels.
        Ignored by other kernels.

    kernel_params : mapping of string to any, optional
        Additional parameters (keyword arguments) for kernel function passed
        as callable object.

    Attributes
    ----------
    dual_coef_ : array, shape = [n_samples] or [n_samples, n_targets]
        Representation of weight vector(s) in kernel space

    X_fit_ : {array-like, sparse matrix}, shape = [n_samples, n_features]
        Training data, which is also required for prediction

    References
    ----------
    * Kevin P. Murphy
      "Machine Learning: A Probabilistic Perspective", The MIT Press
      chapter 14.4.3, pp. 492-493

    See also
    --------
    sklearn.linear_model.Ridge:
        Linear ridge regression.
    sklearn.svm.SVR:
        Support Vector Regression implemented using libsvm.

    Examples
    --------
    >>> from sklearn.kernel_ridge import KernelRidge
    >>> import numpy as np
    >>> n_samples, n_features = 10, 5
    >>> rng = np.random.RandomState(0)
    >>> y = rng.randn(n_samples)
    >>> X = rng.randn(n_samples, n_features)
    >>> clf = KernelRidge(alpha=1.0)
    >>> clf.fit(X, y) # doctest: +NORMALIZE_WHITESPACE
    KernelRidge(alpha=1.0, coef0=1, degree=3, gamma=None, kernel='linear',
                kernel_params=None)
    """
    def __init__(self, alpha=1, kernel="linear", gamma=None, degree=3, coef0=1,
                 kernel_params=None):
        self.alpha = alpha
        self.kernel = kernel
        self.gamma = gamma
        self.degree = degree
        self.coef0 = coef0
        self.kernel_params = kernel_params

    def _get_kernel(self, X, Y=None):
        if callable(self.kernel):
            params = self.kernel_params or {}
        else:
            params = {"gamma": self.gamma,
                      "degree": self.degree,
                      "coef0": self.coef0}
        return pairwise_kernels(X, Y, metric=self.kernel,
                                filter_params=True, **params)

    @property
    def _pairwise(self):
        return self.kernel == "precomputed"

    def fit(self, X, y=None, sample_weight=None):
        """Fit Kernel Ridge regression model

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training data

        y : array-like, shape = [n_samples] or [n_samples, n_targets]
            Target values

        sample_weight : float or array-like of shape [n_samples]
            Individual weights for each sample, ignored if None is passed.

        Returns
        -------
        self : returns an instance of self.
        """
        # Convert data
        X, y = check_X_y(X, y, accept_sparse=("csr", "csc"), multi_output=True,
                         y_numeric=True)
        if sample_weight is not None and not isinstance(sample_weight, float):
            sample_weight = check_array(sample_weight, ensure_2d=False)

        K = self._get_kernel(X)
        alpha = np.atleast_1d(self.alpha)

        ravel = False
        if len(y.shape) == 1:
            y = y.reshape(-1, 1)
            ravel = True

        copy = self.kernel == "precomputed"
        self.dual_coef_ = _solve_cholesky_kernel(K, y, alpha,
                                                 sample_weight,
                                                 copy)
        if ravel:
            self.dual_coef_ = self.dual_coef_.ravel()

        self.X_fit_ = X

        return self

    def predict(self, X):
        """Predict using the kernel ridge model

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Samples.

        Returns
        -------
        C : array, shape = [n_samples] or [n_samples, n_targets]
            Returns predicted values.
        """
        check_is_fitted(self, ["X_fit_", "dual_coef_"])
        K = self._get_kernel(X, self.X_fit_)
        return np.dot(K, self.dual_coef_)


# This code is supporting material for the book
# Building Machine Learning Systems with Python
# by Willi Richert and Luis Pedro Coelho
# published by PACKT Publishing
#
# It is made available under the MIT License

import numpy as np
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.datasets import load_svmlight_file
from sklearn.linear_model import LinearRegression
from sklearn.cross_validation import KFold

# Whether to use Elastic nets (otherwise, ordinary linear regression is used)

# Load data:
data, target = load_svmlight_file('data/E2006.train')

lr = LinearRegression()

# Compute error on training data to demonstrate that we can obtain near perfect
# scores:

lr.fit(data, target)
pred = lr.predict(data)

print('RMSE on training, {:.2}'.format(np.sqrt(mean_squared_error(target, pred))))
print('R2 on training, {:.2}'.format(r2_score(target, pred)))
print('')

pred = np.zeros_like(target)
kf = KFold(len(target), n_folds=5)
for train, test in kf:
    lr.fit(data[train], target[train])
    pred[test] = lr.predict(data[test])

print('RMSE on testing (5 fold), {:.2}'.format(np.sqrt(mean_squared_error(target, pred))))
print('R2 on testing (5 fold), {:.2}'.format(r2_score(target, pred)))


'''
Copyleft Oct 10, 2015 Arya Iranmehr, PhD Student, Bafna's Lab, UC San Diego,  Email: airanmehr@gmail.com
'''

from __future__ import division

import numpy as np;
import pandas as pd;
np.set_printoptions(linewidth=140, precision=5, suppress=True)
import subprocess, uuid, os,sys
import pylab as plt
import UTILS.Util as utl

stdout_old=sys.stdout;sys.stdout=open('/dev/null','w');import simuPOP as sim;sys.stdout=stdout_old # to avoid simuPop welcome message!

def sig(x): return 1./(1+np.exp(-x));
def logit(p): return (np.inf if p==1  else  np.log(p/(1.-p)))

a='';

def fff(msg):
    global a
    a += msg

class MSMS:
    @staticmethod
    def Simulate(n=200, mu=2*1e-9, L=50000, Ne=1e6,r=1e-9,verbose=False,seed=None,intPos=False):
        L=int(L)
        a= MSMS.Song(F=n, mu=mu, L=L, Ne=Ne, r=r,verbose=verbose,seed=seed)
        c=pd.Series(a.columns)
        if c.round().value_counts().max()==1:
            a.columns=c.round().astype(int)
        elif c.astype(int).value_counts().max()==1:
            a.columns = c.astype(int)
        if intPos:
            a.columns=map(int,np.sort(np.random.choice(L, a.shape[1], replace=False)))
        return a

    @staticmethod
    def Song(F=200, mu=2*1e-9, L=50000, Ne=1e6,r=4e-9, uid=None, theta=None, msmsFile=None, dir=None,verbose=False,seed=None):
        """
        Everything is exactly the sam
        """
#         print 'mu: {} r:{} NE:{} ,theta={} '.format(mu,r,Ne,4*Ne*mu*L), theta
        if msmsFile is not None:
            pop=MSMS.load(filename=msmsFile)[0]
        else:
            if theta:
                pop=MSMS.MSMS(n=F, numReps=1, theta=theta, rho=2*Ne*(L-1)*r, L=L, Ne=Ne, uid=uid, dir=dir,verbose=verbose,seed=seed)[0]
            else:
                pop=MSMS.MSMS(n=F, numReps=1, theta=2*Ne*mu*L, rho=2*Ne*(L-1)*r, L=L, Ne=Ne, uid=uid, dir=dir,verbose=verbose,seed=seed)[0]
        pop.r=r
        pop.Ne=Ne
        pop.L=L
        return pop

    @staticmethod
    def MSMS(n, numReps, theta, rho, L, Ne=None,uid=None,oneMutationEvery=None, dir=dir,verbose=False,seed=None):
        """
        Returns a list of dataframe for each replicate
        """
        if dir is None:
            dir= utl.PATH.simout;dir+= 'msms/';
        os.system('mkdir -p ' +dir)
        if oneMutationEvery is not None:
            nSS=L/oneMutationEvery
            theta=nSS/sum(1./np.arange(1,n))
        if uid is None:
            uid=str(uuid.uuid4())
        unique_filename = dir+uid+'.msms'
        if seed is None:
            seed=''
        else:
            seed=' -seed {} '.format(seed)
        cmd="java -jar -Xmx2g ~/bin/msms/lib/msms.jar -ms {} {} -t {:.0f} -r {:.0f} {:.0f} -oFP 0.000000000000E00 {} > {}".format(n, numReps, theta, rho, L, seed,unique_filename)
        if verbose:
            print cmd
        subprocess.call(cmd,shell=True)
        return MSMS.load(unique_filename)

    @staticmethod
    def getSeed(filename):
        file=open(filename);cmd=np.array(file.readline().strip().split(' '));seed=file.readline().strip()
        return seed

    @staticmethod
    def load(filename):
        n, R, L, posUnderSelection = MSMS.getParams(open(filename).readline())
        lines=np.array(map(str.strip,open(filename).readlines()) )
        posIdx= np.where(map(lambda x: x[:len('positions:')]=='positions:',lines))[0]
        try:
            theta = lines[np.where(map(lambda x: 'ThetaW Estimate Summaray:' in x, lines))[0][0]].split(':')[1].strip()
        except:
            theta = None
        POS=[map(lambda x: (float(x)*L), lines[ii].split()[1:]) for ii in posIdx]
        dfs=[pd.DataFrame(map(list ,lines[i +1 +range(n)]),columns=pos ) for i,pos in zip(posIdx,POS)]
        for df in dfs:
            df[df!='0']=1
            df[df=='0']=0
            df.L = L
            if posUnderSelection is not None:
                df.posUnderSelection = posUnderSelection * L
            if theta is not None:
                df.stat = pd.Series(theta.split(), index=['W', 'Pi', 'D']).astype(float)
        return  dfs

    @staticmethod
    def getParams(line):
        """
        Args:
            params: takes the first line of msmsm file

        Returns:
            n,R,L: number of individuals in the sample, the number of the replicates, genome length

        """
        params=np.array(line.strip().split(' '))
        offset=np.where(map(lambda x: 'ms'in x,  params))[0][0]
        if params[offset+1] == '-N':
            i=3
        else:
            i=1
        posUnderSelection = None
        if '-Sp' in params:  posUnderSelection = float(params[np.where(params == '-Sp')[0][0] + 1])
        return int(params[offset + i]), int(params[offset + i + 1]), int(
                params[np.where(params == '-r')[0][0] + 2]), posUnderSelection

    @staticmethod
    def fixDuplicatePositions(pos,L):
        pos=pd.Series(range(len(pos)),index=pos)
        posHits=pos.index.value_counts()
        invalidPOS=posHits[posHits>1]
        if not invalidPOS.shape[0]:
            return pos.index.values
        for invalidPos in invalidPOS.index:
            mini=pos.loc[invalidPos].min()
            maxi=pos.loc[invalidPos].max()
            lowerBound=pos[pos==mini-1].index.max()
            upperBound=pos[pos==maxi+1].index.min();
            if maxi==pos.shape[0]-1: upperBound=L
            if mini==0: lowerBound=0
            validRange=np.arange((upperBound-lowerBound)/2) # only second and third quartiles,
            offset=validRange+validRange.shape[0]/2 # first qunatulw
            newPos=pos.index.values;
            newPos[mini:maxi+1]=np.sort(np.random.choice(offset,pos.loc[invalidPos].shape[0],replace=False))+lowerBound
            pos.index=newPos
        assert pos.index.value_counts().max()==1
        return pos.index.values

    @staticmethod
    def Selection(msms, Ne, n, numReplicates, theta, rho, window_size, s, origin_count, posUnderSelection, gens, path):
        seed = ''
        for ii, gen in enumerate(gens):
            fname = path + '{}.msms'.format(int(gen))
            if (not ii) and s != 0:
                # while (nu0 < 0.95) or (nu0 > 0.99):
                cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SForceKeep -SFC   -oTW  >{}".format(
                        msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                             origin_count / Ne,
                        posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                os.system(cmd)
            else:
                cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC -SForceKeep  -oTW  >{}".format(
                        msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                             origin_count / Ne,
                        posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                os.system(cmd)
            if not ii: seed = MSMS.getSeed(fname)

    @staticmethod
    def SelectionFinale(msms, Ne, n, numReplicates, theta, rho, window_size, s, origin_count, posUnderSelection, gens,
                            path):
        seed = ''
        nu0 = 0
        for ii, gen in enumerate(gens):
            fname = path + '{}.msms'.format(int(gen))
            if (not ii) and s != 0:
                while (nu0 < 0.9):
                    cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SForceKeep -SFC   -oTW  >{}".format(
                            msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                                origin_count / Ne,
                            posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                    os.system(cmd)

                    nu0 = MSMS.load(fname)[0].mean(0).loc[25000]
            else:
                cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC -SForceKeep  -oTW  >{}".format(
                        msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                             origin_count / Ne,
                        posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                os.system(cmd)
            if not ii: seed = MSMS.getSeed(fname)

    @staticmethod
    def SelectionNu(msms, Ne, n, numReplicates, theta, rho, window_size, s, posUnderSelection, nu, path=None):
        seed = ''
        if path is None: path = '~/tmp.msms'
        fname = path + '{}.msms'.format(nu)
        cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SF 0 {}  -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC   -oTW  >{}".format(
                msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, nu, posUnderSelection,
                ('-seed {}'.format(seed), '')[seed is ''], fname)
        print cmd
        os.system(cmd)
        return MSMS.load(fname)

    @staticmethod
    def SelectionNuForward(msms, Ne, n, numReplicates, theta, rho, window_size, s, origin_count, posUnderSelection,
                               gens, path):
        nu0 = 0
        for ii, gen in enumerate(gens):
            fname = path + '{}.msms'.format(gen)
            if (not ii) and s != 0:
                while (nu0 < 0.95) or (nu0 > 0.99):
                    cmd = "{} -N {} -ms {} {} -t {} -r {} {:.0f} -SAA {} -SaA {} -SI {} 1 {} -Sp {} -oOC -Smark -oFP 0.000000000000E00 {}  -SFC   -oTW  >{}".format(
                            msms, Ne, n, numReplicates, theta, rho, window_size, 2 * Ne * s, Ne * s, gen / (4. * Ne),
                                                                                origin_count / Ne,
                            posUnderSelection, ('-seed {}'.format(seed), '')[seed is ''], fname)
                    os.system(cmd)
                    nu0 = MSMS.load(fname)[0].mean(0).loc[25000]
            print nu0, gen, cmd
            if not ii: seed = MSMS.getSeed(fname)

class Simulation:
    @staticmethod
    def setSeed(seed):
        if seed is None: return
        sim.setRNG('rand', seed + 1);
        np.random.seed(seed)

    @staticmethod
    def load(ExperimentName, s=0.1, L=50000, experimentID=0, nu0=0.005, isFolded=False, All=False, startGeneration=0,
             maxGeneration=50, numReplicates=3, numSamples=5, step=10, replicates=None, coverage=np.inf):
        path='{}{}/simpop/'.format(utl.PATH.simout, ExperimentName) + Simulation.getSimulationName(s=s, L=L, experimentID=experimentID, initialCarrierFreq=nu0, isFolded=isFolded) + '.pkl'
        sim= pd.read_pickle(path)
        sim.savedPath=path
        if replicates is not None:          sim.setReplicates(sorted(replicates))
        elif numReplicates is not None:     sim.setReplicates(range(numReplicates))

        if coverage != np.inf:
            sim.Xi = sim.X
            sim.X = sim.C.loc[coverage] / sim.D.loc[coverage].astype(float)
            sim.X = np.array(map(lambda x: utl.roundto(x, 5), sim.X.reshape(-1) * 1e4)).reshape(sim.X.shape) / 1e4
            sim.CD=sim.getCD(coverage)
            sim.CD.columns.names=['REP','GEN','READ']

        if not All: sim.setSamplingTimes(maxGeneration=min(maxGeneration,sim.getGenerationTimes()[-1]),numSamples=numSamples,step=step,startGeneration=startGeneration)
        return sim

    @staticmethod
    def getSimulationName(s,L,experimentID,initialCarrierFreq,isFolded,msms=False):
        if msms:
            return 'L{:.0f}K.{:04.0f}'.format(L/1000,experimentID)
        if s:
            return 'Nu{:E}.s{:E}.L{:.0f}K.{:04.0f}{}'.format(np.round(float(initialCarrierFreq), 3), s, L / 1000,
                                                             experimentID, ('', '.Folded')[isFolded])
        else:
            return 'Nu{:E}.s{:E}.L{:.0f}K.{:04.0f}{}'.format(0, s * 100, L / 1000, experimentID,
                                                             ('', '.Folded')[isFolded])

    def setReplicates(self,replicates):
        self.numReplicates=len(replicates)
        self.X=self.X[:,:,replicates]
        self.C = self.C.apply(lambda x: x[:, :, replicates])
        self.D = self.D.apply(lambda x: x[:, :, replicates])

    def __init__(self, outpath=utl.PATH.simout, N=1000, generationStep=10, maxGeneration=None,
                 s=0.05, r=4e-9, Ne=1e6, mu=2e-9, F=200, h=0.5, L=50000, startGeneration=0, numReplicates=3, H0=None,
                 foldInitialAFs=False, save=True, foutName=None,
                 doForwardSimulationNow=True, experimentID=-1,
                 msmsFile=None,initialCarrierFreq=0, ExperimentName=None, simulateNeutrallyFor=0,
                 initialNeutralGenerations=0, ignoreInitialNeutralGenerations=True,
                 makeSureSelectedSiteDontGetLost=True, onlyKeep=None, verbose=0, sampingTimes=None, minIncrease=0,
                 model=None,initDiploidPop=None,posUnderSelection=-1,haplotypes=False,seed=None,recombinator=None
                 ):
        """
        A General Simulation Class; with params
        H0: Dataframe F x m for F individuals and m segregation sites ;  Initial Haplotypes; dataframe with columns as positions
        """
        self.recombinator=recombinator
        if seed is not None:
            Simulation.setSeed(seed)
        self.s = s;
        self.r = r;
        self.Ne = Ne;
        self.mu = mu;
        self.F = F;
        self.h = h;
        self.L = int(L);
        self.startGeneration = startGeneration;
        self.numReplicates = numReplicates;
        self.posUnderSelection = -1
        self.initDiploidPop = initDiploidPop
        self.initialCarrierFreq= initialCarrierFreq if initialCarrierFreq else 1./self.F
        if foutName is not None:
            self.uid=foutName
            self.uidMSMS=None
        elif experimentID>=0:
            self.uid=Simulation.getSimulationName(self.s, self.L, self.experimentID, initialCarrierFreq=self.initialCarrierFreq, isFolded=self.foldInitialAFs)
            self.uidMSMS=Simulation.getSimulationName(self.s, self.L, self.experimentID, initialCarrierFreq=self.initialCarrierFreq, isFolded=self.foldInitialAFs,msms=True)
        else:
            self.uid=str(uuid.uuid4())
            self.uidMSMS=self.uid

        if H0 is None:
            self.simulateH0()
            H0=self.H0
        else:
            self.setH0(H0);

        if posUnderSelection >= 0:
            if self.positions is None:
                self.positions=map(int, self.initDiploidPop.lociPos())
            self.set_posUnderSelection(posUnderSelection)
        assert ExperimentName != None
        self.save=save
        self.model=model
        self.minIncrease = minIncrease
        self.samplingTimes=sampingTimes
        self.initialNeutralGenerations=initialNeutralGenerations
        self.onlyKeep=onlyKeep
        self.makeSureSelectedSiteDontGetLost=makeSureSelectedSiteDontGetLost
        self.ignoreInitialNeutralGenerations=ignoreInitialNeutralGenerations
        self.msmsFile=msmsFile;self.outpath=outpath; self.outpath=outpath ; self.N=N; self.generationStep=generationStep; self.maxGeneration= maxGeneration;
        self.foldInitialAFs=foldInitialAFs;self.doForwardSimulationNow=doForwardSimulationNow;self.experimentID=experimentID
        self.simulateNeutrallyFor=simulateNeutrallyFor

        self.setH0(H0);
        if not os.path.exists(self.outpath) : os.makedirs(self.outpath)
        self.outpath+=ExperimentName
        if not os.path.exists(self.outpath) : os.makedirs(self.outpath)
        self.outpathmsms=self.outpath+'/msms/';self.outpath+='/simpop/'
        if not os.path.exists(self.outpath) : os.makedirs(self.outpath)
        if not os.path.exists(self.outpathmsms) : os.makedirs(self.outpathmsms)
        if self.maxGeneration is None: self.maxGeneration=Simulation.getFixationTime(self.s, Ne=self.F, roundto10=True)
        self.theta=2*self.Ne*self.mu*self.L
        self.pops=[]

        if self.model is None:
            import simuPOP.demography as dmg
            self.model=dmg.LinearGrowthModel(T=self.maxGeneration, N0=self.N, NT=self.N)

        if self.doForwardSimulationNow:
            self.forwardSimulation()

    @staticmethod
    def simulateSingleLoci(nu0=0.005, T=100, s=0.1, N=1000,verbose=True,h=0.5,seed=None):
        if verbose:
            print '.',
        step = 1
        Simulation.setSeed(seed)
        pop = sim.Population(size=N, ploidy=2, loci=[1],infoFields=['fitness']);sim.initGenotype(pop, prop=[1-nu0,nu0]);simulator = sim.Simulator(pop.clone(), rep=1);
        # sim.stat(pop, alleleFreq=[0]);        print pop.dvars().alleleFreq[0][1]
        global a;a = "0;;{}\n".format(nu0)
        simulator.evolve(initOps=[sim.InitSex()],
                         preOps=sim.MapSelector(loci=0, fitness={(0, 0): 1, (0, 1): 1 + s *h, (1, 1): 1 + s}),
                         matingScheme=sim.RandomMating(), postOps=[sim.Stat(alleleFreq=[0], step=step),
                                                                   sim.PyEval("'%d;;' % (gen+1)", reps=0, step=step,
                                                                              output=fff), sim.PyEval(
                    r"'{}\n'.format(map(lambda x: round(x[1],5),alleleFreq.values())[0])", step=step, output=fff)],
                         gen=T)
        return pd.DataFrame(zip(*map(lambda x: x.split(';;'), a.strip().split('\n')))).T.set_index(0)[1].astype(float)

    def createInitialDiploidPopulation(self):
        """
        initHaps : np 2D array which m x nSS where m i number of individual haps and nSS is number of SS
        return a homozygote diploid population which every haplotype is copied n times
        """
        if self.initDiploidPop is not None: return self.initDiploidPop
        assert int(2*self.N/self.F)==2*self.N/float(self.F)  # N should be  a multiplier of F
        nSS=self.H0.shape[1];n=int(self.N/self.F)
        try:
            pop = sim.Population(size=self.N, ploidy=2, loci=nSS,lociPos=list(self.positions), infoFields='fitness')
        except:
            import traceback
            print(traceback.format_exc())
            print list(self.positions), nSS,n,self.H0.shape[0]
            exit()
        assert (self.N % self.H0.shape[0]) ==0
        H= [[list(h.values),list(h.values)] for _ in range(n) for _,h in self.H0.iterrows()]
        for (i,h) in zip(pop.individuals(),H): # for each indv assing first and second chromosome
            i.setGenotype(h[0],0 );i.setGenotype(h[1],1 ) #homozygote population of diploid
        # sim.stat(pop, alleleFreq=range(nSS));print np.array([pop.dvars().alleleFreq[x][1] for x in range(nSS)])
        return pop

    @staticmethod
    def getGT(pop, i=None, pos=None):
        if i == None and pos == None:
            df = pd.concat([pd.DataFrame([list(i.genotype(0)) for i in pop.individuals()]),
                            pd.DataFrame([list(i.genotype(1)) for i in pop.individuals()])],
                           keys=[0, 1]).sort_index().reorder_levels([1, 0]).sort_index()
            df.columns = map(int, pop.lociPos())
            return df
        i = np.where(np.array(pop.lociPos()).astype(int) == pos)[0][0]
        a, b = [], []
        for ind in pop.individuals():
            a += [ind.genotype(0)[i]]
            b += [ind.genotype(1)[i]]
        return pd.concat([pd.Series(a), pd.Series(b)], keys=[0, 1]).reorder_levels([1, 0]).sort_index()
    @staticmethod
    def createDiploidPopulationFromDataFrame(df):
        """
        initHaps : np 2D array which m x nSS where m i number of individual haps and nSS is number of SS
        return a homozygote diploid population which every haplotype is copied n times
        """
        pop = sim.Population(size=df.shape[0]/2, ploidy=2, loci=df.shape[1], lociPos=list(df.columns), infoFields='fitness')
        for j,i in enumerate(pop.individuals()): # for each indv assing first and second chromosome
            i.setGenotype(df.loc[j].loc[0].tolist(),0 );i.setGenotype(df.loc[j].loc[1].tolist(),1 )
        return pop

    @staticmethod
    def _simualtePop(pop, s=0, h=0.5, r=2e-8, siteUnderSelection=0,gen=1,recombinator=None,seed=None):
        "Gets population and returns population"
        Simulation.setSeed(seed)
        simulator = sim.Simulator(pop.clone(), rep=1)
        if recombinator is None:recombinator=sim.Recombinator(intensity=r)
        simulator.evolve(
            initOps=[sim.InitSex()],
            preOps=sim.MapSelector(loci=siteUnderSelection, fitness={(0, 0): 1, (0, 1): 1 + s * h, (1, 1): 1 + s}),
            matingScheme=sim.RandomMating(ops=recombinator),
            gen=gen)
        return simulator.population(0).clone()

    @staticmethod
    def _simualte(pop,s,h,r,siteUnderSelection,positions,startGeneration,generationStep,maxGeneration,model=None,makeSureSelectedSiteDontGetLost=True):
        "Gets population and returns Dataframe, Static method"
        N = int(pop.popSize())
        if model is None:
            import simuPOP.demography as dmg
            model = dmg.LinearGrowthModel(T=maxGeneration, N0=N, NT=N)
        simulator = sim.Simulator(pop.clone(), rep=1)
        global a;a = ""
        pops=[]
        step=1# this is slow but safe, dont change it
        simulator.evolve(
            initOps=[sim.InitSex()],
            preOps=sim.MapSelector(loci=siteUnderSelection, fitness={(0, 0): 1, (0, 1): 1 + s * h, (1, 1): 1 + s}),
            matingScheme=sim.RandomMating(ops=sim.Recombinator(intensity=r),subPopSize=model),
            postOps=[sim.Stat(alleleFreq=range(int(pop.numLoci()[0])), step=step), sim.PyEval("'Gen %4d;;' % (gen+1)", reps=0,step= step, output=fff), sim.PyEval(r"'{},'.format(map(lambda x: round(x[1],5),alleleFreq.values()))", step=step, output=fff),sim.PyOutput('\n', reps=-1, step=step, output=fff)],
            gen = maxGeneration)
        # idx=np.arange(self.generationStep-1,self.maxGeneration,self.generationStep)+self.initialNeutralGenerations
        print a
        _,data=zip(*map(lambda x: x.split(';;'),a.strip().split('\n')))
        data=np.array(map(eval,data))[:,0,:]
        print data
        # if data[-1, self.siteUnderSelection] >= self.initialCarrierFreq + self.minIncrease or self.s == 0 or not self.makeSureSelectedSiteDontGetLost:
        if data[-1, siteUnderSelection] or s == 0 or not makeSureSelectedSiteDontGetLost:
            try:
                pops+=[simulator.extract(0) ]
            except:
                print 'Error'
            return data[int(startGeneration/generationStep):,:]
        else:
            return Simulation._simualte()

    def simualte(self):
        "Gets population and returns Dataframe, Class method"
        import simuPOP.demography as dmg
        # model=dmg.ExponentialGrowthModel(T=50, N0=1000, NT=200)
        simulator = sim.Simulator(self.initDiploidPop.clone(), rep=1)
        # sim.dump(self.initDiploidPop)
        global a;a = ""
        if self.recombinator is None:
            self.recombinator=sim.Recombinator(intensity=self.r)
        step=1# this is slow but safe, dont change it
        simulator.evolve(
            initOps=[sim.InitSex()],
            preOps=sim.MapSelector(loci=self.siteUnderSelection, fitness={(0,0):1, (0,1):1+self.s*self.h, (1,1):1+self.s}),
            matingScheme=sim.RandomMating(ops=self.recombinator,subPopSize=self.model),

            postOps=[sim.Stat(alleleFreq=range(len(self.positions)), step=step),
                     sim.PyEval("'Gen %4d;;' % (gen+1)", reps=0,step= step, output=fff), sim.PyEval(r"'{},'.format(map(lambda x: round(x[1],5),alleleFreq.values()))", step=step, output=fff),sim.PyOutput('\n', reps=-1, step=step, output=fff)],
            gen = self.maxGeneration)
        # idx=np.arange(self.generationStep-1,self.maxGeneration,self.generationStep)+self.initialNeutralGenerations
        _,data=zip(*map(lambda x: x.split(';;'),a.strip().split('\n')))
        data=np.array(map(eval,data))[:,0,:]
        # if data[-1, self.siteUnderSelection] >= self.initialCarrierFreq + self.minIncrease or self.s == 0 or not self.makeSureSelectedSiteDontGetLost:
        if data[-1, self.siteUnderSelection] or self.s == 0 or not self.makeSureSelectedSiteDontGetLost:
            try:
                self.pops+=[simulator.extract(0) ]
            except:
                print 'Error'
            return data[int(self.startGeneration/self.generationStep):,:]
        else:
            # print pd.Series(data[:, self.siteUnderSelection])
            return self.simualte()

    def simulateH0(self):
        self.H0=MSMS.Song(F=self.F, L=self.L, Ne=self.Ne, r=self.r, mu=self.mu,uid=self.uidMSMS)

    def set_siteUnderSelection(self,x):
        self.siteUnderSelection=x
        self.posUnderSelection=self.positions[self.siteUnderSelection]

    def set_posUnderSelection(self,x):
        self.posUnderSelection=x
        self.siteUnderSelection=np.where(self.positions==self.posUnderSelection)[0][0]

    def setH0(self,H0):
        self.H0=H0
        self.positions=self.H0.columns.values
        self.F=self.H0.shape[0]

    def set_BeneficialLoci(self,selectionOnRandomSite=False,siteUnderSelection=None,posUnderSelection =None):
        if selectionOnRandomSite:
            self.set_siteUnderSelection(np.random.randint(0,self.H0.shape[1]))
        elif siteUnderSelection is not None:
            self.set_siteUnderSelection(siteUnderSelection)
        elif posUnderSelection is not None:
            self.set_siteUnderSelection(posUnderSelection)
        else:
            if not self.s:
                self.set_siteUnderSelection(self.X0.argmax())
            else:
                sites=np.sort(np.where(self.X0== self.initialCarrierFreq)[0]);
                if not len(sites):
                    sites=np.sort(np.where(( self.X0 <= self.initialCarrierFreq +0.025) & ( self.X0 >= self.initialCarrierFreq -0.025) ) [0]);
                    if not len(sites):
                        print 'Try again. No site at freq ',self.initialCarrierFreq, self.uid; return
                self.set_siteUnderSelection(sites[np.random.randint(0,len(sites))])

    def createInitHaps(self):
        assignPositions=True
        if self.H0 is None:
            H0 = MSMS.Song(F=self.F, L=self.L, Ne=self.Ne, r=self.r, mu=self.mu, uid=self.uidMSMS,
                           msmsFile=self.msmsFile, dir=self.outpathmsms)
        else:
            H0 = self.H0
            assignPositions=False
        if self.foldInitialAFs:
            idx = H0.mean(0) > 0.5
            H0.iloc[:, idx.values] = 1 - H0.iloc[:, idx.values]
        self.setH0(H0)
        if assignPositions:
            self.positions_msms = self.H0.columns.values.copy(True)
            self.positions = sorted(np.random.choice(self.L, self.H0.shape[1], replace=False))
        self.H0 = pd.DataFrame(self.H0.values, columns=self.positions)
        self.X0 = self.H0.mean(0).values

    def forwardSimulation(self):
        """
        returns np 3D array T x nSS x R which T=|{t_1,t_2,..}| (nnumber of times), nSS is number of SS , and R is the number of replicates
        """
        import numpy as np
        # df = pd.DataFrame([list(i.genotype(j)) for j in range(2) for i in self.initDiploidPop.individuals()])
        if self.posUnderSelection<0 and self.initDiploidPop is None:
            self.createInitHaps()
            self.set_BeneficialLoci()
            self.initDiploidPop=self.createInitialDiploidPopulation()
        elif self.initDiploidPop is None:
            self.createInitHaps()
            self.initDiploidPop = self.createInitialDiploidPopulation()
            # self.X0=self.H0.mean().values
        else:
            self.X0=Simulation.getGT(self.initDiploidPop).mean().values

        # df = pd.DataFrame([list(i.genotype(j)) for j in range(2) for i in self.initDiploidPop.individuals()])
        # print pd.concat([df.mean(),self.H0.mean().reset_index(drop=True)],1)
        self.X=np.array([self.simualte() for _ in range(self.numReplicates)]).swapaxes(0, 2).swapaxes(0, 1)
        self.X=np.append(np.tile(self.X0[:,None],(1,self.X.shape[2]))[None,:,:],self.X,axis=0)

        self.sampleDepths()
        if self.save:
            pd.to_pickle(self,self.outpath+self.uid+'.pkl')
        # self.createDF()

    def getGenerationTimes(self,step=None,includeZeroGeneration=True):
        if step is None: step=self.generationStep
        times= np.arange(0,self.maxGeneration-self.startGeneration+1,step)
        if includeZeroGeneration:
            return times
        else:
            return times[1:]

    def getTrueGenerationTimes(self,step=None,includeZeroGeneration=True):
        if step is None: step=self.generationStep
        times= np.arange(self.startGeneration,self.maxGeneration+1,step)
        if includeZeroGeneration:
            return times
        else:
            return times[1:]

    @staticmethod
    def getFixationTime(s,Ne=200,roundto10=True):
        if s==0: s=0.01
        t=-4*int(logit(1./Ne)/s)
        if roundto10:
            return (t//10 +1)*10
        else:
            return t

    @staticmethod
    def sampleInitSamplingTime(s,Ne=200,phase=0,samplingWindow=50,startOfEpoch=False):
        fix=Simulation.getFixationTime(s, Ne=Ne)
        if phase==0:    lower,upper=(0, fix-samplingWindow)
        if phase==1:    lower,upper=(0, fix/3-samplingWindow)
        if phase==2:    lower,upper=(fix/3, 2*fix/3-samplingWindow)
        if phase==3:    lower,upper=(2*fix/3, fix-samplingWindow)
        if startOfEpoch:
            rnd=lower
        else:
            rnd=np.random.randint(lower,max(lower,upper)+1)
        return int(rnd)//10 *10
    @staticmethod
    def sampleStartTimesforAlls(samplingWindow=50):
        S=[0.1, 0.05, 0.02, 0.01,0]
        for phase in [1,2,3]:
            pd.DataFrame([[Simulation.sampleInitSamplingTime(s, phase=phase, samplingWindow=samplingWindow, startOfEpoch=True) for _ in range(100)] for s in S], index=S).T.to_pickle('/home/arya/out/startSamplingTimes.phase{}.sampleWin{}.pkl'.format(phase, samplingWindow))

    def setSamplingTimes(self,maxGeneration=None,numSamples=5,step=None,startGeneration=None):
            GT=pd.Series(range(len(self.getTrueGenerationTimes(includeZeroGeneration=True))),index=self.getTrueGenerationTimes(includeZeroGeneration=True))
            if startGeneration is not None:   self.startGeneration=startGeneration
            if maxGeneration is not None:   self.maxGeneration = maxGeneration
            if step is not None:self.generationStep=step
            else:   self.generationStep=(self.maxGeneration-self.startGeneration)/numSamples
            i = GT.loc[self.getTrueGenerationTimes(includeZeroGeneration=True)[:self.X.shape[0]]].values
            self.X = self.X[i, :, :]
            self.C = self.C.apply(lambda x: x[i, :, :])
            self.D = self.D.apply(lambda x: x[i, :, :])
            self.X0=self.X[0,:,0]

    @staticmethod
    def getSamplingTimeBasedOnFreq(sim,phase,samplingWin=50):
        carrier_freq=[0.1,0.5,0.9][phase-1]
        a= np.where(sim.X[:,sim.siteUnderSelection,:].mean(1)>carrier_freq)[0]
        ft=sim.getTrueGenerationTimes().max()
        if len(a):
            t= sim.getTrueGenerationTimes()[np.where(sim.X[:,sim.siteUnderSelection,:].mean(1)>carrier_freq)[0].min()]
        else:
            t=sim.getTrueGenerationTimes().max()
        return min(t,ft-samplingWin)

    @staticmethod
    def Load(s=0.1, experimentID=0, nu0=0.005, numReplicates=3, step=10, ModelName='TimeSeries', samplingWindow=50,
             L=50000, depthRate=30):
        if not s: nu0=0.005
        sim = Simulation.load(s=s, experimentID=experimentID % 100, nu0=nu0, numReplicates=numReplicates, step=step,
                              ExperimentName=ModelName, All=True, L=L, replicates=range(numReplicates),
                              coverage=depthRate)
        sim.experimentID=experimentID
        startGen=0
        sim.setSamplingTimes(maxGeneration=min(startGen+samplingWindow,sim.getTrueGenerationTimes()[-1]),step=step,startGeneration=startGen)
        sim.createDF()
        return sim

    def getHardSweepMutations(self):
        MAF=1./self.H0.shape[0]
        dups=self.H0[self.H0.duplicated()]
        x0=pd.Series(self.X0, index=self.positions)
        hard=[]
        for _,dup in dups.iterrows():
            numDup=self.H0.apply(lambda x:(x==dup).all(),axis=1).sum()
            hard=np.append(hard, (dup*x0==numDup*MAF).replace({False:None}).dropna().index.values)
        hard=np.sort(np.append(hard,(x0==MAF).replace({False:None}).dropna().index.values).astype(int))
        return hard

    @property
    def df(self):
        reps=range(self.numReplicates)
        self.df=pd.concat([pd.DataFrame(self.X[:,:,r],columns=self.positions,index=pd.MultiIndex.from_product([[r],range(self.X.shape[0])],names=['REP','TIME'])).T for r in reps],axis=1)
        if self.numReplicates==1:
            self.df=self.df[0]
        return self.df
    def computeCDi(self, EE, depthRate):
        E = EE.loc[depthRate]
        index = pd.Series(range(E.shape[0]), E.index)
        C = pd.concat([pd.DataFrame(self.C.loc[depthRate][:, :, r], columns=self.H0.columns,
                                    index=pd.MultiIndex.from_product([[r], self.getTrueGenerationTimes()],
                                                                     names=['REP', 'GEN'])).T for r in
                       range(self.numReplicates)], axis=1)
        D = pd.concat([pd.DataFrame(self.D.loc[depthRate][:, :, r], columns=self.H0.columns,
                                    index=pd.MultiIndex.from_product([[r], self.getTrueGenerationTimes()],
                                                                     names=['REP', 'GEN'])).T for r in
                       range(self.numReplicates)], axis=1)
        self.cd = pd.concat([pd.Series(zip(C[i], D[i])) for i in C.columns], axis=1)
        self.cd.columns = C.columns;
        self.cd.index = C.index
        self.cdi = self.cd.applymap(lambda x: index.loc[x])

    def sampleDepths(self,depths = [30, 100, 300]):
        self.D = pd.Series(None, index=depths)
        self.C = pd.Series(None, index=depths)
        for depthRate in depths:
            self.D.loc[depthRate] = np.random.poisson(depthRate,
                                                      self.X.shape[0] * self.X.shape[1] * self.X.shape[2]).reshape(
                self.X.shape).astype(object)
            self.C.loc[depthRate] = np.array([np.random.binomial(d, x) for x, d in
                                              zip(self.X.reshape(-1), self.D.loc[depthRate].reshape(-1))]).reshape(
                    self.X.shape).astype(object)

    @staticmethod
    def sampleDepthX(X,cov):
        D= np.random.poisson(cov,X.size)
        C= np.array([np.random.binomial(d, x) for x, d in zip(X, D)])
        return C,D
    @staticmethod
    def sampleDepthXSeries(X,cov):
        C,D=Simulation.sampleDepthX(X.values,cov)
        a=pd.DataFrame([C,D],columns=X.index,index=['C','D']).T
        return a

    @staticmethod
    def computeCDdf(a, E):
        index = pd.Series(range(E.shape[0]), E.index)
        def f(x):
            try:
                return index.loc[x]
            except:
                return -1
        z=a.groupby(level=[0,1],axis=1).apply(lambda x: x.apply(lambda y:(y.iloc[0],y.iloc[1]),1)).applymap(f)
        return z[(z<0).sum(1)==0]
    def getCD(self,coverage):
        T=self.getTrueGenerationTimes()
        Ti=T
        if T[-1]!=self.C[coverage].shape[0]-1: Ti=range(self.C[coverage].shape[0])
        C=pd.concat([pd.DataFrame(self.C[coverage][Ti,:,i],columns=self.positions,index=T).T for i in range(self.numReplicates)],1,keys=range(self.C[coverage].shape[2]))
        D=pd.concat([pd.DataFrame(self.D[coverage][Ti,:,i],columns=self.positions,index=T).T for i in range(self.numReplicates)],1,keys=range(self.C[coverage].shape[2]))
        CD=pd.concat([C,D],1,keys=['C','D']).reorder_levels([1,2,0],1).sort_index(1)
        CD.columns.names=['REP','GEN','READ']
        return CD

    @staticmethod
    def Recombinator(rate, loci):
        """
        Recombination at loci, after variant index. Loci can take value in [0, NumSNPs-1]
        Args:
            rate: recombination rate
            loci: index of the loci in which rec is is being performed
        Returns: recombinator which is an argument of Simulation, _simulation2 and evolve. It can be list of loci
        """
        if not isinstance(loci, list):
            loci = [loci]
        return sim.Recombinator(intensity=rate, loci=loci)

class POP:
    @staticmethod
    def createISOGenicDiploidPopulation(df):
        """
        initHaps : np 2D array which m x nSS where m i number of individual haps and nSS is number of SS
        return a homozygote diploid population which every haplotype is copied n times
        """
        pop = sim.Population(size=df.shape[0], ploidy=2, loci=df.shape[1], lociPos=list(df.columns),
                             infoFields='fitness')
        for (i, (_, h)) in zip(pop.individuals(), df.iterrows()):
            i.setGenotype(h.tolist(), 0);
            i.setGenotype(h.tolist(), 1)
        return pop

    @staticmethod
    def toDF(pop):
        x = pd.concat(map(pd.DataFrame, [map(list, [i.genotype(0), i.genotype(1)]) for i in pop.allIndividuals()]),
                      keys=range(pop.popSize()))
        x.columns = list(pop.lociPos())
        return x

    @staticmethod
    def freq(pop):
        sim.stat(pop, alleleFreq=range(pop.numLoci()[0]), vars=['alleleFreq'])
        return pd.Series(pd.DataFrame(pop.vars()['alleleFreq']).loc[1].reindex().values,map(int,pop.lociPos())).fillna(0)

    @staticmethod
    def Haplotypes(pop,counts=False,unique=True):
        if isinstance(pop,sim.Population):
            a=POP.toDF(pop)
        else:
            a=pop

        H=a.reset_index(drop=True)
        H.columns=map(int,H.columns)
        b=H.loc[H.sum(1).sort_values().index].astype(str).apply(lambda x: ''.join(x), 1).reset_index(drop=True)
        if counts:
            return b.value_counts().sort_index()
        else:
            if unique:
                b=b.drop_duplicates()
            return b.loc[b.sort_values().index].reset_index(drop=True)

    @staticmethod
    def establish(H, ba, k=5):
        N = H.shape[0]
        car = H[H[ba] == 1]
        n = car.shape[0]
        return pd.concat([car.iloc[np.random.choice(n, k)], H.iloc[np.random.choice(N, N - k)]]).reset_index(drop=True)

class Drift:
    @staticmethod
    def nextGeneration(N,x):
        return (np.random.random(N)<=x).mean()

    @staticmethod
    def sampleReads(D,x):
        return [Drift.sampleReadsDerived(D,x),D]

    @staticmethod
    def sampleReadsDerived(D,x):
        return (np.random.random(D)<=x).sum()

    @staticmethod
    def simulateAF(N,x,T):
        Xt=[]
        for i in range(1, T[-1]+1):
            x=Drift.nextGeneration(N,x)
            if i in T:Xt.append(x)
        return Xt

    @staticmethod
    def simulatePoolCD(N,n,cd):
        x=cd[0].C/float(cd[0].D)
        D=cd.xs('D',level=1)
        Xt=[]
        for i in range(1, D.index[-1]+1):
            x=Drift.nextGeneration(N,x)
            if i in D.index:
                y=Drift.nextGeneration(n,x)
                Xt.append(Drift.sampleReads(D[i], y))
        return pd.DataFrame([[cd[0].C,cd[0].D]]+Xt,index=D.index,columns=['C','D'])

    @staticmethod
    def simulatePoolDerivd(N,n,cd):
        x=cd[0].C/float(cd[0].D)
        D=cd.xs('D',level=1)
        Xt=[]
        for i in range(1, D.index[-1]+1):
            x=Drift.nextGeneration(N,x)
            if i in D.index:
                Xt+=[Drift.sampleReadsDerived(D[i], Drift.nextGeneration(n,x))]
        return [cd[0].C]+Xt

    @staticmethod
    def simulatePools(N,cd,M):
        return pd.concat([Drift.simulatePool(N,cd) for _ in range(M)],keys=range(M))

    @staticmethod
    def simulateAFs(N,x,T,M):
        return pd.DataFrame([Drift.simulateAF(N,x,T) for _ in range(M)],columns=T)


#!/usr/bin/env python
#    RegCM postprocessing tool
#    Copyright (C) 2014 Aliou, Addisu, Kanhu, Andrey

#    This program 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, either version 3 of the License, or
#    (at your option) any later version.

#    This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
import matplotlib.pyplot as plt
import cartopy
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from value import Value

class Plotter(object):

    def __init__(self, value):
        self._value = value
        self.lat, self.lon = value.latlon

    def plot(self, coastlines=True,
                   countries=True,
                   places=True,
                   title=None,
                   levels = None):
        if levels is not None:
            l_min, l_max = levels
            l = (l_max - l_min) / 10
            levels = range(l_min, l_max + l, l)
        projection = ccrs.PlateCarree()
        self.fig, self.ax = plt.subplots(subplot_kw={'projection': projection})

        if coastlines:
            self.ax.coastlines('10m')
        if countries:
            countries = cfeature.NaturalEarthFeature(
                        scale='110m', category='cultural', name='admin_0_countries')
            self.ax.add_feature(countries, color='r', alpha=0.1)
        if places:
            places = cfeature.NaturalEarthFeature(
                        scale='110m', category='cultural', name='populated_places')
            self.ax.add_feature(places, color='b', hatch='o')

        cx = self.ax.contourf(self.lon, self.lat, self._value.data, transform=ccrs.PlateCarree(),cmap='bwr', levels=levels)

        # To mask out OCEAN or LAND
        #ax.add_feature(cfeature.OCEAN)
        #ax.add_feature(cfeature.LAND)

        self.ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                        linewidth=1, color='blue', alpha=0.5, linestyle='-')

        self.fig.colorbar(cx)

        times = self._value.limits['time']
        plt.title(self._value.title + ' [' + self._value.units + ']\n' +
            'mean between ' + str(times[0]) + ' and ' + str(times[1]) + '\n')

    def show(self):
        plt.show()

    def save(self, filename, format):
        plt.savefig(filename + '.' + format)

    def close(self):
        plt.close(self.fig)

if __name__ == "__main__":
    pass


"""
======================================================
Compute source power spectral density (PSD) in a label
======================================================

Returns an STC file containing the PSD (in dB) of each of the sources
within a label.
"""
# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#
# License: BSD (3-clause)

import matplotlib.pyplot as plt

import mne
from mne import io
from mne.datasets import sample
from mne.minimum_norm import read_inverse_operator, compute_source_psd

print(__doc__)

###############################################################################
# Set parameters
data_path = sample.data_path()
raw_fname = data_path + '/MEG/sample/sample_audvis_raw.fif'
fname_inv = data_path + '/MEG/sample/sample_audvis-meg-oct-6-meg-inv.fif'
fname_label = data_path + '/MEG/sample/labels/Aud-lh.label'

# Setup for reading the raw data
raw = io.read_raw_fif(raw_fname, verbose=False)
events = mne.find_events(raw, stim_channel='STI 014')
inverse_operator = read_inverse_operator(fname_inv)
raw.info['bads'] = ['MEG 2443', 'EEG 053']

# picks MEG gradiometers
picks = mne.pick_types(raw.info, meg=True, eeg=False, eog=True,
                       stim=False, exclude='bads')

tmin, tmax = 0, 120  # use the first 120s of data
fmin, fmax = 4, 100  # look at frequencies between 4 and 100Hz
n_fft = 2048  # the FFT size (n_fft). Ideally a power of 2
label = mne.read_label(fname_label)

stc = compute_source_psd(raw, inverse_operator, lambda2=1. / 9., method="dSPM",
                         tmin=tmin, tmax=tmax, fmin=fmin, fmax=fmax,
                         pick_ori="normal", n_fft=n_fft, label=label,
                         dB=True)

stc.save('psd_dSPM')

###############################################################################
# View PSD of sources in label
plt.plot(stc.times, stc.data.T)
plt.xlabel('Frequency (Hz)')
plt.ylabel('PSD (dB)')
plt.title('Source Power Spectrum (PSD)')
plt.show()


#!/usr/bin/env python
"""
Minard's data from Napoleon's 1812-1813  Russian Campaign.
http://www.math.yorku.ca/SCS/Gallery/minard/minard.txt

"""
__author__ = """Aric Hagberg (hagberg@lanl.gov)"""
#    Copyright (C) 2006 by
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.

import string
import networkx as nx

def minard_graph():
    data1="""\
24.0,54.9,340000,A,1
24.5,55.0,340000,A,1
25.5,54.5,340000,A,1
26.0,54.7,320000,A,1
27.0,54.8,300000,A,1
28.0,54.9,280000,A,1
28.5,55.0,240000,A,1
29.0,55.1,210000,A,1
30.0,55.2,180000,A,1
30.3,55.3,175000,A,1
32.0,54.8,145000,A,1
33.2,54.9,140000,A,1
34.4,55.5,127100,A,1
35.5,55.4,100000,A,1
36.0,55.5,100000,A,1
37.6,55.8,100000,A,1
37.7,55.7,100000,R,1
37.5,55.7,98000,R,1
37.0,55.0,97000,R,1
36.8,55.0,96000,R,1
35.4,55.3,87000,R,1
34.3,55.2,55000,R,1
33.3,54.8,37000,R,1
32.0,54.6,24000,R,1
30.4,54.4,20000,R,1
29.2,54.3,20000,R,1
28.5,54.2,20000,R,1
28.3,54.3,20000,R,1
27.5,54.5,20000,R,1
26.8,54.3,12000,R,1
26.4,54.4,14000,R,1
25.0,54.4,8000,R,1
24.4,54.4,4000,R,1
24.2,54.4,4000,R,1
24.1,54.4,4000,R,1"""
    data2="""\
24.0,55.1,60000,A,2
24.5,55.2,60000,A,2
25.5,54.7,60000,A,2
26.6,55.7,40000,A,2
27.4,55.6,33000,A,2
28.7,55.5,33000,R,2
29.2,54.2,30000,R,2
28.5,54.1,30000,R,2
28.3,54.2,28000,R,2"""
    data3="""\
24.0,55.2,22000,A,3
24.5,55.3,22000,A,3
24.6,55.8,6000,A,3
24.6,55.8,6000,R,3
24.2,54.4,6000,R,3
24.1,54.4,6000,R,3"""
    cities="""\
24.0,55.0,Kowno
25.3,54.7,Wilna
26.4,54.4,Smorgoni
26.8,54.3,Moiodexno
27.7,55.2,Gloubokoe
27.6,53.9,Minsk
28.5,54.3,Studienska
28.7,55.5,Polotzk
29.2,54.4,Bobr
30.2,55.3,Witebsk
30.4,54.5,Orscha
30.4,53.9,Mohilow
32.0,54.8,Smolensk
33.2,54.9,Dorogobouge
34.3,55.2,Wixma
34.4,55.5,Chjat
36.0,55.5,Mojaisk
37.6,55.8,Moscou
36.6,55.3,Tarantino
36.5,55.0,Malo-Jarosewii"""

    c={}
    for line in cities.split('\n'):
        x,y,name=line.split(',')
        c[name]=(float(x),float(y))

    g=[]

    for data in [data1,data2,data3]:
        G=nx.Graph()
        i=0
        G.pos={} # location
        G.pop={} # size
        last=None
        for line in data.split('\n'):
            x,y,p,r,n=line.split(',')
            G.pos[i]=(float(x),float(y))
            G.pop[i]=int(p)
            if last is None:
                last=i
            else:
                G.add_edge(i,last,{r:int(n)})
                last=i
            i=i+1
        g.append(G)

    return g,c

if __name__ == "__main__":

    (g,city)=minard_graph()

    try:
        import matplotlib.pyplot as plt
        plt.figure(1,figsize=(11,5))
        plt.clf()
        colors=['b','g','r']
        for G in g:
            c=colors.pop(0)
            node_size=[int(G.pop[n]/300.0) for n in G]
            nx.draw_networkx_edges(G,G.pos,edge_color=c,width=4,alpha=0.5)
            nx.draw_networkx_nodes(G,G.pos,node_size=node_size,node_color=c,alpha=0.5)
            nx.draw_networkx_nodes(G,G.pos,node_size=5,node_color='k')

        for c in city:
            x,y=city[c]
            plt.text(x,y+0.1,c)
        plt.savefig("napoleon_russian_campaign.png")
    except ImportError:
        pass


#!/usr/bin/env python
import zarafa
import matplotlib.pyplot as plt

def opt_args():
    parser = zarafa.parser('skpc')
    parser.add_option('--save', dest='save', action='store',  help='Save plot to file (png)')
    return parser.parse_args()

def b2m(bytes):
    return (bytes / 1024) / 1024

def main():
    options, args = opt_args()
    users = list(zarafa.Server(options).users())

    width = 0.35       # the width of the bars

    fig, ax = plt.subplots()
    ind = range(0, len(users))

    data = [b2m(user.store.size) for user in users]
    rects1 = ax.bar(ind, data, width, color='r')

    data = [len(list(user.store.folders())) for user in users]
    rects2 = ax.bar([offset + width for offset in ind], data, width, color='g')

    data =[sum(folder.count for folder in user.store.folders()) for user in users]
    rects3 = ax.bar([offset + width * 2 for offset in ind], data, width, color='b')

    ax.legend( (rects1[0], rects2[0], rects3[0]), ('Store size (Mb)', 'Folders', 'Items') )

    ax.set_ylabel('Values')
    ax.set_title('Store size, Folder, Items per user')
    ax.set_xticks([offset + width for offset in ind])
    ax.set_xticklabels([user.name for user in users])

    def autolabel(rects):
        # attach some text labels
        for rect in rects:
            height = rect.get_height()
            ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(height),
                    ha='center', va='bottom')

    autolabel(rects1)
    autolabel(rects2)
    autolabel(rects3)

    if options.save:
        plt.savefig(options.save)
    else:
        plt.show()

if __name__ == '__main__':
    main()