import os
import io
import sys
import dill
import copy
from datetime import datetime

from .evaluator import Evaluator
from .utils import (
    post_to_platform,
    get_current_notebook,
    strip_output,
    get_current_notebook,
    mkdir_p,
)

class DataScienceFramework(object):
    def __init__(
        self,
        model,
        problem_class,
        x_test,
        y_test,
        name=None,
        description=None,
        evaluator=Evaluator,
    ):
        # assign variables to class
        self.name = name
        self.description = description
        self.model = model
        self.problem_class = problem_class
        self.y_test = list(y_test)
        self.x_test = list(x_test)
        self.framework = model.__module__.split(".")[0]

        # get environment data
        self._meta_data = self.meta_data()

        self.y_pred = self.predict()
        # initialize evaluator
        self.evaluator = Evaluator(self.problem_class)

    # class methods
    @classmethod
    def load(cls, model_id):
        # use hard coded string to load for now
        with open(".model_cache/sklearn_model_cache.pkl", "rb") as file:
            instance = dill.load(file)
            instance.model = instance.parse_model(io.BytesIO(instance.model_serialized))
        return instance

    @classmethod
    def project_models(cls):
        query = """
            query($service_name: String!) {
                runnableInstance(serviceName: $service_name) {
                    runnable {
                        project {
                            name
                            models {
                                edges {
                                    node {
                                        id
                                        name
                                        description
                                        problemClass
                                        framework
                                        objectClass
                                        language
                                        languageVersion
                                        createdAt
                                        updatedAt
                                        rank
                                        hyperParameters
                                        structure
                                        author {
                                            fullName
                                        }
                                        metrics {
                                            edges {
                                                node {
                                                    key
                                                    value
                                                }
                                            }
                                        }
                                        diagnostics {
                                            edges {
                                                node {
                                                    ... on ModelDiagnosticROC {
                                                    title
                                                    falsePositiveRates
                                                    truePositiveRates
                                                    thresholds
                                                    }
                                                    ... on ModelDiagnosticResidual {
                                                    title
                                                    observations
                                                    residuals
                                                    }
                                                    ... on ModelDiagnosticConfusionMatrix {
                                                    title
                                                    matrix
                                                    }
                                                }
                                            }
                                        }
                                        parameters {
                                            edges {
                                                node {
                                                    key
                                                    value
                                                    confidenceInterval {
                                                    positive
                                                    negative
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
            """

        response = post_to_platform(
            {"query": query, "variables": {"service_name": os.environ["SERVICE_NAME"]}}
        )
        response_data = response.json()["data"]
        models = list(
            map(
                lambda edge: edge["node"],
                response_data["runnableInstance"]["runnable"]["project"]["models"][
                    "edges"
                ],
            )
        )
        return models

    # framework dependent functions
    def predict(self):
        """ Make prediction based on x_test """
        raise NotImplementedError

    def framework_version(self):
        """ Return version of the framework been used. """
        raise NotImplementedError

    def object_class(self):
        """ Return name of the model object.  """
        raise NotImplementedError

    def parameter(self):
        """ Get parameter from model. """
        raise NotImplementedError

    def hyperparameter(self):
        """ Get hyper parameter from model. """
        raise NotImplementedError

    def serialize_model(self):
        """ Default methods for serialize model. """
        return dill.dumps(self.model)

    def parse_model(self, model_file):
        """ Default methods for reading in model. """
        return dill.load(model_file)

    # base framework functions
    def meta_data(self):
        """ Capture environment meta data. """
        meta_data_obj = {
            "name": self.name,
            "description": self.description,
            "framework": self.framework,
            "createdAt": datetime.now().isoformat(),
            "sessionName": os.environ["SERVICE_NAME"],
            "language": "python",
            "languageVersion": ".".join(map(str, sys.version_info[0:3])),
        }
        return meta_data_obj

    def diagnostics(self):
        """ Return diagnostics of model. """
        return [fn(self.y_test, self.y_pred) for fn in self.evaluator.diagnostics]

    def metrics(self):
        """ Return evaluation of model performance. """
        return [fn(self.y_test, self.y_pred) for fn in self.evaluator.metrics]

    def summary(self):
        """ Return all infomation that will be stored. """
        model_meta = {
            "diagnostics": self.diagnostics(),
            "metrics": self.metrics(),
            "parameters": self.parameter(),
            "frameworkVersion": self.framework_version(),
            "hyperParameters": self.hyperparameter(),
            "problemClass": self.problem_class,
            "objectClass": self.object_class(),
        }

        model_meta.update(self._meta_data)
        return model_meta

    def save(self):
        """ Save all information to platform. """
        self.model_serialized = self.serialize_model()

        # save model object locally for now
        #mkdir_p(".model_cache")
        #with open(".model_cache/sklearn_model_cache.pkl", "w") as file:
        #   dill.dump(self, file)

        model_meta = self.summary()

        model_meta.update(
            {
                "data": {"y_pred": list(self.y_pred), "y_test": list(self.y_test)},
                "notebook": get_current_notebook(),
            }
        )

        query = """
            mutation($input: CreateModelInput!) {
                createModel(input: $input) {
                    clientMutationId
                }
            }
            """

        return post_to_platform({"query": query, "variables": {"input": model_meta}})


#!/usr/bin/python

import sys, string, re,  os, commands, time, math

#from scipy import stats

#import scipy as sp

import numpy as np

#import matplotlib as mpl

#from matplotlib import pyplot as plt

class SNP:

    def __init__(self, name, num_env, t):

        self.name = name
        self.num_env = [False] * num_env
        self.bf_list = [[0 for i in range(t)] for j in range(num_env)]
        self.rel_signal = []
        self.sum_signals = 0
        self.lg_info = []
        self.chr = 99
        self.lg = 99

    def get_name(self):
        return self.name

    def get_num_env(self):
        return self.num_env

    def set_num_env(self, n):
        self.num_env[n] = True

    def add_to_list(self, bf, k, i):
        self.bf_list[k][i] = bf

    def set_signal(self, gamma):
        self.rel_signal.append(gamma)
        self.sum_signals += gamma #Add to the total of signals

    #Return the bf signal in variable k
    def get_signal(self, k):
        return self.rel_signal[k]

    #Return the bf signal list
    def get_signals(self):
        return self.rel_signal

    def get_sum_signals(self):
        return self.sum_signals

    def print_env(self):
        print self.num_env

    def get_median_bf(self, k):
        #print self.bf_list[k]
        bfs = np.array(self.bf_list[k])
        median = np.median(bfs)
        return median

    def get_avg_bf(self, k):
        #print self.bf_list[k]
        bfs = np.array(self.bf_list[k])
        avg = np.average(bfs)
        return avg

    def add_bf(self, bf):
        self.sum_bf += bf

    def get_sum_bf(self):
        return self.sum_bf

    def get_num_runs(self):
        return self.num_runs

    def get_bf_list(self):
        return self.bf_list

    def get_bf_list(self):
        return self.bf_list

    def set_lg_info(self, info):
        self.lg_info.append(info)

    def get_lg_info(self):
        return self.lg_info

    def set_chr(self, ch):
        self.chr = ch

    def get_chr(self):
        return self.chr

    def set_linkage_group(self, lg):
        self.lg = lg

    def get_linkage_group(self):
        return self.lg

def compute_average_bf(num_var, num_tests):

    N = int(num_var)
    t = int(num_tests)

    snp_dict = {}

    for i in range (0, t):
        filename = "results/bf_results_t" + str(i) + ".bf"
        data = open( filename, "r")
        print filename
        lines = data.readlines()
        for line in lines:
            cols = line.split("\t")
            snp_name = cols[0][0:-2]
            if i > 9:
                snp_name = snp_name[0:-1]
            if snp_name in snp_dict:
                snp = snp_dict[snp_name]
                for k in range(0, N):
                    snp.add_to_list(float(cols[k+1]), k, i)
            else:
                snp = SNP(snp_name, N, t)
                snp_dict[snp_name] = snp
                for k in range(0, N):
                    snp.add_to_list(float(cols[k+1]), k, i)

        data.close()

    print "################LENGTH:" + str(len(snp_dict))

    FILE1 = open("results/median_bf.txt", "w")
    FILE2 = open("results/average_bf.txt", "w")

    #bf_median = "marker\tsal1\tsal2\ttemp1\ttemp2\tox1\tox2\n"
    #bf_avg = "marker\tsal1\tsal2\ttemp1\ttemp2\tox1\tox2\n"
    bf_median = ""
    bf_avg = ""
    for key in snp_dict:
        snp = snp_dict[key]
        bf_avg += snp.get_name()
        bf_median += snp.get_name()
        for k in range(0, N):
            bf_a = snp.get_avg_bf(k)
            bf_m = snp.get_median_bf(k)
            bf_avg += "\t" + str(bf_a)
            bf_median += "\t" + str(bf_m)
        bf_avg += "\n"
        bf_median += "\n"
    FILE1.write(bf_median)
    FILE2.write(bf_avg)
    FILE1.close()
    FILE2.close()

if __name__ == '__main__':
    # Terminate if too few arguments
    if len(sys.argv) < 3:
        print 'usage: %s <number of vars> <num tests>' % sys.argv[0]
        sys.exit(-1)
    main(sys.argv[1], sys.argv[2])


"""
Several basic tests for hierarchical clustering procedures

"""
# Authors: Vincent Michel, 2010, Gael Varoquaux 2012,
#          Matteo Visconti di Oleggio Castello 2014
# License: BSD 3 clause
from tempfile import mkdtemp
import shutil
from functools import partial

