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

from sklearn.datasets import load_digits, load_iris
from sklearn.cross_validation import cross_val_score, train_test_split

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]).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.float)
    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)


# Copyright (c) 2019, MD2K Center of Excellence
# - Nasir Ali <nasir.ali08@gmail.com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import numpy as np
import pandas as pd
from geopy.distance import great_circle
from pyspark.sql.functions import pandas_udf, PandasUDFType
from shapely.geometry.multipoint import MultiPoint
from sklearn.cluster import DBSCAN
from pyspark.sql.types import StructField, StructType, StringType, FloatType

EPSILON_CONSTANT = 1000
LATITUDE = 0
LONGITUDE = 1
ACCURACY = -1
GPS_ACCURACY_THRESHOLD = 41.0
KM_PER_RADIAN = 6371.0088
GEO_FENCE_DISTANCE = 2
MINIMUM_POINTS_IN_CLUSTER = 500

def get_centermost_point(cluster: object) -> object:
    """

    :param cluster:
    :return:
    :rtype: object
    """
    centroid = (
        MultiPoint(cluster).centroid.x, MultiPoint(cluster).centroid.y)
    centermost_point = min(cluster, key=lambda point: great_circle(point,
                                                                   centroid).m)
    return tuple(centermost_point)

schema = StructType([
    StructField("user", StringType()),
    StructField("latitude", FloatType()),
    StructField("longitude", FloatType())
])

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def gps_clusters(data: object) -> object:
    """
    Computes the clusters

    :rtype: object
    :param list data: list of interpolated gps data
    :param float geo_fence_distance: Maximum distance between points in a
    cluster
    :param int min_points_in_cluster: Minimum number of points in a cluster
    :return: list of cluster-centroids coordinates
    """
    geo_fence_distance = GEO_FENCE_DISTANCE
    min_points_in_cluster = MINIMUM_POINTS_IN_CLUSTER

    data = data[data.accuracy < GPS_ACCURACY_THRESHOLD]

    id = data.user.iloc[0]
    dataframe = pd.DataFrame(
        {'latitude': data.latitude, 'longitude': data.longitude})
    coords = dataframe.as_matrix(columns=['latitude', 'longitude'])

    epsilon = geo_fence_distance / (
            EPSILON_CONSTANT * KM_PER_RADIAN)
    db = DBSCAN(eps=epsilon, min_samples=min_points_in_cluster,
                algorithm='ball_tree', metric='haversine').fit(
        np.radians(coords))
    cluster_labels = db.labels_
    num_clusters = len(set(cluster_labels))
    clusters = pd.Series(
        [coords[cluster_labels == n] for n in range(-1, num_clusters)])
    clusters = clusters.apply(lambda y: np.nan if len(y) == 0 else y)
    clusters.dropna(how='any', inplace=True)
    centermost_points = clusters.map(get_centermost_point)
    centermost_points = np.array(centermost_points)
    all_centroid = []
    for cols in centermost_points:
        cols = np.array(cols)
        cols.flatten()
        cs = ([id, cols[LATITUDE], cols[LONGITUDE]])
        all_centroid.append(cs)
    df = pd.DataFrame(all_centroid, columns=['user', 'latitude', 'longitude'])
    return df


import matplotlib.pyplot as pl
import numpy as np
import json
import glob

pl.close('all')
pl.ion()

def get_timings(kernel_name, gridsizes):

    dt = np.zeros_like(gridsizes, dtype=float)

    for i,gridsize in enumerate(gridsizes):
        with open( '{0}_{1:03d}.json'.format(kernel_name, gridsize) ) as f:
            data = json.load(f)
            timings = data[0]

            fastest = 1e9
            for timing in timings:
                fastest = min(fastest, timing['time'])
            dt[i] = fastest

    return dt

if __name__ == '__main__':
    gridsize = np.arange(32, 513, 32)
    normal = get_timings('diff_c_g',      gridsize)
    smem   = get_timings('diff_c_g_smem', gridsize)

    fac = gridsize**3

    pl.figure(figsize=(8,4))
    pl.subplot(121)
    pl.plot(gridsize, normal/fac, 'k-x', label='non smem')
    pl.plot(gridsize, smem  /fac, 'r-x', label='smem')
    pl.ylim(0, 2e-7)
    pl.ylabel('time/gridpoint (s)')
    pl.xlabel('gridpoints (-)')
    pl.legend()
    pl.grid()

    pl.subplot(122)
    pl.plot(gridsize, normal/smem, 'k-x')
    pl.ylabel('non_smem/smem (-)')
    pl.xlabel('gridpoints (-)')
    pl.grid()

    pl.tight_layout()


"""
=========================
Multi-dimensional scaling
=========================

An illustration of the metric and non-metric MDS on generated noisy data.

The reconstructed points using the metric MDS and non metric MDS are slightly
shifted to avoid overlapping.
"""

# Author: Nelle Varoquaux <nelle.varoquaux@gmail.com>
# License: BSD

print(__doc__)
import numpy as np

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

from sklearn import manifold
from sklearn.metrics import euclidean_distances
from sklearn.decomposition import PCA

n_samples = 20
seed = np.random.RandomState(seed=3)
X_true = seed.randint(0, 20, 2 * n_samples).astype(np.float)
X_true = X_true.reshape((n_samples, 2))
# Center the data
X_true -= X_true.mean()

similarities = euclidean_distances(X_true)

# Add noise to the similarities
noise = np.random.rand(n_samples, n_samples)
noise = noise + noise.T
noise[np.arange(noise.shape[0]), np.arange(noise.shape[0])] = 0
similarities += noise

mds = manifold.MDS(n_components=2, max_iter=3000, eps=1e-9, random_state=seed,
                   dissimilarity="precomputed", n_jobs=1)
pos = mds.fit(similarities).embedding_

nmds = manifold.MDS(n_components=2, metric=False, max_iter=3000, eps=1e-12,
                    dissimilarity="precomputed", random_state=seed, n_jobs=1,
                    n_init=1)
npos = nmds.fit_transform(similarities, init=pos)

# Rescale the data
pos *= np.sqrt((X_true ** 2).sum()) / np.sqrt((pos ** 2).sum())
npos *= np.sqrt((X_true ** 2).sum()) / np.sqrt((npos ** 2).sum())

# Rotate the data
clf = PCA(n_components=2)
X_true = clf.fit_transform(X_true)

pos = clf.fit_transform(pos)

npos = clf.fit_transform(npos)

fig = plt.figure(1)
ax = plt.axes([0., 0., 1., 1.])

s = 100
plt.scatter(X_true[:, 0], X_true[:, 1], color='navy', s=s, lw=0,
            label='True Position')
plt.scatter(pos[:, 0], pos[:, 1], color='turquoise', s=s, lw=0, label='MDS')
plt.scatter(npos[:, 0], npos[:, 1], color='darkorange', s=s, lw=0, label='NMDS')
plt.legend(scatterpoints=1, loc='best', shadow=False)

similarities = similarities.max() / similarities * 100
similarities[np.isinf(similarities)] = 0

# Plot the edges
start_idx, end_idx = np.where(pos)
# a sequence of (*line0*, *line1*, *line2*), where::
#            linen = (x0, y0), (x1, y1), ... (xm, ym)
segments = [[X_true[i, :], X_true[j, :]]
            for i in range(len(pos)) for j in range(len(pos))]
values = np.abs(similarities)
lc = LineCollection(segments,
                    zorder=0, cmap=plt.cm.Blues,
                    norm=plt.Normalize(0, values.max()))
lc.set_array(similarities.flatten())
lc.set_linewidths(0.5 * np.ones(len(segments)))
ax.add_collection(lc)

plt.show()


#!/usr/bin/env python

'''
Compute mutual information between individual features
and labels
'''

import argparse
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from common import *

def opts():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('features', type=load_npz,
        help='Training data features (npz)')
    parser.add_argument('labels', type=load_npz,
        help='Training data labels (npz)')
    parser.add_argument('output',
        help='Output file with plots (pdf)')
    return parser

if __name__ == "__main__":
    args = opts().parse_args()
    print "Loading labels"
    labels = args.labels['labels']
    header = args.labels['header']
    pdf = PdfPages(args.output)

    #print "Plotting boolean features conditioned on labels"
    #bf = args.features['bfeatures']
    #n = bf.shape[1]
    #m = np.zeros((n,11))
    #m[:,0] = np.sum(bf==-1, axis=0)
    #m[:,1] = np.sum(bf==0, axis=0)
    #m[:,2] = np.sum(bf==1, axis=0)
    #fig = plt.figure()
    #pdf.savefig(fig)
    #plt.close()

    print "Plotting float features conditioned on labels"
    ff = args.features['ffeatures']
    n = ff.shape[1]
    x = np.arange(n)
    for i, l in enumerate(labels.T):
        print "label %d" % i
        for j, f in enumerate(ff.T):
            print "...ffeature %d" % j
            fig = plt.figure()
            plt.hist(f[l], normed=True, label='P(f | l)',
                     color='blue', alpha=0.4,
                     range=(f.min(),f.max()), bins=25)
            plt.hist(f[np.logical_not(l)], normed=True, label='P(f | ~l)',
                     color='green', alpha=0.4,
                     range=(f.min(),f.max()), bins=25)
            plt.xlim(f.min(), f.max())
            plt.xlabel('f')
            plt.ylabel('P(f)')
            plt.title('FFeature %d, Label %s' % (j, header[i]))
            pdf.savefig(fig)
            plt.close()

    pdf.close()


#!/usr/bin/env python
# encoding: utf-8
"""
Example of creating a block model using the blockmodel function in NX.  Data used is the Hartford, CT drug users network:

@article{,
	title = {Social Networks of Drug Users in {High-Risk} Sites: Finding the Connections},
	volume = {6},
	shorttitle = {Social Networks of Drug Users in {High-Risk} Sites},
	url = {http://dx.doi.org/10.1023/A:1015457400897},
	doi = {10.1023/A:1015457400897},
	number = {2},
	journal = {{AIDS} and Behavior},
	author = {Margaret R. Weeks and Scott Clair and Stephen P. Borgatti and Kim Radda and Jean J. Schensul},
	month = jun,
	year = {2002},
	pages = {193--206}
}

"""

__author__ = """\n""".join(['Drew Conway <drew.conway@nyu.edu>',
                            'Aric Hagberg <hagberg@lanl.gov>'])

from collections import defaultdict
import networkx as nx
import numpy
from scipy.cluster import hierarchy
from scipy.spatial import distance
import matplotlib.pyplot as plt

def create_hc(G):
    """Creates hierarchical cluster of graph G from distance matrix"""
    path_length=nx.all_pairs_shortest_path_length(G)
    distances=numpy.zeros((len(G),len(G)))
    for u,p in path_length.items():
        for v,d in p.items():
            distances[u][v]=d
    # Create hierarchical cluster
    Y=distance.squareform(distances)
    Z=hierarchy.complete(Y)  # Creates HC using farthest point linkage
    # This partition selection is arbitrary, for illustrive purposes
    membership=list(hierarchy.fcluster(Z,t=1.15))
    # Create collection of lists for blockmodel
    partition=defaultdict(list)
    for n,p in zip(list(range(len(G))),membership):
        partition[p].append(n)
    return list(partition.values())

if __name__ == '__main__':
    G=nx.read_edgelist("hartford_drug.edgelist")

    # Extract largest connected component into graph H
    H=nx.connected_component_subgraphs(G)[0]
    # Makes life easier to have consecutively labeled integer nodes
    H=nx.convert_node_labels_to_integers(H)
    # Create parititions with hierarchical clustering
    partitions=create_hc(H)
    # Build blockmodel graph
    BM=nx.blockmodel(H,partitions)

    # Draw original graph
    pos=nx.spring_layout(H,iterations=100)
    fig=plt.figure(1,figsize=(6,10))
    ax=fig.add_subplot(211)
    nx.draw(H,pos,with_labels=False,node_size=10)
    plt.xlim(0,1)
    plt.ylim(0,1)

    # Draw block model with weighted edges and nodes sized by number of internal nodes
    node_size=[BM.node[x]['nnodes']*10 for x in BM.nodes()]
    edge_width=[(2*d['weight']) for (u,v,d) in BM.edges(data=True)]
    # Set positions to mean of positions of internal nodes from original graph
    posBM={}
    for n in BM:
        xy=numpy.array([pos[u] for u in BM.node[n]['graph']])
        posBM[n]=xy.mean(axis=0)
    ax=fig.add_subplot(212)
    nx.draw(BM,posBM,node_size=node_size,width=edge_width,with_labels=False)
    plt.xlim(0,1)
    plt.ylim(0,1)
    plt.axis('off')
    plt.savefig('hartford_drug_block_model.png')


import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import pylab

########################################################################
# 2017 -  Manny Grewal
# Purpose of this class is to visualise a list of images from the CIFAR dataset

# How many columns to show in a grid
MAX_COLS = 5

#PlotImages method takes an list of Images and their respective labels in the second parameter
#Then it renders them using matplotlib imshow method in a 5 column matrix
def PlotImages(arrayImages,arrayClassLabels,reShapeRequired=False):
    totalImages=len(arrayImages)
    if(reShapeRequired==True):
        arrayImages = np.reshape(arrayImages, (totalImages,32,32,3))

    totalRows= math.ceil(totalImages/MAX_COLS)
    fig = plt.figure(figsize=(5,5))
    gs = gridspec.GridSpec(totalImages, MAX_COLS)
    # set the space between subplots and the position of the subplots in the figure
    gs.update(wspace=0.1, hspace=0.4, left = 0.1, right = 0.7, bottom = 0.1, top = 0.9)

    arrayIndex=0
    for g in gs:
        if(arrayIndex<totalImages):
            axes=plt.subplot(g)
            axes.set_axis_off()
            axes.set_title(arrayClassLabels[arrayIndex])
            axes.imshow(arrayImages[arrayIndex])
            arrayIndex+=1
    #plt.show()


from .pls_ import _PLS

__all__ = ['CCA']

class CCA(_PLS):
    """CCA Canonical Correlation Analysis.

    CCA inherits from PLS with mode="B" and deflation_mode="canonical".

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

    Parameters
    ----------
    n_components : int, (default 2).
        number of components to keep.

    scale : boolean, (default True)
        whether to scale the data?

    max_iter : an integer, (default 500)
        the maximum number of iterations of the NIPALS inner loop

    tol : non-negative real, default 1e-06.
        the tolerance used in the iterative algorithm

    copy : boolean
        Whether the deflation be done on a copy. Let the default value
        to True unless you don't care about side effects

    Attributes
    ----------
    x_weights_ : array, [p, n_components]
        X block weights vectors.

    y_weights_ : array, [q, n_components]
        Y block weights vectors.

    x_loadings_ : array, [p, n_components]
        X block loadings vectors.

    y_loadings_ : array, [q, n_components]
        Y block loadings vectors.

    x_scores_ : array, [n_samples, n_components]
        X scores.

    y_scores_ : array, [n_samples, n_components]
        Y scores.

    x_rotations_ : array, [p, n_components]
        X block to latents rotations.

    y_rotations_ : array, [q, n_components]
        Y block to latents rotations.

    n_iter_ : array-like
        Number of iterations of the NIPALS inner loop for each
        component.

    Notes
    -----
    For each component k, find the weights u, v that maximizes
    max corr(Xk u, Yk v), such that ``|u| = |v| = 1``

    Note that it maximizes only the correlations between the scores.

    The residual matrix of X (Xk+1) block is obtained by the deflation on the
    current X score: x_score.

    The residual matrix of Y (Yk+1) block is obtained by deflation on the
    current Y score.

    Examples
    --------
    >>> from sklearn.cross_decomposition import CCA
    >>> X = [[0., 0., 1.], [1.,0.,0.], [2.,2.,2.], [3.,5.,4.]]
    >>> Y = [[0.1, -0.2], [0.9, 1.1], [6.2, 5.9], [11.9, 12.3]]
    >>> cca = CCA(n_components=1)
    >>> cca.fit(X, Y)
    ... # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
    CCA(copy=True, max_iter=500, n_components=1, scale=True, tol=1e-06)
    >>> X_c, Y_c = cca.transform(X, Y)

    References
    ----------

    Jacob A. Wegelin. A survey of Partial Least Squares (PLS) methods, with
    emphasis on the two-block case. Technical Report 371, Department of
    Statistics, University of Washington, Seattle, 2000.

    In french but still a reference:
    Tenenhaus, M. (1998). La regression PLS: theorie et pratique. Paris:
    Editions Technic.

    See also
    --------
    PLSCanonical
    PLSSVD
    """

    def __init__(self, n_components=2, scale=True,
                 max_iter=500, tol=1e-06, copy=True):
        _PLS.__init__(self, n_components=n_components, scale=scale,
                      deflation_mode="canonical", mode="B",
                      norm_y_weights=True, algorithm="nipals",
                      max_iter=max_iter, tol=tol, copy=copy)


#!/usr/bin/env python

from numpy import zeros,inner
import numpy as np
import re
from pyDFTutils.ase_utils import symbol_number
import matplotlib.pyplot as plt
def fix_line(line):
    line=re.sub("(\d)-(\d)", r'\1 -\2',line)
    return line