import numpy as np
from scipy import sparse
from scipy.cluster import hierarchy

from sklearn.utils.testing import assert_true
from sklearn.utils.testing import assert_raises
from sklearn.utils.testing import assert_equal
from sklearn.utils.testing import assert_almost_equal
from sklearn.utils.testing import assert_array_almost_equal
from sklearn.utils.testing import assert_raise_message
from sklearn.utils.testing import ignore_warnings

from sklearn.cluster import ward_tree
from sklearn.cluster import AgglomerativeClustering, FeatureAgglomeration
from sklearn.cluster.hierarchical import (_hc_cut, _TREE_BUILDERS,
                                          linkage_tree)
from sklearn.feature_extraction.image import grid_to_graph
from sklearn.metrics.pairwise import PAIRED_DISTANCES, cosine_distances,\
    manhattan_distances, pairwise_distances
from sklearn.metrics.cluster import normalized_mutual_info_score
from sklearn.neighbors.graph import kneighbors_graph
from sklearn.cluster._hierarchical import average_merge, max_merge
from sklearn.utils.fast_dict import IntFloatDict
from sklearn.utils.testing import assert_array_equal
from sklearn.utils.testing import assert_warns

def test_linkage_misc():
    # Misc tests on linkage
    rng = np.random.RandomState(42)
    X = rng.normal(size=(5, 5))
    assert_raises(ValueError, AgglomerativeClustering(linkage='foo').fit, X)
    assert_raises(ValueError, linkage_tree, X, linkage='foo')
    assert_raises(ValueError, linkage_tree, X, connectivity=np.ones((4, 4)))

    # Smoke test FeatureAgglomeration
    FeatureAgglomeration().fit(X)

    # test hiearchical clustering on a precomputed distances matrix
    dis = cosine_distances(X)

    res = linkage_tree(dis, affinity="precomputed")
    assert_array_equal(res[0], linkage_tree(X, affinity="cosine")[0])

    # test hiearchical clustering on a precomputed distances matrix
    res = linkage_tree(X, affinity=manhattan_distances)
    assert_array_equal(res[0], linkage_tree(X, affinity="manhattan")[0])

def test_structured_linkage_tree():
    # Check that we obtain the correct solution for structured linkage trees.
    rng = np.random.RandomState(0)
    mask = np.ones([10, 10], dtype=np.bool)
    # Avoiding a mask with only 'True' entries
    mask[4:7, 4:7] = 0
    X = rng.randn(50, 100)
    connectivity = grid_to_graph(*mask.shape)
    for tree_builder in _TREE_BUILDERS.values():
        children, n_components, n_leaves, parent = \
            tree_builder(X.T, connectivity)
        n_nodes = 2 * X.shape[1] - 1
        assert_true(len(children) + n_leaves == n_nodes)
        # Check that ward_tree raises a ValueError with a connectivity matrix
        # of the wrong shape
        assert_raises(ValueError,
                      tree_builder, X.T, np.ones((4, 4)))
        # Check that fitting with no samples raises an error
        assert_raises(ValueError,
                      tree_builder, X.T[:0], connectivity)

def test_unstructured_linkage_tree():
    # Check that we obtain the correct solution for unstructured linkage trees.
    rng = np.random.RandomState(0)
    X = rng.randn(50, 100)
    for this_X in (X, X[0]):
        # With specified a number of clusters just for the sake of
        # raising a warning and testing the warning code
        with ignore_warnings():
            children, n_nodes, n_leaves, parent = assert_warns(
                UserWarning, ward_tree, this_X.T, n_clusters=10)
        n_nodes = 2 * X.shape[1] - 1
        assert_equal(len(children) + n_leaves, n_nodes)

    for tree_builder in _TREE_BUILDERS.values():
        for this_X in (X, X[0]):
            with ignore_warnings():
                children, n_nodes, n_leaves, parent = assert_warns(
                    UserWarning, tree_builder, this_X.T, n_clusters=10)

            n_nodes = 2 * X.shape[1] - 1
            assert_equal(len(children) + n_leaves, n_nodes)

def test_height_linkage_tree():
    # Check that the height of the results of linkage tree is sorted.
    rng = np.random.RandomState(0)
    mask = np.ones([10, 10], dtype=np.bool)
    X = rng.randn(50, 100)
    connectivity = grid_to_graph(*mask.shape)
    for linkage_func in _TREE_BUILDERS.values():
        children, n_nodes, n_leaves, parent = linkage_func(X.T, connectivity)
        n_nodes = 2 * X.shape[1] - 1
        assert_true(len(children) + n_leaves == n_nodes)

def test_agglomerative_clustering():
    # Check that we obtain the correct number of clusters with
    # agglomerative clustering.
    rng = np.random.RandomState(0)
    mask = np.ones([10, 10], dtype=np.bool)
    n_samples = 100
    X = rng.randn(n_samples, 50)
    connectivity = grid_to_graph(*mask.shape)
    for linkage in ("ward", "complete", "average"):
        clustering = AgglomerativeClustering(n_clusters=10,
                                             connectivity=connectivity,
                                             linkage=linkage)
        clustering.fit(X)
        # test caching
        try:
            tempdir = mkdtemp()
            clustering = AgglomerativeClustering(
                n_clusters=10, connectivity=connectivity,
                memory=tempdir,
                linkage=linkage)
            clustering.fit(X)
            labels = clustering.labels_
            assert_true(np.size(np.unique(labels)) == 10)
        finally:
            shutil.rmtree(tempdir)
        # Turn caching off now
        clustering = AgglomerativeClustering(
            n_clusters=10, connectivity=connectivity, linkage=linkage)
        # Check that we obtain the same solution with early-stopping of the
        # tree building
        clustering.compute_full_tree = False
        clustering.fit(X)
        assert_almost_equal(normalized_mutual_info_score(clustering.labels_,
                                                         labels), 1)
        clustering.connectivity = None
        clustering.fit(X)
        assert_true(np.size(np.unique(clustering.labels_)) == 10)
        # Check that we raise a TypeError on dense matrices
        clustering = AgglomerativeClustering(
            n_clusters=10,
            connectivity=sparse.lil_matrix(
                connectivity.toarray()[:10, :10]),
            linkage=linkage)
        assert_raises(ValueError, clustering.fit, X)

    # Test that using ward with another metric than euclidean raises an
    # exception
    clustering = AgglomerativeClustering(
        n_clusters=10,
        connectivity=connectivity.toarray(),
        affinity="manhattan",
        linkage="ward")
    assert_raises(ValueError, clustering.fit, X)

    # Test using another metric than euclidean works with linkage complete
    for affinity in PAIRED_DISTANCES.keys():
        # Compare our (structured) implementation to scipy
        clustering = AgglomerativeClustering(
            n_clusters=10,
            connectivity=np.ones((n_samples, n_samples)),
            affinity=affinity,
            linkage="complete")
        clustering.fit(X)
        clustering2 = AgglomerativeClustering(
            n_clusters=10,
            connectivity=None,
            affinity=affinity,
            linkage="complete")
        clustering2.fit(X)
        assert_almost_equal(normalized_mutual_info_score(clustering2.labels_,
                                                         clustering.labels_),
                            1)

    # Test that using a distance matrix (affinity = 'precomputed') has same
    # results (with connectivity constraints)
    clustering = AgglomerativeClustering(n_clusters=10,
                                         connectivity=connectivity,
                                         linkage="complete")
    clustering.fit(X)
    X_dist = pairwise_distances(X)
    clustering2 = AgglomerativeClustering(n_clusters=10,
                                          connectivity=connectivity,
                                          affinity='precomputed',
                                          linkage="complete")
    clustering2.fit(X_dist)
    assert_array_equal(clustering.labels_, clustering2.labels_)

def test_ward_agglomeration():
    # Check that we obtain the correct solution in a simplistic case
    rng = np.random.RandomState(0)
    mask = np.ones([10, 10], dtype=np.bool)
    X = rng.randn(50, 100)
    connectivity = grid_to_graph(*mask.shape)
    agglo = FeatureAgglomeration(n_clusters=5, connectivity=connectivity)
    agglo.fit(X)
    assert_true(np.size(np.unique(agglo.labels_)) == 5)

    X_red = agglo.transform(X)
    assert_true(X_red.shape[1] == 5)
    X_full = agglo.inverse_transform(X_red)
    assert_true(np.unique(X_full[0]).size == 5)
    assert_array_almost_equal(agglo.transform(X_full), X_red)

    # Check that fitting with no samples raises a ValueError
    assert_raises(ValueError, agglo.fit, X[:0])

def assess_same_labelling(cut1, cut2):
    """Util for comparison with scipy"""
    co_clust = []
    for cut in [cut1, cut2]:
        n = len(cut)
        k = cut.max() + 1
        ecut = np.zeros((n, k))
        ecut[np.arange(n), cut] = 1
        co_clust.append(np.dot(ecut, ecut.T))
    assert_true((co_clust[0] == co_clust[1]).all())