class procar_reader():
    def __init__(self,fname='PROCAR'):
        self.read(fname=fname)

    def get_dos(self,iion,orb_name,iband):
        dos=inner(self.dos_array[iion,self.orb_dict[orb_name],iband,:],self.weight)
        return dos

    def filter_band(self,iion,orb_name,thr=0.01):
        """
        return a energy array 2D. energy(iband,nkpt).
        """
        band_ids=[]
        for iband in range(self.nbands):
            d=self.get_dos(iion,orb_name,iband)
            print(d)
            if d>thr:
                band_ids.append(iband)
        return self.energy[np.array(band_ids,dtype=int)]

    def plot_band(self,iion,orb_name,thr=0.01):
        earray=self.filter_band(iion,orb_name,thr=thr)
        for k_e in earray:
            plt.plot(k_e)
        plt.ylim(-3,2)
        #plt.show()
    def plot_band_alpha(self,iion,orb_name,color='k'):
        for iband in range(self.nbands):
            d=self.get_dos(iion,orb_name,iband)
            print(d)
            plt.plot(self.energy[iband],color,linewidth=d*50,alpha=0.5)

    def read(self,fname='PROCAR'):
        lines=open(fname).readlines()
        iline=0
        self.has_phase=bool(lines[iline].rfind('phase'))
        iline=1
        p=lines[iline].split()
        self.nkpts=int(p[3])
        self.nbands=int(p[7])
        self.nions=int(p[11])

        self.dos_label=lines[7].split()[1:-1]

        self.norb=len(self.dos_label)
        self.orb_dict=dict(list(zip(self.dos_label,list(range(self.norb)))))
        print(self.orb_dict)

        self.dos_array=zeros((self.nions,self.norb,self.nbands,self.nkpts),dtype='float')

        self.weight=zeros(self.nkpts,dtype='float')
        self.energy=zeros((self.nbands,self.nkpts))
        self.band_occ=zeros((self.nbands,self.nkpts))
        self.kpts=zeros((self.nkpts,3))

        iline+=1
        for ikpts in range(self.nkpts):
            iline+=1

            line_k=fix_line( lines[iline]).split()

            self.kpts[ikpts]=[float(x) for x in line_k[3:6]]
            self.weight[ikpts]=float(line_k[-1])

            iline+=2
            for iband in range(self.nbands):

                line_b=lines[iline].split()

                self.energy[iband,ikpts]=float(line_b[4])
                self.band_occ[iband,ikpts]=float(line_b[7])

                iline+=3
                for iion in range(self.nions):

                    #print iline
                    line_dos=lines[iline].strip().split()
                    #print iline
                    #print line_dos
                    self.dos_array[iion,:,iband,ikpts]=[float(x) for x in line_dos[1:-1]]
                    iline+=1
                    #if self.has_phase:
                    #iline+=1+self.nions*2
                iline+=3
        self.efermi=np.max(self.energy[self.band_occ>0.5])
        print(self.efermi)
        self.energy=self.energy-self.efermi
def test(iion=0,orb_name='dx2',thr=0.005):
    p=procar_reader()
    #for e in p.filter_band(0,orb_name,thr=thr):
    #    plt.plot(e,'.',color='green')
    #for e in p.filter_band(1,'dx2',thr=thr):
    #    plt.plot(e,'-',color='red')
        #plt.plot(p.filter_band(iion,'dz2',thr=thr))
    p.plot_band_alpha(1,'dx2',color='r')
    p.plot_band_alpha(1,'dz2',color='g')
    plt.ylim(-5,5)
    plt.show()

if __name__=='__main__':
    test()


'''
Easy integration of DataFrame into pyqt framework

Copyright: Jev Kuznetsov
Licence: BSD

'''
from PyQt4.QtCore import (QAbstractTableModel,Qt,QVariant,QModelIndex,SIGNAL)
from PyQt4.QtGui import (QApplication,QDialog,QVBoxLayout, QHBoxLayout, QTableView, QPushButton,
                         QWidget,QTableWidget, QHeaderView, QFont,QMenu,QAbstractItemView)

from pandas import DataFrame, Index

class DataFrameModel(QAbstractTableModel):
    ''' data model for a DataFrame class '''
    def __init__(self,parent=None):
        super(DataFrameModel,self).__init__(parent)
        self.df = DataFrame()
        self.columnFormat = {} # format columns

    def setFormat(self,fmt):
        """
        set string formatting for the output
        example : format = {'close':"%.2f"}
        """

        self.columnFormat = fmt

    def setDataFrame(self,dataFrame):
        self.df = dataFrame

        self.signalUpdate()

    def signalUpdate(self):
        ''' tell viewers to update their data (this is full update, not efficient)'''
        self.layoutChanged.emit()

    def __repr__(self):
        return str(self.df)

    def setData(self,index,value, role=Qt.EditRole):

        if index.isValid():
            row,column = index.row(), index.column()
            dtype = self.df.dtypes.tolist()[column] # get column dtype

            if np.issubdtype(dtype,np.float):
                val,ok = value.toFloat()
            elif np.issubdtype(dtype,np.int):
                val,ok = value.toInt()
            else:
                val = value.toString()
                ok = True

            if ok:
                self.df.iloc[row,column] = val
                return True

        return False

    def flags(self, index):
        if not index.isValid():
            return Qt.ItemIsEnabled
        return Qt.ItemFlags(
                QAbstractTableModel.flags(self, index)|
                Qt.ItemIsEditable)

    def appendRow(self, index, data=0):
        self.df.loc[index,:] = data
        self.signalUpdate()

    def deleteRow(self, index):
        idx = self.df.index[index]
        #self.beginRemoveRows(QModelIndex(), index,index)
        #self.df = self.df.drop(idx,axis=0)
        #self.endRemoveRows()
        #self.signalUpdate()

    #------------- table display functions -----------------
    def headerData(self,section,orientation,role=Qt.DisplayRole):
        if role != Qt.DisplayRole:
            return QVariant()

        if orientation == Qt.Horizontal:
            try:
                return self.df.columns.tolist()[section]
            except (IndexError, ):
                return QVariant()
        elif orientation == Qt.Vertical:
            try:
                #return self.df.index.tolist()
                return str(self.df.index.tolist()[section])
            except (IndexError, ):
                return QVariant()

    def data(self, index, role=Qt.DisplayRole):
        if role != Qt.DisplayRole:
            return QVariant()

        if not index.isValid():
            return QVariant()

        col = self.df.ix[:,index.column()] # get a column slice first to get the right data type
        elm = col[index.row()]
        #elm = self.df.ix[index.row(),index.column()]

        if self.df.columns[index.column()] in self.columnFormat.keys():
            return QVariant(self.columnFormat[self.df.columns[index.column()]] % elm )
        else:
            return QVariant(str(elm))

    def sort(self,nCol,order):

        self.layoutAboutToBeChanged.emit()
        if order == Qt.AscendingOrder:
            self.df = self.df.sort(columns=self.df.columns[nCol], ascending=True)
        elif order == Qt.DescendingOrder:
            self.df = self.df.sort(columns=self.df.columns[nCol], ascending=False)

        self.layoutChanged.emit()

    def rowCount(self, index=QModelIndex()):
        return self.df.shape[0]

    def columnCount(self, index=QModelIndex()):
        return self.df.shape[1]

class TableView(QTableView):
    """ extended table view """
    def __init__(self,name='TableView1', parent=None):
        super(TableView,self).__init__(parent)
        self.name = name
        self.setSelectionBehavior(QAbstractItemView.SelectRows)

    def contextMenuEvent(self, event):
        menu = QMenu(self)

        Action = menu.addAction("delete row")
        Action.triggered.connect(self.deleteRow)

        menu.exec_(event.globalPos())

    def deleteRow(self):
        print "Action triggered from " + self.name

        print 'Selected rows:'
        for idx in self.selectionModel().selectedRows():
            print idx.row()
           # self.model.deleteRow(idx.row())

class DataFrameWidget(QWidget):
    ''' a simple widget for using DataFrames in a gui '''
    def __init__(self,name='DataFrameTable1', parent=None):
        super(DataFrameWidget,self).__init__(parent)
        self.name = name

        self.dataModel = DataFrameModel()
        self.dataModel.setDataFrame(DataFrame())

        self.dataTable = QTableView()
        #self.dataTable.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.dataTable.setSortingEnabled(True)

        self.dataTable.setModel(self.dataModel)
        self.dataModel.signalUpdate()

        #self.dataTable.setFont(QFont("Courier New", 8))

        layout = QVBoxLayout()
        layout.addWidget(self.dataTable)
        self.setLayout(layout)

    def setFormat(self,fmt):
        """ set non-default string formatting for a column """
        for colName, f in fmt.iteritems():
            self.dataModel.columnFormat[colName]=f

    def fitColumns(self):
        self.dataTable.horizontalHeader().setResizeMode(QHeaderView.Stretch)

    def setDataFrame(self,df):
        self.dataModel.setDataFrame(df)

    def resizeColumnsToContents(self):
        self.dataTable.resizeColumnsToContents()

    def insertRow(self,index, data=None):
        self.dataModel.appendRow(index,data)