def test_scikit_vs_scipy():
    # Test scikit linkage with full connectivity (i.e. unstructured) vs scipy
    n, p, k = 10, 5, 3
    rng = np.random.RandomState(0)

    # Not using a lil_matrix here, just to check that non sparse
    # matrices are well handled
    connectivity = np.ones((n, n))
    for linkage in _TREE_BUILDERS.keys():
        for i in range(5):
            X = .1 * rng.normal(size=(n, p))
            X -= 4. * np.arange(n)[:, np.newaxis]
            X -= X.mean(axis=1)[:, np.newaxis]

            out = hierarchy.linkage(X, method=linkage)

            children_ = out[:, :2].astype(np.int)
            children, _, n_leaves, _ = _TREE_BUILDERS[linkage](X, connectivity)

            cut = _hc_cut(k, children, n_leaves)
            cut_ = _hc_cut(k, children_, n_leaves)
            assess_same_labelling(cut, cut_)

    # Test error management in _hc_cut
    assert_raises(ValueError, _hc_cut, n_leaves + 1, children, n_leaves)

def test_connectivity_propagation():
    # Check that connectivity in the ward tree is propagated correctly during
    # merging.
    X = np.array([(.014, .120), (.014, .099), (.014, .097),
                  (.017, .153), (.017, .153), (.018, .153),
                  (.018, .153), (.018, .153), (.018, .153),
                  (.018, .153), (.018, .153), (.018, .153),
                  (.018, .152), (.018, .149), (.018, .144)])
    connectivity = kneighbors_graph(X, 10, include_self=False)
    ward = AgglomerativeClustering(
        n_clusters=4, connectivity=connectivity, linkage='ward')
    # If changes are not propagated correctly, fit crashes with an
    # IndexError
    ward.fit(X)

def test_ward_tree_children_order():
    # Check that children are ordered in the same way for both structured and
    # unstructured versions of ward_tree.

    # test on five random datasets
    n, p = 10, 5
    rng = np.random.RandomState(0)

    connectivity = np.ones((n, n))
    for i in range(5):
        X = .1 * rng.normal(size=(n, p))
        X -= 4. * np.arange(n)[:, np.newaxis]
        X -= X.mean(axis=1)[:, np.newaxis]

        out_unstructured = ward_tree(X)
        out_structured = ward_tree(X, connectivity=connectivity)

        assert_array_equal(out_unstructured[0], out_structured[0])

def test_ward_linkage_tree_return_distance():
    # Test return_distance option on linkage and ward trees

    # test that return_distance when set true, gives same
    # output on both structured and unstructured clustering.
    n, p = 10, 5
    rng = np.random.RandomState(0)

    connectivity = np.ones((n, n))
    for i in range(5):
        X = .1 * rng.normal(size=(n, p))
        X -= 4. * np.arange(n)[:, np.newaxis]
        X -= X.mean(axis=1)[:, np.newaxis]

        out_unstructured = ward_tree(X, return_distance=True)
        out_structured = ward_tree(X, connectivity=connectivity,
                                   return_distance=True)

        # get children
        children_unstructured = out_unstructured[0]
        children_structured = out_structured[0]

        # check if we got the same clusters
        assert_array_equal(children_unstructured, children_structured)

        # check if the distances are the same
        dist_unstructured = out_unstructured[-1]
        dist_structured = out_structured[-1]

        assert_array_almost_equal(dist_unstructured, dist_structured)

        for linkage in ['average', 'complete']:
            structured_items = linkage_tree(
                X, connectivity=connectivity, linkage=linkage,
                return_distance=True)[-1]
            unstructured_items = linkage_tree(
                X, linkage=linkage, return_distance=True)[-1]
            structured_dist = structured_items[-1]
            unstructured_dist = unstructured_items[-1]
            structured_children = structured_items[0]
            unstructured_children = unstructured_items[0]
            assert_array_almost_equal(structured_dist, unstructured_dist)
            assert_array_almost_equal(
                structured_children, unstructured_children)

    # test on the following dataset where we know the truth
    # taken from scipy/cluster/tests/hierarchy_test_data.py
    X = np.array([[1.43054825, -7.5693489],
                  [6.95887839, 6.82293382],
                  [2.87137846, -9.68248579],
                  [7.87974764, -6.05485803],
                  [8.24018364, -6.09495602],
                  [7.39020262, 8.54004355]])
    # truth
    linkage_X_ward = np.array([[3., 4., 0.36265956, 2.],
                               [1., 5., 1.77045373, 2.],
                               [0., 2., 2.55760419, 2.],
                               [6., 8., 9.10208346, 4.],
                               [7., 9., 24.7784379, 6.]])

    linkage_X_complete = np.array(
        [[3., 4., 0.36265956, 2.],
         [1., 5., 1.77045373, 2.],
         [0., 2., 2.55760419, 2.],
         [6., 8., 6.96742194, 4.],
         [7., 9., 18.77445997, 6.]])

    linkage_X_average = np.array(
        [[3., 4., 0.36265956, 2.],
         [1., 5., 1.77045373, 2.],
         [0., 2., 2.55760419, 2.],
         [6., 8., 6.55832839, 4.],
         [7., 9., 15.44089605, 6.]])

    n_samples, n_features = np.shape(X)
    connectivity_X = np.ones((n_samples, n_samples))

    out_X_unstructured = ward_tree(X, return_distance=True)
    out_X_structured = ward_tree(X, connectivity=connectivity_X,
                                 return_distance=True)

    # check that the labels are the same
    assert_array_equal(linkage_X_ward[:, :2], out_X_unstructured[0])
    assert_array_equal(linkage_X_ward[:, :2], out_X_structured[0])

    # check that the distances are correct
    assert_array_almost_equal(linkage_X_ward[:, 2], out_X_unstructured[4])
    assert_array_almost_equal(linkage_X_ward[:, 2], out_X_structured[4])

    linkage_options = ['complete', 'average']
    X_linkage_truth = [linkage_X_complete, linkage_X_average]
    for (linkage, X_truth) in zip(linkage_options, X_linkage_truth):
        out_X_unstructured = linkage_tree(
            X, return_distance=True, linkage=linkage)
        out_X_structured = linkage_tree(
            X, connectivity=connectivity_X, linkage=linkage,
            return_distance=True)

        # check that the labels are the same
        assert_array_equal(X_truth[:, :2], out_X_unstructured[0])
        assert_array_equal(X_truth[:, :2], out_X_structured[0])

        # check that the distances are correct
        assert_array_almost_equal(X_truth[:, 2], out_X_unstructured[4])
        assert_array_almost_equal(X_truth[:, 2], out_X_structured[4])

def test_connectivity_fixing_non_lil():
    # Check non regression of a bug if a non item assignable connectivity is
    # provided with more than one component.
    # create dummy data
    x = np.array([[0, 0], [1, 1]])
    # create a mask with several components to force connectivity fixing
    m = np.array([[True, False], [False, True]])
    c = grid_to_graph(n_x=2, n_y=2, mask=m)
    w = AgglomerativeClustering(connectivity=c, linkage='ward')
    assert_warns(UserWarning, w.fit, x)

def test_int_float_dict():
    rng = np.random.RandomState(0)
    keys = np.unique(rng.randint(100, size=10).astype(np.intp))
    values = rng.rand(len(keys))

    d = IntFloatDict(keys, values)
    for key, value in zip(keys, values):
        assert d[key] == value

    other_keys = np.arange(50).astype(np.intp)[::2]
    other_values = 0.5 * np.ones(50)[::2]
    other = IntFloatDict(other_keys, other_values)
    # Complete smoke test
    max_merge(d, other, mask=np.ones(100, dtype=np.intp), n_a=1, n_b=1)
    average_merge(d, other, mask=np.ones(100, dtype=np.intp), n_a=1, n_b=1)

def test_connectivity_callable():
    rng = np.random.RandomState(0)
    X = rng.rand(20, 5)
    connectivity = kneighbors_graph(X, 3, include_self=False)
    aglc1 = AgglomerativeClustering(connectivity=connectivity)
    aglc2 = AgglomerativeClustering(
        connectivity=partial(kneighbors_graph, n_neighbors=3, include_self=False))
    aglc1.fit(X)
    aglc2.fit(X)
    assert_array_equal(aglc1.labels_, aglc2.labels_)

def test_connectivity_ignores_diagonal():
    rng = np.random.RandomState(0)
    X = rng.rand(20, 5)
    connectivity = kneighbors_graph(X, 3, include_self=False)
    connectivity_include_self = kneighbors_graph(X, 3, include_self=True)
    aglc1 = AgglomerativeClustering(connectivity=connectivity)
    aglc2 = AgglomerativeClustering(connectivity=connectivity_include_self)
    aglc1.fit(X)
    aglc2.fit(X)
    assert_array_equal(aglc1.labels_, aglc2.labels_)

def test_compute_full_tree():
    # Test that the full tree is computed if n_clusters is small
    rng = np.random.RandomState(0)
    X = rng.randn(10, 2)
    connectivity = kneighbors_graph(X, 5, include_self=False)

    # When n_clusters is less, the full tree should be built
    # that is the number of merges should be n_samples - 1
    agc = AgglomerativeClustering(n_clusters=2, connectivity=connectivity)
    agc.fit(X)
    n_samples = X.shape[0]
    n_nodes = agc.children_.shape[0]
    assert_equal(n_nodes, n_samples - 1)

    # When n_clusters is large, greater than max of 100 and 0.02 * n_samples.
    # we should stop when there are n_clusters.
    n_clusters = 101
    X = rng.randn(200, 2)
    connectivity = kneighbors_graph(X, 10, include_self=False)
    agc = AgglomerativeClustering(n_clusters=n_clusters,
                                  connectivity=connectivity)
    agc.fit(X)
    n_samples = X.shape[0]
    n_nodes = agc.children_.shape[0]
    assert_equal(n_nodes, n_samples - n_clusters)

def test_n_components():
    # Test n_components returned by linkage, average and ward tree
    rng = np.random.RandomState(0)
    X = rng.rand(5, 5)

    # Connectivity matrix having five components.
    connectivity = np.eye(5)

    for linkage_func in _TREE_BUILDERS.values():
        assert_equal(ignore_warnings(linkage_func)(X, connectivity)[1], 5)

def test_agg_n_clusters():
    # Test that an error is raised when n_clusters <= 0

    rng = np.random.RandomState(0)
    X = rng.rand(20, 10)
    for n_clus in [-1, 0]:
        agc = AgglomerativeClustering(n_clusters=n_clus)
        msg = ("n_clusters should be an integer greater than 0."
               " %s was provided." % str(agc.n_clusters))
        assert_raise_message(ValueError, msg, agc.fit, X)


"""Check the log N log F slope for future surveys."""
import numpy as np
import matplotlib.pyplot as plt
from copy import copy

from frbpoppy import CosmicPopulation, Survey, LargePopulation, SurveyPopulation, hist
from frbpoppy import unpickle, pprint
import frbpoppy.direction_dists as did
import frbpoppy.galacticops as go

from tests.convenience import plot_aa_style, rel_path

from tests.rates.alpha_real import EXPECTED

MAKE = True
SURVEYS = ('parkes-htru',
           'wsrt-apertif',
           'fast-crafts',
           'puma-full',
           'chord',
           'ska1-low',
           'ska1-mid')
SIZE = 5e4

if MAKE:
    # Calculate the fraction of the sky that the survey covers
    surv_f_area = {}
    for name in SURVEYS:
        pop = CosmicPopulation.simple(5e5)
        pop.gen_direction()
        survey = Survey(name)
        mask = survey.in_region(pop.frbs.ra, pop.frbs.dec,
                                pop.frbs.gl, pop.frbs.gb)
        in_surv_region = np.sum(mask)
        tot_region = len(mask)

        area_sky = 4*np.pi*(180/np.pi)**2   # In sq. degrees
        f_area = (survey.beam_size/area_sky)*(tot_region/in_surv_region)

        surv_f_area[name] = f_area
        print(f'{name} covers {f_area*100}% of the sky')

    surv_pops = []

    for name in SURVEYS:
        # Set up survey
        survey = Survey(name)
        if name in ('parkes-htru', 'wsrt-apertif'):
            survey.set_beam(model=name)

        # Set up CosmicPopulation
        pop = CosmicPopulation.optimal(SIZE, generate=False)

        # Only generate FRBs in the survey region
        pop.set_direction(model='uniform',
                          min_ra=survey.ra_min,
                          max_ra=survey.ra_max,
                          min_dec=survey.dec_min,
                          max_dec=survey.dec_max)

        # Parkes also has galactic limits:
        if name == 'parkes-htru':
            pop.gen_index()
            pop.gen_dist()
            pop.gen_time()

            # Generate FRBs just within the galactic constraints
            pop.gen_direction()

            # Gather ra, dec coordinate limits
            lims = {'min_ra': survey.ra_min, 'max_ra': survey.ra_max,
                    'min_dec': survey.dec_min, 'max_dec': survey.dec_max}

            def sample(n_gen):
                ra, dec = did.uniform(n_srcs=n_gen, **lims)
                gl, gb = go.radec_to_lb(ra, dec, frac=True)
                coords = [ra, dec, gl, gb]
                return coords

            def accept(coords):
                return survey.in_region(*coords)

            coords = sample(int(SIZE))
            mask = accept(coords)
            reject, = np.where(~mask)
            while reject.size > 0:
                fill = sample(reject.size)
                mask = accept(fill)
                for i in range(len(coords)):
                    coords[i][reject[mask]] = fill[i][mask]
                reject = reject[~mask]

            # Assign the values
            frbs = pop.frbs
            frbs.ra, frbs.dec = coords[0], coords[1]
            frbs.gl, frbs.gb = coords[2], coords[3]

            # Continue with generation
            pop.gen_gal_coords()
            pop.gen_dm()
            pop.gen_w()
            pop.gen_lum()
            pop.gen_si()
        else:
            pop.generate()

        surv_pop = SurveyPopulation(pop, survey, scale_by_area=False)
        surv_pop.source_rate.f_area = surv_f_area[name]
        surv_pop.source_rate.scale_by_area()
        # surv_pop.save()
        surv_pops.append(surv_pop)

else:
    surv_pops = []
    for name in SURVEYS:
        surv_pops.append(unpickle(f'optimal_{name}'))

# Start plot
plot_aa_style(cols=2)
plt.rcParams["figure.figsize"] = (3.556*3, 3.556)
fig, axes = plt.subplots(1, 3)

for ax in axes.flatten():
    ax.set_aspect('auto')

# Get norm pop
y = 0
ys = []
names = []
rates = []
norm_sim_rate = surv_pops[0].source_rate.det
norm_real_rate = EXPECTED['parkes-htru'][0] / EXPECTED['parkes-htru'][1]
norm_rate = norm_sim_rate / norm_real_rate

for i, surv_pop in enumerate(surv_pops):

    name = surv_pop.name.split('_')[-1]
    pprint(name)
    if surv_pop.n_sources() == 0:
        print(surv_pop.source_rate)
        print(f'{name} | no FRBs in population')
        continue

    names.append(name)
    ys.append(y)

    # Dimensions measure plot
    ax = axes[0]
    ax.set_xlabel(r'DM ($\textrm{pc}\ \textrm{cm}^{-3}$)')
    ax.set_ylabel(r'\#')
    ax.set_yscale('log')

    bins, values = hist(surv_pop.frbs.dm, bin_type='lin', norm='frac',
                        n_bins=20)
    values = values.astype(np.float64)
    values *= float(surv_pop.source_rate.f_area)*1e6
    ax.step(bins, values, where='mid', label=name)

    # Fluence plot
    ax = axes[1]
    ax.set_xlabel('S/N')
    ax.set_xscale('log')
    ax.set_ylabel(r'\#(${>}\text{S/N}$)')
    ax.set_yscale('log')

    # Update fluence plot
    bins, values = hist(surv_pop.frbs.snr, bin_type='log', norm='frac',
                        n_bins=25)

    # Cumulative sum
    values = np.cumsum(values[::-1])[::-1]
    values = values.astype(np.float64)
    values *= float(surv_pop.source_rate.f_area)*1e6
    ax.step(bins, values, where='mid', label=name)

    # Plot rates
    ax = axes[2]
    ax.set_xscale('log')
    ax.set_xlabel(r'Rate (day$^{-1}$)')
    rate = surv_pop.source_rate.det/norm_rate
    print(f'rate: {rate}')
    line = ax.errorbar(rate, y,
                       fmt='x',
                       label=rf'{name}')
    ax.grid()
    rates.append(rate)

    y += 1

ax.yaxis.tick_right()
ax.set_yticks(ys)
ax.set_yticklabels(names)
colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
for i, y in enumerate(ax.get_yticklabels()):
    y.set_color(colors[i])
ax.invert_yaxis()  # labels read top-to-bottom

# Add thin grey horizontal lines
x_lim = ax.get_xlim()
ax.set_xlim(x_lim)
for i, y in enumerate(ys):
    ax.plot((x_lim[0], rates[i]), (y, y), color='k', lw=0.5, zorder=0, ls='--')

for e in list(zip(SURVEYS, rates)):
    pprint(e)

euclidean_lines = True
if euclidean_lines:
    xlims = axes[1].get_xlim()
    ylims = axes[1].get_ylim()
    axes[1].set_xlim(xlims)
    axes[1].set_ylim(ylims)
    xs = np.logspace(np.log10(xlims[0]),
                     np.log10(xlims[1]),
                     100)
    for n in range(-10, 15):
        ys = 10**((np.log10(xs)+n)*-1.5)
        axes[1].plot(xs, ys, 'k:', linewidth=0.25)

# plt.legend()
plt.tight_layout()
plt.savefig(rel_path('./plots/future_surveys.pdf'))


import os
import numpy as np
import sonnet as snt
import tensorflow as tf
import matplotlib.pyplot as plt

from utils.data_utils_kitti import wrap_angle, compute_statistics, split_data, make_batch_iterator, make_repeating_batch_iterator, rotation_matrix, load_data_for_stats
from utils.method_utils import atan2, compute_sq_distance
from utils.plotting_utils import plot_maze, show_pause
from datetime import datetime

if tf.__version__ == '1.1.0-rc1' or tf.__version__ == '1.2.0':
    from tensorflow.python.framework import ops
    @ops.RegisterGradient("FloorMod")
    def _mod_grad(op, grad):
        x, y = op.inputs
        gz = grad
        x_grad = gz
        y_grad = None  # tf.reduce_mean(-(x // y) * gz, axis=[0], keep_dims=True)[0]
        return x_grad, y_grad