#-----------------stand alone test code

def testDf():
    ''' creates test dataframe '''
    data = {'int':[1,2,3],'float':[1./3,2.5,3.5],'string':['a','b','c'],'nan':[np.nan,np.nan,np.nan]}
    return DataFrame(data, index=Index(['AAA','BBB','CCC']))[['int','float','string','nan']]

class Form(QDialog):
    def __init__(self,parent=None):
        super(Form,self).__init__(parent)

        df = testDf() # make up some data
        self.table = DataFrameWidget(parent=self)
        self.table.setDataFrame(df)
        #self.table.resizeColumnsToContents()
        self.table.fitColumns()
        self.table.setFormat({'float': '%.2f'})

        #buttons
       #but_add = QPushButton('Add')
        but_test = QPushButton('Test')
        but_test.clicked.connect(self.testFcn)
        hbox = QHBoxLayout()
        #hbox.addself.table(but_add)
        hbox.addWidget(but_test)

        layout = QVBoxLayout()
        layout.addWidget(self.table)
        layout.addLayout(hbox)

        self.setLayout(layout)

    def testFcn(self):
        print 'test function'
        self.table.insertRow('foo')

if __name__=='__main__':
    import sys
    import numpy as np

    app = QApplication(sys.argv)
    form = Form()
    form.show()
    app.exec_()


import operation as op
import random
import math
import Helper
import GRDEI
import RDE
import GDE
import Wave
import os
import numpy
import matplotlib.pyplot as plt
numpy.set_printoptions(threshold=numpy.nan)

#def encode(payload_path,cover_path,threshold,segment_size,partition_segment_size,method):
#    file_name = cover_path.split("\\")[len(cover_path.split("\\"))-1]
#    path = cover_path.replace(file_name,'')
#    print(" PROSES ENCODING ")
#    print(" METODE YANG DIGUNAKAN : " + method)
#    payload = Helper.payloadIO.open(payload_path)
#    bin_payload = op.operation.stringToBinary(payload)
#    print "besar payload : "
#    payload_size = len(bin_payload)
#    print(payload_size)
#    medium = Helper.WavIO.open(cover_path)
#    print "bitrate: "
#    print(medium.bitrate)

#    samples = op.operation.numToBinary(medium.samples)
#    (M1,M2,Partisi) = op.operation.intel_partition(samples,partition_segment_size)
#    intM1 = op.operation.binaryTonum(M1)
#    intM2 = op.operation.binaryTonum(M2)
#    if method == "GDE" :
#        kapasitas_M1 = GDE.checkCapacity(intM1,segment_size,threshold)
#        kapasitas_M2 = GDE.checkCapacity(intM2,segment_size,threshold)
#    elif method == "GRDEI":
#        kapasitas_M1 = GRDEI.checkCapacity(intM1,segment_size,threshold)
#        kapasitas_M2 = GRDEI.checkCapacity(intM2,segment_size,threshold)
#    print(" kapasitas segmen 1 : " + str(kapasitas_M1))
#    print(" kapasitas segmen 2 : " + str(kapasitas_M2))
#    capacity = kapasitas_M1+kapasitas_M2
#    print "Kapasitas penyimpanan :"
#    print(capacity)
#    if capacity >= payload_size:
#        print "stegano dapat dilakukan"
#        payload_seg1 = bin_payload[:kapasitas_M1]
#        payload_seg2 = bin_payload[kapasitas_M1:len(bin_payload)]
#        if method == "GDE" :
#            (encoded_1,locMap_1) = GDE.encode(intM1,payload_seg1,segment_size,threshold)
#            (encoded_2,locMap_2) = GDE.encode(intM2,payload_seg2,segment_size,threshold)
#        elif method == "GRDEI":
#            (encoded_1,locMap_1,reduceMap_1) = GRDEI.encode(intM1,payload_seg1,segment_size,threshold)
#            (encoded_2,locMap_2,reduceMap_2) = GRDEI.encode(intM2,payload_seg2,segment_size,threshold)
#        encoded_1 = op.operation.numToBinary(encoded_1)
#        encoded_2 = op.operation.numToBinary(encoded_2)
#        for i in range(len(encoded_1)):
#            encoded_1[i]=encoded_1[i][8:16]
#        for i in range(len(encoded_2)):
#            encoded_2[i]=encoded_2[i][8:16]
#        encoded_1_bin = numpy.array(encoded_1,dtype=int)
#        encoded_2_bin = numpy.array(encoded_2,dtype=int)
#        _M = op.operation.reconstructPartition(encoded_1_bin,encoded_2_bin,Partisi)
#        _M_int2= numpy.asarray(op.operation.binaryTonum(_M),dtype=numpy.uint16)
#        _M_int = numpy.asarray(op.operation.binaryTonum(_M),dtype=numpy.int16)
#        new_wav = Wave.Wave(method+"_encoded_"+str(file_name))
#        new_wav.samples = _M_int
#        new_wav.bitrate = medium.bitrate

#        time_axis = numpy.linspace(0,len(medium.samples)/medium.bitrate,num=len(medium.samples))
#        plt.subplot(3,1,1)
#        plt.title("perbandingan wav asli dan hasil encode")
#        plt.plot(time_axis,numpy.array(medium.samples,dtype= "int16"))
#        plt.ylabel("WAV asli")
#        plt.subplot(3,1,2)
#        plt.plot(time_axis,numpy.array(_M_int,dtype= "int16"))
#        plt.ylabel("Hasil Encode")
#        plt.subplot(3,1,3)
#        plt.plot(time_axis,numpy.array(_M_int,dtype= "int16"))
#        plt.subplot(3,1,3)
#        plt.plot(time_axis,numpy.array(medium.samples,dtype= "int16"))

#        plt.ylabel("perbandingan")
#        plt.xlabel("Waktu (s)")
#        plt.savefig("original-encoded.png")
#        #print(medium.samples)
#        #raw_input()
#        #print(new_wav.samples)
#        #new_wav.print_info()
#        #medium.print_info()
#        Helper.WavIO.write(path,new_wav)
#        if method == "GDE" :
#            return(locMap_1,locMap_2,Partisi)
#        elif method == "GRDEI":
#            return(locMap_1,locMap_2,reduceMap_1,reduceMap_2,Partisi)

#    else:
#        print "kapasitas tidak mencukupi"
#        return -1

def encode(intM1,intM2,payload_seg1,payload_seg2,threshold,segment_size,partition_segment_size,Partisi,method):

    if method == "GDE" :
        (encoded_1,locMap_1) = GDE.encode(intM1,payload_seg1,segment_size,threshold)
        (encoded_2,locMap_2) = GDE.encode(intM2,payload_seg2,segment_size,threshold)
    elif method == "GRDEI":
        (encoded_1,locMap_1,reduceMap_1) = GRDEI.encode(intM1,payload_seg1,segment_size,threshold)
        (encoded_2,locMap_2,reduceMap_2) = GRDEI.encode(intM2,payload_seg2,segment_size,threshold)
    encoded_1 = op.operation.numToBinary(encoded_1)
    encoded_2 = op.operation.numToBinary(encoded_2)
    for i in range(len(encoded_1)):
        encoded_1[i]=encoded_1[i][8:16]
    for i in range(len(encoded_2)):
        encoded_2[i]=encoded_2[i][8:16]
    encoded_1_bin = numpy.array(encoded_1,dtype=int)
    encoded_2_bin = numpy.array(encoded_2,dtype=int)
    _M = op.operation.reconstructPartition(encoded_1_bin,encoded_2_bin,Partisi)
    _M_int2= numpy.asarray(op.operation.binaryTonum(_M),dtype=numpy.uint16)

    if method == "GDE" :
        return(_M_int2,locMap_1,locMap_2,Partisi)
    elif method == "GRDEI":
        return(_M_int2,locMap_1,locMap_2,reduceMap_1,reduceMap_2,Partisi)

def decode(file_path,segment_size,payload_size,method,Partition,locMap_1,locMap_2,reduceMap_1 = None , reduceMap_2 = None):
        file_name = file_path.split("\\")[len(file_path.split("\\"))-1]
        path = file_path.replace(file_name,'')
        print(" PROSES DECODING ")
        print(" METODE YANG DIGUNAKAN : " )
        medium_encoded = Helper.WavIO.open(file_path)
        print "bitrate: "
        print(medium_encoded.bitrate)
        samples_decode = op.operation.numToBinary(medium_encoded.samples)
        (_M1,_M2,P) = op.operation.intel_partition(samples_decode,0,Partition)
        _intM1 = op.operation.binaryTonum(_M1)
        _intM2 = op.operation.binaryTonum(_M2)
        if method == "GDE" :
            (decoded_M1,message1) = GDE.decode(_intM1,segment_size,locMap_1)
            (decoded_M2,message2) = GDE.decode(_intM2,segment_size,locMap_2)
        elif method == "GRDEI":
            (decoded_M1,message1) = GRDEI.decode(_intM1,segment_size,locMap_1,reduceMap_1)
            (decoded_M2,message2) = GRDEI.decode(_intM2,segment_size,locMap_2,reduceMap_2)

        message_decoded = []
        message_decoded.extend(message1)
        message_decoded.extend(message2)
        message_decoded = message_decoded[:payload_size]
        print(len(message_decoded))
        message_write = op.operation.revStringToBinary(message_decoded)
        Helper.payloadIO.write(path+"payload_decoded.txt",message_write)
        decoded_1 = op.operation.numToBinary(decoded_M1)
        decoded_2 = op.operation.numToBinary(decoded_M2)
        for i in range(len(decoded_1)):
            decoded_1[i]=decoded_1[i][8:16]
        for i in range(len(decoded_2)):
            decoded_2[i]=decoded_2[i][8:16]
        decoded_1_bin = numpy.array(decoded_1,dtype=int)
        decoded_2_bin = numpy.array(decoded_2,dtype=int)
        M_awal = op.operation.reconstructPartition(decoded_1_bin,decoded_2_bin,Partition)
        M__awal = numpy.asarray(op.operation.binaryTonum(M_awal),dtype=numpy.int16)

        new_wav2 = Wave.Wave(file_name.replace("encoded","decoded"))
        new_wav2.samples = M__awal
        new_wav2.bitrate = medium_encoded.bitrate
        plt.clf()
        time_axis = numpy.linspace(0,len(medium_encoded.samples)/medium_encoded.bitrate,num=len(medium_encoded.samples))
        plt.subplot(2,1,1)
        plt.title("perbandingan hasil encode  dan hasil decode")
        plt.plot(time_axis,numpy.array(medium_encoded.samples,dtype= "int16"))
        plt.ylabel("WAV encoded")
        plt.subplot(2,1,2)
        plt.plot(time_axis,numpy.array(M__awal,dtype= "int16"))
        plt.ylabel("WAV Hasil Dencode")
        plt.xlabel("Waktu (s)")
        plt.savefig("encoded-decoded.png")
        Helper.WavIO.write(path,new_wav2)

def multilayer_encode(payload_path,cover_path,threshold,segment_size,partition_segment_size,method,n_layer):
    locmap_list = []
    reducemap_list = []
    capacities = []
    payload_sizes = []
    total_capacity = 0
    file_name = cover_path.split("\\")[len(cover_path.split("\\"))-1]
    path = cover_path.replace(file_name,'')
    print(" PROSES ENCODING ")
    print(" METODE YANG DIGUNAKAN : " + str(method))
    payload = Helper.payloadIO.open(payload_path)
    bin_payload = op.operation.stringToBinary(payload)
    print "besar payload : "
    payload_size = len(bin_payload)
    print(payload_size)
    medium = Helper.WavIO.open(cover_path)
    print "bitrate: "
    print(medium.bitrate)
    audio_sample = medium.samples
    #for i in range(n_layer):
    #    samples = op.operation.numToBinary(audio_sample)
    #    (M1,M2,Partisi) = op.operation.intel_partition(samples,partition_segment_size)
    #    intM1 = op.operation.binaryTonum(M1)
    #    intM2 = op.operation.binaryTonum(M2)
    #    if method == "GDE" :
    #        kapasitas_M1 = GDE.checkCapacity(intM1,segment_size,threshold)
    #        kapasitas_M2 = GDE.checkCapacity(intM2,segment_size,threshold)
    #    elif method == "GRDEI":
    #        kapasitas_M1 = GRDEI.checkCapacity(intM1,segment_size,threshold)
    #        kapasitas_M2 = GRDEI.checkCapacity(intM2,segment_size,threshold)
    #    print(" kapasitas segmen 1 : " + str(kapasitas_M1))
    #    print(" kapasitas segmen 2 : " + str(kapasitas_M2))
    #    capacity = kapasitas_M1+kapasitas_M2
    #    capacities.append((kapasitas_M1,kapasitas_M2))
    #    print "Kapasitas penyimpanan layer ke "+str(i)+":"+str(capacity)
    #    payload_seg1 = [1 for i in range(kapasitas_M1)]
    #    payload_seg2 = [1 for i in range(kapasitas_M2)]
    #    total_capacity+=capacity
    #    audio_sample = encode(intM1,intM2,payload_seg1,payload_seg2,threshold,segment_size,partition_segment_size,Partisi,method)[0]
    #print("total kapasitas "+str(total_capacity))
    #if(total_capacity > payload_size):
    #    print("stegano dapat dilakukan")

    #time_axis = numpy.linspace(0,len(medium.samples)/medium.bitrate,num=len(medium.samples))
    #plt.subplot(n_layer+1,1,1)
    #plt.title("perbandingan wav asli dan hasil encode multi-layer")
    #plt.plot(time_axis,numpy.array(medium.samples,dtype= "int16"))
    #plt.ylabel("WAV asli")

    P = op.operation.intel_partition(op.operation.numToBinary(audio_sample),partition_segment_size)[2]
    payload_counter = 0
    for i in range(n_layer):
        print("layer ke " + str(i))
        samples = op.operation.numToBinary(audio_sample)
        (M1,M2,Partisi) = op.operation.intel_partition(samples,partition_segment_size,P)
        intM1 = op.operation.binaryTonum(M1)
        intM2 = op.operation.binaryTonum(M2)
        if method == "GDE" :
            kapasitas_M1 = GDE.checkCapacity(intM1,segment_size,threshold)
            kapasitas_M2 = GDE.checkCapacity(intM2,segment_size,threshold)
        elif method == "GRDEI":
            kapasitas_M1 = GRDEI.checkCapacity(intM1,segment_size,threshold)
            kapasitas_M2 = GRDEI.checkCapacity(intM2,segment_size,threshold)
        capacities.append((kapasitas_M1,kapasitas_M2))
        kapasitas_total_layer = kapasitas_M1+kapasitas_M2
        total_capacity += kapasitas_total_layer
        if((payload_size - payload_counter) > kapasitas_total_layer):
            payload_i = bin_payload[payload_counter:kapasitas_total_layer]
            payload_sizes.append(kapasitas_total_layer)
            payload_counter+=kapasitas_total_layer
            print(len(payload_i))
        else:
            payload_i = bin_payload[payload_counter:payload_size]
            payload_sizes.append((payload_size-payload_counter))
            payload_counter+=(payload_size-payload_counter)
            print(len(payload_i))

        payload_seg1 = payload_i[:kapasitas_M1]
        print(len(payload_seg1))
        payload_seg2 = payload_i[kapasitas_M1:len(payload_i)]
        print(len(payload_seg2))
        if method == "GDE" :
            (audio_sample,locMap_1,locMap_2,Partisi)= encode(intM1,intM2,payload_seg1,payload_seg2,threshold,segment_size,partition_segment_size,P,method)
            locmap_list.append((locMap_1,locMap_2))
        elif method == "GRDEI":
            (audio_sample,locMap_1,locMap_2,reduceMap_1,reduceMap_2,Partisi) =  encode(intM1,intM2,payload_seg1,payload_seg2,threshold,segment_size,partition_segment_size,P,method)
            locmap_list.append((locMap_1,locMap_2))
            reducemap_list.append((reduceMap_1,reduceMap_2))
            #audio_sample =  numpy.asarray(audio_sample2,dtype="uint16")
        #plt.subplot(n_layer+1,1,i+2)
        #plt.plot(time_axis,numpy.array(audio_sample,dtype= "int16"))
        #plt.ylabel("Hasil Encode layer ke "+str(i+1))
        #plt.xlabel("Waktu (s)")
        #plt.savefig("original-multiencode-"+str(n_layer)+".png")
    if(payload_counter>=payload_size):
        print("stegano berhasil dilakukan")
        print(total_capacity)
        _M_int = numpy.asarray(audio_sample,dtype=numpy.int16)
        new_wav = Wave.Wave(method+"_encoded_"+str(file_name))
        new_wav.samples = _M_int
        new_wav.bitrate = medium.bitrate
        Helper.WavIO.write(path,new_wav)
        if method == "GDE" :
            return(locmap_list,payload_sizes,P)
        elif method == "GRDEI":
            return(locmap_list,reducemap_list,payload_sizes,P)
    else:
        print("kapasitas tidak mencukupi")