class DPF():

    def __init__(self, init_with_true_state, learn_odom, use_proposer, propose_ratio, proposer_keep_ratio, min_obs_likelihood, learn_gaussian_mle):
        """
        :param init_with_true_state:
        :param learn_odom:
        :param use_proposer:
        :param propose_ratio:
        :param particle_std:
        :param proposer_keep_ratio:
        :param min_obs_likelihood:
        """

        # store hyperparameters which are needed later
        self.init_with_true_state = init_with_true_state
        self.learn_odom = learn_odom
        self.use_proposer = use_proposer and not init_with_true_state  # only use proposer if we do not initializet with true state
        self.propose_ratio = propose_ratio if not self.init_with_true_state else 0.0

        # define some more parameters and placeholders
        self.state_dim = 5
        self.action_dim = 3
        self.observation_dim = 6
        self.placeholders = {'o': tf.placeholder('float32', [None, None, 50, 150, self.observation_dim], 'observations'),
                             'a': tf.placeholder('float32', [None, None, 3], 'actions'),
                             's': tf.placeholder('float32', [None, None, 5], 'states'),
                             'num_particles': tf.placeholder('float32'),
                             'keep_prob': tf.placeholder_with_default(tf.constant(1.0), []),
                             'is_training': tf.placeholder_with_default(tf.constant(False), [])
                             }
        self.num_particles_float = self.placeholders['num_particles']
        self.num_particles = tf.to_int32(self.num_particles_float)

        # build learnable modules
        self.build_modules(min_obs_likelihood, proposer_keep_ratio, learn_gaussian_mle)

    def build_modules(self, min_obs_likelihood, proposer_keep_ratio, learn_gaussian_mle):
        """
        :param min_obs_likelihood:
        :param proposer_keep_ratio:
        :return: None
        """

        # MEASUREMENT MODEL

        # conv net for encoding the image
        self.encoder = snt.Sequential([
            snt.nets.ConvNet2D([16, 16, 16, 16], [[7, 7], [5, 5], [5, 5], [5, 5]], [[1,1], [1, 2], [1, 2], [2, 2]], [snt.SAME], activate_final=True, name='encoder/convnet'),
            snt.BatchFlatten(),
            lambda x: tf.nn.dropout(x,  self.placeholders['keep_prob']),
            snt.Linear(128, name='encoder/linear'),
            tf.nn.relu
        ])

        # observation likelihood estimator that maps states and image encodings to probabilities
        self.obs_like_estimator = snt.Sequential([
            snt.Linear(128, name='obs_like_estimator/linear'),
            tf.nn.relu,
            snt.Linear(128, name='obs_like_estimator/linear'),
            tf.nn.relu,
            snt.Linear(1, name='obs_like_estimator/linear'),
            tf.nn.sigmoid,
            lambda x: x * (1 - min_obs_likelihood) + min_obs_likelihood
        ], name='obs_like_estimator')

        # motion noise generator used for motion sampling
        if learn_gaussian_mle:
            self.mo_noise_generator = snt.nets.MLP([32, 32, 4], activate_final=False, name='mo_noise_generator')
        else:
            self.mo_noise_generator = snt.nets.MLP([32, 32, 2], activate_final=False, name='mo_noise_generator')

        # odometry model (if we want to learn it)
        if self.learn_odom:
            self.mo_transition_model = snt.nets.MLP([128, 128, 128, self.state_dim], activate_final=False, name='mo_transition_model')

        # particle proposer that maps encodings to particles (if we want to use it)
        if self.use_proposer:
            self.particle_proposer = snt.Sequential([
                snt.Linear(128, name='particle_proposer/linear'),
                tf.nn.relu,
                lambda x: tf.nn.dropout(x,  proposer_keep_ratio),
                snt.Linear(128, name='particle_proposer/linear'),
                tf.nn.relu,
                snt.Linear(128, name='particle_proposer/linear'),
                tf.nn.relu,
                snt.Linear(128, name='particle_proposer/linear'),
                tf.nn.relu,
                snt.Linear(4, name='particle_proposer/linear'),
                tf.nn.tanh,
            ])

        self.noise_scaler1 = snt.Module(lambda x: x * tf.exp(10 * tf.get_variable('motion_sampler/noise_scaler1', initializer=np.array(0.0, dtype='float32'))))
        self.noise_scaler2 = snt.Module(lambda x: x * tf.exp(10 * tf.get_variable('motion_sampler/noise_scaler2', initializer=np.array(0.0, dtype='float32'))))

    def custom_build(self, inputs):
        """A custom build method to wrap into a sonnet Module."""
        outputs = snt.Conv2D(output_channels=16, kernel_shape=[7, 7], stride=[1, 1])(inputs)
        outputs = tf.nn.relu(outputs)
        outputs = snt.Conv2D(output_channels=16, kernel_shape=[5, 5], stride=[1, 2])(outputs)
        outputs = tf.nn.relu(outputs)
        outputs = snt.Conv2D(output_channels=16, kernel_shape=[5, 5], stride=[1, 2])(outputs)
        outputs = tf.nn.relu(outputs)
        outputs = snt.Conv2D(output_channels=16, kernel_shape=[5, 5], stride=[2, 2])(outputs)
        outputs = tf.nn.relu(outputs)
        outputs = tf.nn.dropout(outputs,  self.placeholders['keep_prob'])
        outputs = snt.BatchFlatten()(outputs)
        outputs = snt.Linear(128)(outputs)
        outputs = tf.nn.relu(outputs)

        return outputs

    def measurement_update(self, encoding, particles, means, stds):
        """
        Compute the likelihood of the encoded observation for each particle.

        :param encoding: encoding of the observation
        :param particles:
        :param means:
        :param stds:
        :return: observation likelihood
        """

        # prepare input (normalize particles poses and repeat encoding per particle)
        particle_input = self.transform_particles_as_input(particles, means, stds)
        encoding_input = tf.tile(encoding[:, tf.newaxis, :], [1,  tf.shape(particles)[1], 1])
        input = tf.concat([encoding_input, particle_input], axis=-1)

        # estimate the likelihood of the encoded observation for each particle, remove last dimension
        obs_likelihood = snt.BatchApply(self.obs_like_estimator)(input)[:, :, 0]

        return obs_likelihood

    def transform_particles_as_input(self, particles, means, stds):
        return ((particles - means['s']) / stds['s'])[..., 3:5]

    def propose_particles(self, encoding, num_particles, state_mins, state_maxs):
        duplicated_encoding = tf.tile(encoding[:, tf.newaxis, :], [1, num_particles, 1])
        proposed_particles = snt.BatchApply(self.particle_proposer)(duplicated_encoding)
        proposed_particles = tf.concat([
            proposed_particles[:,:,:1] * (state_maxs[0] - state_mins[0]) / 2.0 + (state_maxs[0] + state_mins[0]) / 2.0,
            proposed_particles[:,:,1:2] * (state_maxs[1] - state_mins[1]) / 2.0 + (state_maxs[1] + state_mins[1]) / 2.0,
            atan2(proposed_particles[:,:,2:3], proposed_particles[:,:,3:4])], axis=2)
        return proposed_particles

    def motion_update(self, actions, particles, means, stds, state_step_sizes, learn_gaussian_mle, stop_sampling_gradient=False):
        """
        Move particles according to odometry info in actions. Add learned noise.

        :param actions:
        :param particles:
        :param means:
        :param stds:
        :param state_step_sizes:
        :param stop_sampling_gradient:
        :return: moved particles
        """

        # 1. SAMPLE NOISY ACTIONS

        # add dimension for particles
        time_step = 0.103

        if learn_gaussian_mle:
            actions = tf.concat([particles[:, :, 3:4] - means['s'][:, :, 3:4], particles[:, :, 4:5] - means['s'][:, :, 4:5]], axis=-1)

            # prepare input (normalize actions and repeat per particle)
            action_input = actions / stds['s'][:, :, 3:5]
            input = action_input

            # estimate action noise
            delta = snt.BatchApply(self.mo_noise_generator)(input)
            delta = tf.concat([delta[:, :, 0:2] * state_step_sizes[3], delta[:, :, 2:4] * state_step_sizes[4]], axis=-1)
            if stop_sampling_gradient:
                delta = tf.stop_gradient(delta)

            action_vel_f = tf.random_normal(tf.shape(particles[:, :, 3:4]), mean = delta[:, :, 0:1], stddev = delta[:, :, 1:2])
            action_vel_rot = tf.random_normal(tf.shape(particles[:, :, 4:5]), mean = delta[:, :, 2:3], stddev = delta[:, :, 3:4])

            heading = particles[:, :, 2:3]
            sin_heading = tf.sin(heading)
            cos_heading = tf.cos(heading)

            new_x = particles[:, :, 0:1] + cos_heading * particles[:, :, 3:4] * time_step
            new_y = particles[:, :, 1:2] + sin_heading * particles[:, :, 3:4] * time_step
            new_theta = particles[:, :, 2:3] + particles[:, :, 4:5] * time_step
            wrap_angle(new_theta)
            new_v = particles[:, :, 3:4] + action_vel_f
            new_theta_dot = particles[:, :, 4:5] + action_vel_rot

            moved_particles = tf.concat([new_x, new_y, new_theta, new_v, new_theta_dot], axis=-1)

            return moved_particles, delta

        else:

            heading = particles[:, :, 2:3]
            sin_heading = tf.sin(heading)
            cos_heading = tf.cos(heading)

            random_input = tf.random_normal(tf.shape(particles[:, :, 3:5]))
            noise = snt.BatchApply(self.mo_noise_generator)(random_input)
            noise = noise - tf.reduce_mean(noise, axis=1, keep_dims=True)

            new_z = particles[:, :, 0:1] + cos_heading * particles[:, :, 3:4] * time_step
            new_x = particles[:, :, 1:2] + sin_heading * particles[:, :, 3:4] * time_step
            new_theta = wrap_angle(particles[:, :, 2:3] + particles[:, :, 4:5] * time_step)

            new_v = particles[:, :, 3:4] + noise[:, :, :1] * state_step_sizes[3]
            new_theta_dot = particles[:, :, 4:5] + noise[:, :, 1:] * state_step_sizes[4]

            moved_particles = tf.concat([new_z, new_x, new_theta, new_v, new_theta_dot], axis=-1)

            return moved_particles

    def compile_training_stages(self, sess, batch_iterators, particle_list, particle_probs_list, encodings, means, stds, state_step_sizes, state_mins, state_maxs, learn_gaussian_mle, learning_rate, plot_task):

        # TRAINING!
        losses = dict()
        train_stages = dict()
        std = 0.25

        # TRAIN ODOMETRY

        if self.learn_odom:

            # apply model
            motion_samples = self.motion_update(self.placeholders['a'][:,0],
                                                self.placeholders['s'][:, :1],
                                                means, stds, state_step_sizes,
                                                stop_sampling_gradient=True)

            # define loss and optimizer
            sq_distance = compute_sq_distance(motion_samples, self.placeholders['s'][:, 1:2], state_step_sizes)
            losses['motion_mse'] = tf.reduce_mean(sq_distance, name='loss')
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)

            # put everything together
            train_stages['train_odom'] = {
                         'train_op': optimizer.minimize(losses['motion_mse']),
                         'batch_iterator_names': {'train': 'train1', 'val': 'val1'},
                         'monitor_losses': ['motion_mse'],
                         'validation_loss': 'motion_mse',
                         'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val2']), motion_samples, plot_task, state_step_sizes) if e % 1 == 0 else None
                         }

        # TRAIN MOTION MODEL

        if learn_gaussian_mle:
            motion_samples, motion_params = self.motion_update(self.placeholders['a'][:,1],
                                                tf.tile(self.placeholders['s'][:, :1], [1, 1, 1]),
                                                means, stds, state_step_sizes, learn_gaussian_mle)

            # define loss and optimizer
            diff_in_states = self.placeholders['s'][:, 1:2] - self.placeholders['s'][:, :1]
            activations_vel_f = (1 / 32) / tf.sqrt(2 * np.pi * motion_params[:, :, 1] ** 2) * tf.exp(
                -(diff_in_states[:, :, 3] - motion_params[:, :, 0]) ** 2 / (2.0 * motion_params[:, :, 1] ** 2))
            activations_vel_rot = (1 / 32) / tf.sqrt(2 * np.pi * motion_params[:, :, 3] ** 2) * tf.exp(
                -(diff_in_states[:, :, 4] - motion_params[:, :, 2]) ** 2 / (2.0 * motion_params[:, :, 3] ** 2))
            losses['motion_mle'] = tf.reduce_mean(-tf.log(1e-16 + (tf.reduce_sum(activations_vel_f, axis=-1, name='loss1') * tf.reduce_sum(activations_vel_rot, axis=-1, name='loss2'))))
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)

            # put everything together
            train_stages['train_motion_sampling'] = {
                         'train_op': optimizer.minimize(losses['motion_mle']),
                         'batch_iterator_names': {'train': 'train2', 'val': 'val2'},
                         'monitor_losses': ['motion_mle'],
                         'validation_loss': 'motion_mle',
                         'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val2']), motion_samples, plot_task, state_step_sizes) if e % 1 == 0 else None
                         }

        else:
            motion_samples = self.motion_update(self.placeholders['a'][:,1],
                                    tf.tile(self.placeholders['s'][:, :1], [1, self.num_particles, 1]),
                                    means, stds, state_step_sizes, learn_gaussian_mle)

            # define loss and optimizer
            sq_distance = compute_sq_distance(motion_samples, self.placeholders['s'][:, 1:2], state_step_sizes)
            activations_sample = (1 / self.num_particles_float) / tf.sqrt(2 * np.pi * std ** 2) * tf.exp(
                -sq_distance / (2.0 * std ** 2))
            losses['motion_mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations_sample, axis=-1, name='loss')))
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)

            # put everything together
            train_stages['train_motion_sampling'] = {
                         'train_op': optimizer.minimize(losses['motion_mle']),
                         'batch_iterator_names': {'train': 'train2', 'val': 'val2'},
                         'monitor_losses': ['motion_mle'],
                         'validation_loss': 'motion_mle',
                         'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val2']), motion_samples, plot_task, state_step_sizes) if e % 1 == 0 else None
                         }

        # TRAIN MEASUREMENT MODEL

        # apply model for all pairs of observations and states in that batch
        test_particles = tf.tile(self.placeholders['s'][tf.newaxis, :, 0], [self.batch_size, 1, 1])
        measurement_model_out = self.measurement_update(encodings[:, 0], test_particles, means, stds)

        # define loss (correct -> 1, incorrect -> 0) and optimizer
        correct_samples = tf.diag_part(measurement_model_out)
        incorrect_samples = measurement_model_out - tf.diag(tf.diag_part(measurement_model_out))
        losses['measurement_heuristic'] = tf.reduce_sum(-tf.log(correct_samples)) / tf.cast(self.batch_size, tf.float32) \
                                          + tf.reduce_sum(-tf.log(1.0 - incorrect_samples)) / tf.cast(self.batch_size * (self.batch_size - 1), tf.float32)
        optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)

        # put everything together
        train_stages['train_measurement_model'] = {
                     'train_op': optimizer.minimize(losses['measurement_heuristic']),
                     'batch_iterator_names': {'train': 'train1', 'val': 'val1'},
                     'monitor_losses': ['measurement_heuristic'],
                     'validation_loss': 'measurement_heuristic',
                     'plot': lambda e: self.plot_measurement_model(sess, batch_iterators['val1'], measurement_model_out) if e % 1 == 0 else None
                     }

        # TRAIN PARTICLE PROPOSER

        if self.use_proposer:

            # apply model (but only compute gradients until the encoding,
            # otherwise we would unlearn it and the observation likelihood wouldn't work anymore)
            proposed_particles = self.propose_particles(tf.stop_gradient(encodings[:, 0]), self.num_particles, state_mins, state_maxs)

            # define loss and optimizer
            std = 0.2
            sq_distance = compute_sq_distance(proposed_particles, self.placeholders['s'][:, :1], state_step_sizes)
            activations = (1 / self.num_particles_float) / tf.sqrt(2 * np.pi * std ** 2) * tf.exp(
                -sq_distance / (2.0 * std ** 2))
            losses['proposed_mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations, axis=-1)))
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)

            # put everything together
            train_stages['train_particle_proposer'] = {
                         'train_op': optimizer.minimize(losses['proposed_mle']),
                         'batch_iterator_names': {'train': 'train1', 'val': 'val1'},
                         'monitor_losses': ['proposed_mle'],
                         'validation_loss': 'proposed_mle',
                         'plot': lambda e: self.plot_particle_proposer(sess, next(batch_iterators['val1']), proposed_particles, plot_task) if e % 10 == 0 else None
                         }

        # END-TO-END TRAINING

        # model was already applied further up -> particle_list, particle_probs_list

        # define losses and optimizer
        # first loss (which is being optimized)
        sq_distance = compute_sq_distance(particle_list[:, :, :, 3:5], self.placeholders['s'][:, :, tf.newaxis, 3:5], state_step_sizes[3:5])
        activations = particle_probs_list[:, :] / tf.sqrt(2 * np.pi * self.particle_std ** 2) * tf.exp(
            -sq_distance / (2.0 * self.particle_std ** 2))
        losses['mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations, axis=2, name='loss')))

        # second loss (which we will monitor during execution)
        pred = self.particles_to_state(particle_list, particle_probs_list)

        sq_error = compute_sq_distance(pred[:, -1, 0:2], self.placeholders['s'][:, -1, 0:2], [1., 1.])
        sq_dist = compute_sq_distance(self.placeholders['s'][:, 0, 0:2], self.placeholders['s'][:, -1, 0:2], [1., 1.])
        losses['m/m'] = tf.reduce_mean(sq_error**0.5/sq_dist**0.5)

        sq_error = compute_sq_distance(pred[:, -1, 2:3], self.placeholders['s'][:, -1, 2:3], [np.pi/180.0])
        losses['deg/m'] = tf.reduce_mean(sq_error ** 0.5 / sq_dist ** 0.5)

        # optimizer
        optimizer = tf.train.AdamOptimizer(learning_rate)

        # put everything together
        train_stages['train_e2e'] = {
                     'train_op': optimizer.minimize(losses['mle']),
                     'batch_iterator_names': {'train': 'train', 'val': 'val'},
                     'monitor_losses': ['m/m', 'deg/m', 'mle'],
                     'validation_loss': 'deg/m',
                     'plot': lambda e: self.plot_particle_filter(sess, next(batch_iterators['val_ex']), particle_list,
                                                                 particle_probs_list, state_step_sizes, plot_task) if e % 1 == 0 else None
                     }

        return losses, train_stages

    def load(self, sess, model_path, model_file='best_validation', statistics_file='statistics.npz', connect_and_initialize=True, modules=('encoder', 'mo_noise_generator', 'mo_transition_model', 'obs_like_estimator', 'particle_proposer')):

        if type(modules) not in [type(list()), type(tuple())]:
            raise Exception('modules must be a list or tuple, not a ' + str(type(modules)))

        # build the tensorflow graph
        if connect_and_initialize:
            # load training data statistics (which are needed to build the tf graph)
            statistics = dict(np.load(os.path.join(model_path, statistics_file)))
            for key in statistics.keys():
                if statistics[key].shape == ():
                    statistics[key] = statistics[key].item()  # convert 0d array of dictionary back to a normal dictionary

            # connect all modules into the particle filter
            self.connect_modules(**statistics)
            init = tf.global_variables_initializer()
            sess.run(init)

        # load variables
        all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
        vars_to_load = []
        loaded_modules = set()
        for v in all_vars:
            for m in modules:
                if m in v.name:
                    vars_to_load.append(v)
                    loaded_modules.add(m)

        print('Loading all modules')

        saver = tf.train.Saver()
        saver.restore(sess, os.path.join(model_path, model_file))

    # def fit(self, sess, data, model_path, train_individually, train_e2e, split_ratio, seq_len, batch_size, epoch_length, num_epochs, patience, learning_rate, dropout_keep_ratio, num_particles, particle_std, plot_task=None, plot=False):
    def fit(self, sess, data, model_path, train_individually, train_e2e, split_ratio, seq_len, batch_size, epoch_length, num_epochs, patience, learning_rate, dropout_keep_ratio, num_particles, particle_std, learn_gaussian_mle, plot_task=None, plot=False):
        if plot:
            plt.ion()

        self.particle_std = particle_std

        mean_loss_for_plot = np.zeros((1,))

        means, stds, state_step_sizes, state_mins, state_maxs = compute_statistics(data)

        data = split_data(data, ratio=split_ratio)

        epoch_lengths = {'train': epoch_length, 'val': epoch_length*2}
        batch_iterators = {'train': make_batch_iterator(data['train'], seq_len=seq_len, batch_size=batch_size),
                           'val': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=seq_len),
                           'train_ex': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=seq_len),
                           'val_ex': make_batch_iterator(data['val'], batch_size=batch_size, seq_len=seq_len),
                           'train1': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=1),
                           'train2': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=2),
                            'val1': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=1),
                            'val2': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=2),
                        }

        # build the tensorflow graph by connecting all modules in the particles filter
        particles, particle_probs, encodings, particle_list, particle_probs_list = self.connect_modules(means, stds, state_mins, state_maxs, state_step_sizes, learn_gaussian_mle)

        # define losses and train stages for different ways of training (e.g. training individual models and e2e training)
        losses, train_stages = self.compile_training_stages(sess, batch_iterators, particle_list, particle_probs_list,
                                                            encodings, means, stds, state_step_sizes, state_mins,
                                                            state_maxs, learn_gaussian_mle, learning_rate, plot_task)

        # initialize variables
        init = tf.global_variables_initializer()
        sess.run(init)

        # save statistics and prepare saving variables
        if not os.path.exists(model_path):
            os.makedirs(model_path)
        np.savez(os.path.join(model_path, 'statistics'), means=means, stds=stds, state_step_sizes=state_step_sizes,
                 state_mins=state_mins, state_maxs=state_maxs)
        saver = tf.train.Saver()
        save_path = os.path.join(model_path, 'best_validation')

        # define the training curriculum
        curriculum = []
        if train_individually:
            if self.learn_odom:
                curriculum += ['train_odom']
            curriculum += ['train_measurement_model']
            curriculum += ['train_motion_sampling']
            if self.use_proposer:
                curriculum += ['train_particle_proposer']
        if train_e2e:
            curriculum += ['train_e2e']

        # split data for early stopping
        data_keys = ['train']
        if split_ratio < 1.0:
            data_keys.append('val')

        # define log dict
        log = {c: {dk: {lk: {'mean': [], 'se': []} for lk in train_stages[c]['monitor_losses']} for dk in data_keys} for c in curriculum}

        # go through curriculum
        for c in curriculum:

            stage = train_stages[c]
            best_val_loss = np.inf
            best_epoch = 0
            epoch = 0

            if c == 'train_e2e':
                saver.save(sess, os.path.join(model_path, 'before_e2e/best_validation'))
                np.savez(os.path.join(model_path, 'before_e2e/statistics'), means=means, stds=stds, state_step_sizes=state_step_sizes,
                 state_mins=state_mins, state_maxs=state_maxs)
            while epoch < num_epochs and epoch - best_epoch < patience:
                # training
                for dk in data_keys:
                    # don't train in the first epoch, just evaluate the initial parameters
                    if dk == 'train' and epoch == 0:
                        continue
                    # set up loss lists which will be filled during the epoch
                    loss_lists = {lk: [] for lk in stage['monitor_losses']}
                    for e in range(epoch_lengths[dk]):
                        # t0 = time.time()
                        # pick a batch from the right iterator
                        batch = next(batch_iterators[stage['batch_iterator_names'][dk]])
                        # define the inputs and train/run the model
                        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},
                                      **{self.placeholders['num_particles']: num_particles},
                                      }
                        if dk == 'train':
                            input_dict[self.placeholders['keep_prob']] = dropout_keep_ratio
                            input_dict[self.placeholders['is_training']] = True
                        monitor_losses = {l: losses[l] for l in stage['monitor_losses']}
                        if dk == 'train':
                            s_losses, _ = sess.run([monitor_losses, stage['train_op']], input_dict)
                        else:
                            s_losses = sess.run(monitor_losses, input_dict)

                        for lk in stage['monitor_losses']:
                            loss_lists[lk].append(s_losses[lk])

                    # after each epoch, compute and log statistics
                    for lk in stage['monitor_losses']:
                        log[c][dk][lk]['mean'].append(np.mean(loss_lists[lk]))
                        log[c][dk][lk]['se'].append(np.std(loss_lists[lk], ddof=1) / np.sqrt(len(loss_lists[lk])))

                # check whether the current model is better than all previous models
                if 'val' in data_keys:
                    current_val_loss = log[c]['val'][stage['validation_loss']]['mean'][-1]
                    mean_loss_for_plot = np.append(mean_loss_for_plot,current_val_loss)
                    if current_val_loss < best_val_loss:
                        best_val_loss = current_val_loss
                        best_epoch = epoch
                        # save current model
                        saver.save(sess, save_path)
                        txt = 'epoch {:>3} >> '.format(epoch)
                    else:
                        txt = 'epoch {:>3} == '.format(epoch)
                else:
                    best_epoch = epoch
                    saver.save(sess, save_path)
                    txt = 'epoch {:>3} >> '.format(epoch)

                # after going through all data sets, do a print out of the current result
                for lk in stage['monitor_losses']:
                    txt += '{}: '.format(lk)
                    for dk in data_keys:
                        if len(log[c][dk][lk]['mean']) > 0:
                            txt += '{:.2f}+-{:.2f}/'.format(log[c][dk][lk]['mean'][-1], log[c][dk][lk]['se'][-1])

                    txt = txt[:-1] + ' -- '
                print(txt)

                if plot:
                    stage['plot'](epoch)

                epoch += 1

            # after running out of patience, restore the model with lowest validation loss
            saver.restore(sess, save_path)

        return log

    def predict(self, sess, batch, return_particles=False, **kwargs):
        # define input dict, use the first state only if we do tracking
        input_dict = {self.placeholders['o']: batch['o'],
                      self.placeholders['a']: batch['a'],
                      self.placeholders['num_particles']: 100}
        if self.init_with_true_state:
            input_dict[self.placeholders['s']] = batch['s'][:, :1]

        if return_particles:
            return sess.run([self.pred_states, self.particle_list, self.particle_probs_list], input_dict)
        else:
            return sess.run(self.pred_states, input_dict)

    def connect_modules(self, means, stds, state_mins, state_maxs, state_step_sizes, learn_gaussian_mle=False):

        # get shapes
        self.batch_size = tf.shape(self.placeholders['o'])[0]
        self.seq_len = tf.shape(self.placeholders['o'])[1]
        # we use the static shape here because we need it to build the graph
        self.action_dim = self.placeholders['a'].get_shape()[-1].value

        encodings = snt.BatchApply(self.encoder)((self.placeholders['o'] - means['o']) / stds['o'])

        # initialize particles
        if self.init_with_true_state:
            # tracking with known initial state
            initial_particles = tf.tile(self.placeholders['s'][:, 0, tf.newaxis, :], [1, self.num_particles, 1])
        else:
            # global localization
            if self.use_proposer:
                # propose particles from observations
                initial_particles = self.propose_particles(encodings[:, 0], self.num_particles, state_mins, state_maxs)
            else:
                # sample particles randomly
                initial_particles = tf.concat(
                    [tf.random_uniform([self.batch_size, self.num_particles, 1], state_mins[d], state_maxs[d]) for d in
                     range(self.state_dim)], axis=-1, name='particles')

        initial_particle_probs = tf.ones([self.batch_size, self.num_particles],
                                         name='particle_probs') / self.num_particles_float

        # assumes that samples has the correct size
        def permute_batch(x, samples):
            # get shapes
            batch_size = tf.shape(x)[0]
            num_particles = tf.shape(x)[1]
            sample_size = tf.shape(samples)[1]
            # compute 1D indices into the 2D array
            idx = samples + num_particles * tf.tile(
                tf.reshape(tf.range(batch_size), [batch_size, 1]),
                [1, sample_size])
            # index using the 1D indices and reshape again
            result = tf.gather(tf.reshape(x, [batch_size * num_particles, -1]), idx)
            result = tf.reshape(result, tf.shape(x[:,:sample_size]))
            return result

        def loop(particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i):

            num_proposed_float = tf.round((self.propose_ratio ** tf.cast(i, tf.float32)) * self.num_particles_float)
            num_proposed = tf.cast(num_proposed_float, tf.int32)
            num_resampled_float = self.num_particles_float - num_proposed_float
            num_resampled = tf.cast(num_resampled_float, tf.int32)

            if self.propose_ratio < 1.0:

                # resampling
                basic_markers = tf.linspace(0.0, (num_resampled_float - 1.0) / num_resampled_float, num_resampled)
                random_offset = tf.random_uniform([self.batch_size], 0.0, 1.0 / num_resampled_float)
                markers = random_offset[:, None] + basic_markers[None, :]  # shape: batch_size x num_resampled
                cum_probs = tf.cumsum(particle_probs, axis=1)
                marker_matching = markers[:, :, None] < cum_probs[:, None, :]  # shape: batch_size x num_resampled x num_particles
                samples = tf.cast(tf.argmax(tf.cast(marker_matching, 'int32'), dimension=2), 'int32')
                standard_particles = permute_batch(particles, samples)
                standard_particle_probs = tf.ones([self.batch_size, num_resampled])
                standard_particles = tf.stop_gradient(standard_particles)
                standard_particle_probs = tf.stop_gradient(standard_particle_probs)

                # motion update
                if learn_gaussian_mle:
                    standard_particles, _ = self.motion_update(self.placeholders['a'][:, i], standard_particles, means, stds, state_step_sizes, learn_gaussian_mle)
                else:
                    standard_particles = self.motion_update(self.placeholders['a'][:, i], standard_particles, means, stds, state_step_sizes, learn_gaussian_mle)

                # measurement update
                standard_particle_probs *= self.measurement_update(encodings[:, i], standard_particles, means, stds)

            if self.propose_ratio > 0.0:

                # proposed particles
                proposed_particles = self.propose_particles(encodings[:, i], num_proposed, state_mins, state_maxs)
                proposed_particle_probs = tf.ones([self.batch_size, num_proposed])

            # NORMALIZE AND COMBINE PARTICLES
            if self.propose_ratio == 1.0:
                particles = proposed_particles
                particle_probs = proposed_particle_probs

            elif self.propose_ratio == 0.0:
                particles = standard_particles
                particle_probs = standard_particle_probs

            else:
                standard_particle_probs *= (num_resampled_float / self.num_particles_float) / tf.reduce_sum(standard_particle_probs, axis=1, keep_dims=True)
                proposed_particle_probs *= (num_proposed_float / self.num_particles_float) / tf.reduce_sum(proposed_particle_probs, axis=1, keep_dims=True)
                particles = tf.concat([standard_particles, proposed_particles], axis=1)
                particle_probs = tf.concat([standard_particle_probs, proposed_particle_probs], axis=1)

            # NORMALIZE PROBABILITIES
            particle_probs /= tf.reduce_sum(particle_probs, axis=1, keep_dims=True)

            particle_list = tf.concat([particle_list, particles[:, tf.newaxis]], axis=1)
            particle_probs_list = tf.concat([particle_probs_list, particle_probs[:, tf.newaxis]], axis=1)

            return particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i + 1

        # reshapes and sets the first shape sizes to None (which is necessary to keep the shape consistent in while loop)
        particle_list = tf.reshape(initial_particles,
                                   shape=[self.batch_size, -1, self.num_particles, self.state_dim])
        particle_probs_list = tf.reshape(initial_particle_probs, shape=[self.batch_size, -1, self.num_particles])
        additional_probs_list = tf.reshape(tf.ones([self.batch_size, self.num_particles, 4]), shape=[self.batch_size, -1, self.num_particles, 4])

        # run the filtering process
        particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i = tf.while_loop(
            lambda *x: x[-1] < self.seq_len, loop,
            [initial_particles, initial_particle_probs, particle_list, particle_probs_list, additional_probs_list,
             tf.constant(1, dtype='int32')], name='loop')

        # compute mean of particles
        self.pred_states = self.particles_to_state(particle_list, particle_probs_list)
        self.particle_list = particle_list
        self.particle_probs_list = particle_probs_list

        return particles, particle_probs, encodings, particle_list, particle_probs_list

    def particles_to_state(self, particle_list, particle_probs_list):
        mean_position = tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * particle_list[:, :, :, :2], axis=2)
        mean_orientation = atan2(
            tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * tf.cos(particle_list[:, :, :, 2:3]), axis=2),
            tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * tf.sin(particle_list[:, :, :, 2:3]), axis=2))
        mean_velocity = tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * particle_list[:, :, :, 3:5], axis=2)
        return tf.concat([mean_position, mean_orientation, mean_velocity], axis=2)

    def plot_motion_model(self, sess, batch, motion_samples, task, state_step_sizes):

        # define the inputs and train/run the model
        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},
                      **{self.placeholders['num_particles']: 100},
                      }

        s_motion_samples = sess.run(motion_samples, input_dict)

        plt.figure('Motion Model')
        plt.gca().clear()
        for i in range(min(len(s_motion_samples), 10)):
            plt.scatter(s_motion_samples[i, :, 3] / state_step_sizes[3], s_motion_samples[i, :, 4] / state_step_sizes[4], color='blue', s=1)
            plt.scatter(batch['s'][i, 0, 3] / state_step_sizes[3], batch['s'][i, 0, 4] / state_step_sizes[4], color='black', s=1)
            plt.scatter(batch['s'][i, 1, 3] / state_step_sizes[3], batch['s'][i, 1, 4] / state_step_sizes[4], color='red', s=3)
            plt.plot(batch['s'][i, :2, 3] / state_step_sizes[3], batch['s'][i, :2, 4] / state_step_sizes[4], color='black')

        plt.xlim([0, 200])
        plt.ylim([-50, 50])
        plt.xlabel('translational vel')
        plt.ylabel('angular vel')
        plt.gca().set_aspect('equal')
        plt.pause(0.01)

    def plot_measurement_model(self, sess, batch_iterator, measurement_model_out):

        batch = next(batch_iterator)

        # define the inputs and train/run the model
        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},
                      **{self.placeholders['num_particles']: 100},
                      }

        s_measurement_model_out = sess.run([measurement_model_out], input_dict)

        plt.figure('Measurement Model Output')
        plt.gca().clear()
        plt.imshow(s_measurement_model_out[0], interpolation="nearest", cmap="viridis_r", vmin=0.0, vmax=1.0)

        plt.figure('Measurement Model Input')
        plt.clf()
        plt.scatter(batch['s'][:1, 0, 3], batch['s'][:1, 0, 4], marker='x', c=s_measurement_model_out[0][0,:1], vmin=0, vmax=1.0, cmap='viridis_r')
        plt.scatter(batch['s'][1:, 0, 3], batch['s'][1:, 0, 4], marker='o', c=s_measurement_model_out[0][0,1:], vmin=0, vmax=1.0, cmap='viridis_r')
        plt.xlabel('x_dot')
        plt.ylabel('theta_dot')
        plt.pause(0.01)

    def plot_particle_proposer(self, sess, batch, proposed_particles, task):

        # define the inputs and train/run the model
        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},
                      **{self.placeholders['num_particles']: 100},
                      }

        s_samples = sess.run(proposed_particles, input_dict)

        plt.figure('Particle Proposer')
        plt.gca().clear()
        plot_maze(task)

        for i in range(min(len(s_samples), 10)):
            color = np.random.uniform(0.0, 1.0, 3)
            plt.quiver(s_samples[i, :, 0], s_samples[i, :, 1], np.cos(s_samples[i, :, 2]), np.sin(s_samples[i, :, 2]), color=color, width=0.001, scale=100)
            plt.quiver(batch['s'][i, 0, 0], batch['s'][i, 0, 1], np.cos(batch['s'][i, 0, 2]), np.sin(batch['s'][i, 0, 2]), color=color, scale=50, width=0.003)

        plt.pause(0.01)

    def plot_particle_filter(self, sess, batch, particle_list,
                        particle_probs_list, state_step_sizes, task):

        s_states, s_particle_list, s_particle_probs_list, \
            = sess.run([self.placeholders['s'], particle_list,
                        particle_probs_list], #self.noise_scaler1(1.0), self.noise_scaler2(2.0)],
                       {**{self.placeholders[key]: batch[key] for key in 'osa'},
                        **{self.placeholders['num_particles']: 20},
                        })
        # print('learned motion noise factors {:.2f}/{:.2f}'.format(n1, n2))

        num_steps = s_particle_list.shape[1]

        for s in range(3):

            plt.figure('particle_evolution, example {}'.format(s))
            plt.clf()

            for d in range(5):

                plt.subplot(3, 2, [1, 3, 5, 2, 4][d])

                for i in range(num_steps):

                    plt.scatter(i * np.ones_like(s_particle_list[s, i, :, d]),
                                s_particle_list[s, i, :, d] / (1 if s == 0 else state_step_sizes[d]),
                                c=s_particle_probs_list[s, i, :], cmap='viridis_r', marker='o', s=6, alpha=0.5,
                                linewidths=0.05,
                                vmin=0.0,
                                vmax=0.1)
                    current_state = batch['s'][s, i, d] / (1 if s == 0 else state_step_sizes[d])
                    plt.plot([i], [current_state], 'o', markerfacecolor='None', markeredgecolor='k',
                             markersize=2.5)

                plt.xlabel('Time')
                plt.ylabel('State {}'.format(d))

        show_pause(pause=0.01)