#else:
#    print "kapasitas tidak mencukupi"
#    return -1

def multilayer_decode(file_path,segment_size,payload_sizes,Partition,method,n_layer,locmap_list,reducemap_list = None,):
        full_messages=[]
        file_name = file_path.split("\\")[len(file_path.split("\\"))-1]
        path = file_path.replace(file_name,'')
        print(" PROSES DECODING ")
        print(" METODE YANG DIGUNAKAN : " )
        print(method)
        medium_encoded = Helper.WavIO.open(file_path)
        print "bitrate: "
        print(medium_encoded.bitrate)
        audio_samples = medium_encoded.samples
        for i in range (n_layer):
            print("layer ke " + str(i))
            (locmap_i_1,locmap_i_2) = locmap_list.pop()
            if(method == "GRDEI"):
                (reducemap_i_1,reducemap_i_2) = reducemap_list.pop()
            samples = op.operation.numToBinary(audio_samples)
            (_M1,_M2,P) = op.operation.intel_partition(samples,0,Partition)
            _intM1 = op.operation.binaryTonum(_M1)
            _intM2 = op.operation.binaryTonum(_M2)
            if method == "GDE" :
                (decoded_M1,message1) = GDE.decode(_intM1,segment_size,locmap_i_1)
                (decoded_M2,message2) = GDE.decode(_intM2,segment_size,locmap_i_2)
            elif method == "GRDEI":
                (decoded_M1,message1) = GRDEI.decode(_intM1,segment_size,locmap_i_1,reducemap_i_1)
                (decoded_M2,message2) = GRDEI.decode(_intM2,segment_size,locmap_i_2,reducemap_i_2)
            message_decoded = []
            message_decoded.extend(message1)
            message_decoded.extend(message2)
            payload_i_size = payload_sizes.pop()
            full_messages.insert(0,message_decoded[0:payload_i_size])
            decoded_1 = op.operation.numToBinary(decoded_M1)
            decoded_2 = op.operation.numToBinary(decoded_M2)
            for i in range(len(decoded_1)):
                decoded_1[i]=decoded_1[i][8:16]
            for i in range(len(decoded_2)):
                decoded_2[i]=decoded_2[i][8:16]
            decoded_1_bin = numpy.array(decoded_1,dtype=int)
            decoded_2_bin = numpy.array(decoded_2,dtype=int)
            M_awal = op.operation.reconstructPartition(decoded_1_bin,decoded_2_bin,Partition)
            audio_samples = numpy.asarray(op.operation.binaryTonum(M_awal),dtype=numpy.uint16)
            for i in audio_samples:
                if(i<0):
                    print(i)
        message_write = []
        for i in range(len(full_messages)):
            message_write.extend(full_messages[i])
        message_write = op.operation.revStringToBinary(message_write)
        Helper.payloadIO.write(path+"payload_decoded.txt",message_write)
        new_wav = Wave.Wave(file_name.replace("encoded","decoded"))
        M__awal = numpy.asarray(audio_samples,dtype=numpy.int16)
        new_wav.samples = M__awal
        new_wav.bitrate = medium_encoded.bitrate
        Helper.WavIO.write(path,new_wav)

if __name__ == "__main__":
    #(map1,map2,rmap1,rmap2,p) = encode("D:\\payload.txt","D:\\coba16.wav",100,20,10,"GRDEI")
    (locmap_list,reducemap_list,payload_sizes,partisi) = multilayer_encode("D:\\payload.txt","D:\\coba16.wav",50,20,10,"GRDEI",20)
    #print(len(locmap_list))
   # print(len(reducemap_list))
    multilayer_decode("D:\\GRDEI_encoded_coba16.wav",20,payload_sizes,partisi,"GRDEI",20,locmap_list,reducemap_list)
    #locmap_list.pop()
    #(map1_1,map1_2) = locmap_list.pop()
    #reducemap_list.pop()
    #(rmap1_1,rmap1_2) = reducemap_list.pop()

    #decode("D:\\GRDEI_encoded2_coba16.wav",20,1187,"GRDEI",partisi,map1_1,map1_2,rmap1_1,rmap1_2)
     #print(len(locmap_list),len(reducemap_list))
#   # print "data payload : "
#    #print(payload)

#    #arr = [13954, 4369, 37385, 3995, 2556, 46896, 13816, 17865, 40433, 42503, 27740, 14980, 22323, 27920, 48381, 40456, 58866, 60412, 36991, 30730, 14601, 31475, 50583, 57144, 18332, 46140, 47181, 62996, 19071, 30753, 55953, 62831, 8814, 44566, 2191, 16703, 36414, 55831, 28696, 43850]
#    #samples = op.operation.numToBinary(arr)

#    besar_segmen = 2
#    threshold = 200

#        ########### DECODING ###############


"""
Equivalent layer for griding and upward-continuing gravity data
-------------------------------------------------------------------------

The equivalent layer is one of the best methods for griding and upward
continuing gravity data and much more. The trade-off is that performing this
requires an inversion and later forward modeling, which are more time consuming
and more difficult to tune than the standard griding and FFT-based approaches.

This example uses the equivalent layer in :mod:`fatiando.gravmag.eqlayer` to
grid and upward continue some gravity data. There are more advanced methods in
the module than the one we are showing here. They can be more efficient but
usually require more configuration.

"""
from __future__ import division, print_function
import matplotlib.pyplot as plt
from fatiando.gravmag import prism, sphere
from fatiando.gravmag.eqlayer import EQLGravity
from fatiando.inversion import Damping
from fatiando import gridder, utils, mesher

# First thing to do is make some synthetic data to test the method. We'll use a
# single prism to keep it simple
props = {'density': 500}
model = [mesher.Prism(-5000, 5000, -200, 200, 100, 4000, props)]

# The synthetic data will be generated on a random scatter of points
area = [-8000, 8000, -5000, 5000]
x, y, z = gridder.scatter(area, 300, z=0, seed=42)
# Generate some noisy data from our model
gz = utils.contaminate(prism.gz(x, y, z, model), 0.2, seed=0)

# Now for the equivalent layer. We must setup a layer of point masses where
# we'll estimate a density distribution that fits our synthetic data
layer = mesher.PointGrid(area, 500, (20, 20))
# Estimate the density using enough damping so that won't try to fit the error
eql = EQLGravity(x, y, z, gz, layer) + 1e-22*Damping(layer.size)
eql.fit()
# Now we add the estimated densities to our layer
layer.addprop('density', eql.estimate_)
# and print some statistics of how well the estimated layer fits the data
residuals = eql[0].residuals()
print("Residuals:")
print("  mean:", residuals.mean(), 'mGal')
print("  stddev:", residuals.std(), 'mGal')

# Now I can forward model gravity data anywhere we want. For interpolation, we
# calculate it on a grid. For upward continuation, at a greater height. We can
# even combine both into a single operation.
x2, y2, z2 = gridder.regular(area, (50, 50), z=-1000)
gz_up = sphere.gz(x2, y2, z2, layer)

fig, axes = plt.subplots(1, 2, figsize=(8, 6))

ax = axes[0]
ax.set_title('Original data')
ax.set_aspect('equal')
tmp = ax.tricontourf(y/1000, x/1000, gz, 30, cmap='viridis')
fig.colorbar(tmp, ax=ax, pad=0.1, aspect=30,
             orientation='horizontal').set_label('mGal')
ax.plot(y/1000, x/1000, 'xk')
ax.set_xlabel('y (km)')
ax.set_ylabel('x (km)')

ax = axes[1]
ax.set_title('Gridded and upward continued')
ax.set_aspect('equal')
tmp = ax.tricontourf(y2/1000, x2/1000, gz_up, 30, cmap='viridis')
fig.colorbar(tmp, ax=ax, pad=0.1, aspect=30,
             orientation='horizontal').set_label('mGal')
ax.set_xlabel('y (km)')

plt.tight_layout()
plt.show()


# vim: set fileencoding=utf-8 sw=2 ts=2 et :
from __future__ import absolute_import
from __future__ import with_statement

from logging import getLogger

import networkx as NX
import yaml

from systems.collector import Aggregate, CResource
from systems.registry import get_registry
from systems.typesystem import EResource, Transition, ResourceRef

__all__ = ('Realizer', )

LOGGER = getLogger(__name__)

DESC_LIMIT = 64

def describe(thing):
  return '%s' % str(thing)[:DESC_LIMIT]

class CycleError(Exception):
  pass

class Node(object):
  def __init__(self):
    if type(self) == Node:
      raise TypeError

  def __repr__(self):
    return '<%s>' % self

  def __str__(self):
    return type(self).__name__

class CheckPointNode(Node):
  pass

class ExpandableNode(Node):
  def __init__(self, res):
    super(ExpandableNode, self).__init__()
    if type(self) == ExpandableNode:
      # Abstract class
      raise TypeError
    self._res = res

class BeforeExpandableNode(ExpandableNode):
  def __str__(self):
    return 'Before %s' % self._res

class AfterExpandableNode(ExpandableNode):
  def __str__(self):
    return 'After %s' % self._res

class GraphFirstNode(Node, yaml.YAMLObject):
  yaml_tag = u'GraphFirstNode'

class GraphLastNode(Node, yaml.YAMLObject):
  yaml_tag = u'GraphLastNode'

node_types = (CheckPointNode, BeforeExpandableNode, AfterExpandableNode,
    GraphFirstNode, GraphLastNode,
    Transition, Aggregate, CResource, EResource, ResourceRef)

class ResourceGraph(yaml.YAMLObject):
  """
  A graph of resources and transitions linked by dependencies.

  Resources are positioned as two sentinels in the transition graph.

  Invariant: directed, acyclic.
  """

  def __init__(self, top=None):
    self._graph = NX.DiGraph()
    self._first = GraphFirstNode()
    self._last = GraphLastNode()
    self._graph.add_edge(self._first, self._last)
    # Contains CResource and EResource, despite the name.
    # Used to enforce max one resource per id.
    self.__expandables = {}
    # Received references, by name.
    self.__received_refs = {}
    # What nodes were processed (meaning expanding or collecting)
    self.__processed = set()
    # Pre-bound args pased by ref. Allow putting extra depends on them.
    if top is not None:
      if not isinstance(top, ResourceGraph):
        raise TypeError(top, ResourceGraph)
      self.__top = top
    else:
      self.__top = self

  yaml_tag = u'!ResourceGraph'

  @classmethod
  def from_yaml(cls, loader, ynode):
    rg = cls()
    # Deep because of aliases and anchors, I think.
    mp = loader.construct_mapping(ynode, deep=True)
    pred_rels = mp['nodes']
    for rel in pred_rels:
      rg._add_node(rel['node'], depends=rel['depends'])
    return rg

  @classmethod
  def to_yaml(cls, dumper, rg):
    # This is incomplete.
    pred_rels = [{'node': node, 'depends': list(depends), }
        for (node, depends) in rg._iter_pred_rels()]
    return dumper.represent_mapping(cls.yaml_tag, {
      'nodes': pred_rels,
      })

  def _iter_node_preds(self, node0):
    return (node
        for node in self._graph.predecessors_iter(node0)
        if node not in (self._first, self._last))

  def _iter_pred_rels(self):
    return ((node, self._iter_node_preds(node))
      for node in self.sorted_nodes()
      if node not in (self._first, self._last))

  def sorted_nodes(self):
    return NX.topological_sort(self._graph)

  def sorted_transitions(self):
    return [n for n in self.sorted_nodes()
        if isinstance(n, Transition)]

  def iter_uncollected_resources(self):
    for nod in self._graph.nodes_iter():
      if isinstance(nod, CResource):
        if not nod in self.__processed:
          yield nod

  def iter_unexpanded_resources(self):
    for nod in self._graph.nodes_iter():
      if isinstance(nod, EResource):
        if not nod in self.__processed:
          yield nod

  def iter_unexpanded_aggregates(self):
    for agg in self._graph.nodes_iter():
      if isinstance(agg, Aggregate):
        if not agg in self.__processed:
          yield agg

  def iter_unprocessed(self):
    for nod in self.iter_uncollected_resources():
      yield nod
    for nod in self.iter_unexpanded_resources():
      yield nod
    for nod in self.iter_unexpanded_aggregates():
      yield nod

  def has_unprocessed(self):
    l = list(self.iter_unprocessed())
    return bool(l) # Tests for non-emptiness

  def require_acyclic(self):
    if not NX.is_directed_acyclic_graph(self._graph):
      # XXX NX doesn't have a 1-line method for listing those cycles
      raise CycleError

  def _add_node(self, node, depends=()):
    if not isinstance(node, node_types):
      raise TypeError(node, node_types)
    self._graph.add_node(node)
    self._graph.add_edge(self._first, node)
    self._graph.add_edge(node, self._last)
    for dep in depends:
      depn = self._intern(dep)
      self._add_node_dep(depn, node)
    return node

  def add_checkpoint(self, depends=()):
    return self._add_node(CheckPointNode(), depends)

  def add_transition(self, transition, depends=()):
    if not isinstance(transition, Transition):
      raise TypeError(transition, Transition)
    return self._add_node(transition, depends)

  def _add_aggregate(self, aggregate, depends=()):
    if not isinstance(aggregate, Aggregate):
      raise TypeError(aggregate, Aggregate)
    return self._add_node(aggregate, depends)

  def add_resource(self, resource, depends=()):
    """
    Add a resource.

    If an identical resource exists, it is returned.
    """

    if not isinstance(resource, (CResource, EResource)):
      raise TypeError(resource, (CResource, EResource))

    if resource.identity in self.__expandables:
      # We have this id already.
      # Either it's the exact same resource, or a KeyError is thrown.
      resource = self._intern(resource)
      # XXX Need to bypass _intern for already expanded.
      # XXX When we use add_to_top, we sometimes have to deal
      # with a resource that's already been expanded.
      # Those are not in the graph anymore. How do we refer to them?
    else:
      self.__expandables[resource.identity] = resource
    # Even if already there, we need to add the depends.
    resource = self._add_node(resource, depends)
    # If already there, notice we aliase it.
    return self.make_ref(resource)

  def make_ref(self, res, depends=()):
    res = self._intern(res)
    if not isinstance(res, (CResource, EResource)):
      raise TypeError(res, (CResource, EResource))
    depends = list(depends)
    depends.append(res)
    return self._add_node(ResourceRef(res), depends)

  def make_alias_ref(self, ref, depends=()):
    ref = self._intern(ref)
    if not isinstance(ref, ResourceRef):
      raise TypeError(ref, ResourceRef)
    depends = list(depends)
    depends.append(ref)
    return self._add_node(ResourceRef(ref.unref), depends)

  def add_to_top(self, res):
    """
    Add a resource to the top ResourceGraph.

    Use it to put things that you don't necessarily
    want to be after the outside dependencies the current graph has.
    """

    ref = self.__top.add_resource(res)
    return self._add_node(ref)

  def _add_node_dep(self, node0, node1):
    if not isinstance(node0, node_types):
      raise TypeError(node0, node_types)
    if not isinstance(node1, node_types):
      raise TypeError(node1, node_types)
    if not self._graph.has_node(node0):
      raise KeyError(node0)
    if not self._graph.has_node(node1):
      raise KeyError(node1)
    if self._graph.has_edge(node0, node1):
      return False
    if node0 == node1:
      # Disallow self-loops to keep acyclic invariant.
      # Also they don't make sense.
      raise ValueError(node0)
    # Invariant check
    rev_path = NX.shortest_path(self._graph, node1, node0)
    if rev_path is not False:
      raise CycleError(rev_path)
    self._graph.add_edge(node0, node1)
    return True

  def _intern(self, thing):
    if not isinstance(thing, node_types):
      raise TypeError
    if thing not in self._graph:
      raise KeyError(thing)
    return thing

  def add_dependency(self, elem0, elem1):
    node0 = self._intern(elem0)
    node1 = self._intern(elem1)
    return self._add_node_dep(node0, node1)

  def _is_direct_rconnect(self, r0, r1):
    s0 = self._intern(r0)
    s1 = self._intern(r1)
    # shortest_path is also a test for connectedness.
    return bool(NX.shortest_path(self._graph, s0, s1))

  def resources_connected(self, r0, r1):
    return self._is_direct_rconnect(r0, r1) \
        or self._is_direct_rconnect(r1, r0)

  def draw(self, fname):
    return self.draw_agraph(fname)

  def draw_agraph(self, fname):
    # XXX pygraphviz has steep dependencies (x11 libs)
    # and recommends (texlive) for a headless box.

    # We duplicate the graph, otherwise networkx / pygraphviz
    # would make a lossy conversion (sometimes refusing to convert), by adding
    # nodes as their string representation. Madness, I know.
    gr2 = NX.create_empty_copy(self._graph, False)
    for node in self._graph.nodes_iter():
      gr2.add_node(id(node))
    for (n0, n1) in self._graph.edges_iter():
      gr2.add_edge(id(n0), id(n1))
    names = dict((id(node), { 'label': describe(node)})
        for node in self._graph.nodes_iter())
    gr2.delete_node(id(self._first))
    gr2.delete_node(id(self._last))
    g = NX.to_agraph(gr2, {
        'graph': {
          'nodesep': '0.2',
          'rankdir': 'TB',
          'ranksep': '0.5',
          },
        'node': {
          'shape': 'box',
          },
        },
        names)
    g.write(fname + '.dot')
    # Dot is good for DAGs.
    g.layout(prog='dot')
    g.draw(fname + '.svg')
    with open(fname + '.yaml', 'w') as f:
      yaml.dump(self, f)
    # Fails with the expanded graph, due to instancemethod
    #yaml.load(yaml.dump(self))

  def draw_matplotlib(self, fname):
    # Pyplot is stateful and awkward to use.
    import matplotlib.pyplot as P
    # Disable hold or it definitely won't work (probably a bug).
    P.hold(False)
    NX.draw(self._graph)
    P.savefig(fname)

  def collect_resources(self, r0s, r1):
    """
    Replace an iterable of resources with one new resource.

    May break the acyclic invariant, caveat emptor.
    """

    # The invariant is kept iff the r0s don't have paths linking them.
    # For our use case (collectors), we could allow paths provided they are
    # internal to r0s. This introduces self-loops that we would then remove.

    for r0 in r0s:
      r0 = self._intern(r0)
      if r0 in self.__processed:
        raise RuntimeError

    if r1 in self._graph:
      raise ValueError(r1)
    r1 = self._add_aggregate(r1)

    for r0 in r0s:
      r0 = self._intern(r0)
      self._move_edges(r0, r1)
      self.__processed.add(r0)
    self.require_acyclic()

  def _move_edges(self, n0, n1):
    if n0 == n1:
      raise RuntimeError
    n0 = self._intern(n0)
    n1 = self._intern(n1)
    # list is used as a temporary
    # add after delete in case of same.
    for pred in list(self._graph.predecessors_iter(n0)):
      self._graph.delete_edge(pred, n0)
      self._graph.add_edge(pred, n1)
    for succ in list(self._graph.successors_iter(n0)):
      self._graph.delete_edge(n0, succ)
      self._graph.add_edge(n1, succ)
    self._graph.delete_node(n0)
    # Can't undo. Invariant will stay broken.

  def _split_node(self, res):
    res = self._intern(res)
    before = self._add_node(BeforeExpandableNode(res))
    after = self._add_node(AfterExpandableNode(res))
    self._graph.add_edge(before, after)
    for pred in list(self._graph.predecessors_iter(res)):
      self._graph.delete_edge(pred, res)
      self._graph.add_edge(pred, before)
    for succ in list(self._graph.successors_iter(res)):
      self._graph.delete_edge(res, succ)
      self._graph.add_edge(after, succ)
    self._graph.delete_node(res)
    return before, after

  def _receive_by_ref(self, name, ref):
    if name in self.__received_refs:
      raise RuntimeError(name, ref)
    ref = self._add_node(ref)
    self.__received_refs[name] = ref
    return ref

  def _pass_by_ref(self, subgraph, name, ref):
    # The origin/value distinction is important
    # for aliased arguments (two refs, same val).
    ref = self._intern(ref)
    if not isinstance(ref, ResourceRef):
      raise TypeError(ref, ResourceRef)

    subgraph._receive_by_ref(name, ref)

  def expand_resource(self, res):
    """
    Replace res by a small resource graph.

    The resource_graph is inserted in the main graph
    between the sentinels that represent the resource.
    """

    res = self._intern(res)

    # We're processing from the outside in.
    if res in self.__processed:
      raise RuntimeError

    resource_graph = ResourceGraph(self.__top)
    if isinstance(res, EResource):
      for (name, ref) in res.iter_passed_by_ref():
        # ref will be present in both graphs.
        self._pass_by_ref(resource_graph, name, ref)
    elif isinstance(res, Aggregate):
      pass
    else:
      raise TypeError(res)
    res.expand_into(resource_graph)
    # We expand from the outside in
    if bool(resource_graph.__processed):
      raise RuntimeError

    # Do not skip sentinels.
    for n in resource_graph._graph.nodes_iter():
      self._add_node(n)
    for (n0, n1) in resource_graph._graph.edges_iter():
      self._add_node_dep(n0, n1)

    for (id1, res1) in resource_graph.__expandables.iteritems():
      # We expand from the outside in.
      assert res1 not in self.__processed
      if id1 in self.__expandables:
        # Pass by reference if you must use the same resource
        # in different contexts.
        raise RuntimeError('ResourceBase collision.', res, res1)
      else:
        self.__expandables[id1] = res1

    before, after = self._split_node(res)
    self.__processed.add(res)
    self._move_edges(resource_graph._first, before)
    self._move_edges(resource_graph._last, after)
    # What may break the invariant:
    # Passing a ref to res, and making res depend on ref.
    # ref ends up on both sides of ref.before.
    self.require_acyclic()

class Realizer(object):
  """
  A graph of realizables linked by dependencies.
  """

  def __init__(self, expandable):
    self.__resources = ResourceGraph()
    self.__expandable = expandable
    self.__state = 'init'

  def require_state(self, state):
    """
    Raise an exception if we are not in the required state.
    """

    if self.__state != state:
      raise RuntimeError(u'Realizer state should be «%s»' % state)

  def ensure_frozen(self):
    """
    Build the finished dependency graph.

    Merge identical realizables, collect what can be.
    """

    if self.__state == 'frozen':
      return
    # Order is important
    self.require_state('init')
    self.__expandable.expand_into(self.__resources)
    #self.__resources.draw('/tmp/freezing')
    self._expand()
    #self.__resources.draw('/tmp/pre-collect')
    self._collect()
    self._expand_aggregates()
    assert not bool(list(self.__resources.iter_unprocessed()))
    self.__state = 'frozen'
    #self.__resources.draw('/tmp/frozen')

  def _collect(self):
    # Collects compatible nodes into merged nodes.

    def can_merge(part0, part1):
      for n0 in part0:
        for n1 in part1:
          if self.__resources.resources_connected(n0, n1):
            return False
      return True

    def possibly_merge(partition):
      # Merge once if possible. Return true if did merge.
      e = dict(enumerate(partition))
      n = len(partition)
      # Loop over the triangle of unordered pairs
      for i in xrange(n):
        for j in xrange(i + 1, n):
          part0, part1 = e[i], e[j]
          if can_merge(part0, part1):
            partition.add(part0.union(part1))
            partition.remove(part0)
            partition.remove(part1)
            return True
      return False

    reg = get_registry()
    for collector in reg.collectors:
      # Pre-partition is made of parts acceptable for the collector.
      pre_partition = collector.partition(
          [r for r in self.__resources.iter_uncollected_resources()
            if collector.filter(r)])
      for part in pre_partition:
        # Collector parts are split again, the sub-parts are merged
        # when dependencies allow.
        # Not a particularly efficient algorithm, just simple.
        # Gives one solution among many possibilities.
        partition = set(frozenset((r, ))
            for r in part
            for part in pre_partition)
        while possibly_merge(partition):
          pass

        # Let the collector handle the rest
        for part in partition:
          if not bool(part):
            # Test for emptiness.
            # Aggregate even singletons.
            continue
          merged = collector.collect(part)
          self.__resources.collect_resources(part, merged)
    assert not bool(list(self.__resources.iter_uncollected_resources()))

  def _expand(self):
    # Poor man's recursion
    while True:
      fresh = set(r
          for r in self.__resources.iter_unexpanded_resources())
      if bool(fresh) == False: # Test for emptiness
        break
      for r in fresh:
        self.__resources.expand_resource(r)
    assert not bool(list(self.__resources.iter_unexpanded_resources()))

  def _expand_aggregates(self):
    for a in list(self.__resources.iter_unexpanded_aggregates()):
      self.__resources.expand_resource(a)
    assert not bool(list(self.__resources.iter_unexpanded_aggregates()))
    # Enforce the rule that aggregates can only expand into transitions.
    if self.__resources.has_unprocessed():
      raise RuntimeError(list(self.__resources.iter_unprocessed()))

  def realize(self):
    """
    Realize all realizables and transitions in dependency order.
    """

    self.ensure_frozen()
    for t in self.__resources.sorted_transitions():
      t.realize()
    self.__state = 'realized'