import unittest

import networkx as nx
import pandas as pd
import numpy as np
import numpy.testing as np_test

from pgmpy.models import BayesianModel, MarkovModel
import pgmpy.tests.help_functions as hf
from pgmpy.factors.discrete import TabularCPD, JointProbabilityDistribution, DiscreteFactor
from pgmpy.independencies import Independencies
from pgmpy.estimators import BayesianEstimator, BaseEstimator, MaximumLikelihoodEstimator

class TestBaseModelCreation(unittest.TestCase):

    def setUp(self):
        self.G = BayesianModel()

    def test_class_init_without_data(self):
        self.assertIsInstance(self.G, nx.DiGraph)

    def test_class_init_with_data_string(self):
        self.g = BayesianModel([('a', 'b'), ('b', 'c')])
        self.assertListEqual(sorted(self.g.nodes()), ['a', 'b', 'c'])
        self.assertListEqual(hf.recursive_sorted(self.g.edges()),
                             [['a', 'b'], ['b', 'c']])

    def test_class_init_with_data_nonstring(self):
        BayesianModel([(1, 2), (2, 3)])

    def test_add_node_string(self):
        self.G.add_node('a')
        self.assertListEqual(self.G.nodes(), ['a'])

    def test_add_node_nonstring(self):
        self.G.add_node(1)

    def test_add_nodes_from_string(self):
        self.G.add_nodes_from(['a', 'b', 'c', 'd'])
        self.assertListEqual(sorted(self.G.nodes()), ['a', 'b', 'c', 'd'])

    def test_add_nodes_from_non_string(self):
        self.G.add_nodes_from([1, 2, 3, 4])

    def test_add_edge_string(self):
        self.G.add_edge('d', 'e')
        self.assertListEqual(sorted(self.G.nodes()), ['d', 'e'])
        self.assertListEqual(self.G.edges(), [('d', 'e')])
        self.G.add_nodes_from(['a', 'b', 'c'])
        self.G.add_edge('a', 'b')
        self.assertListEqual(hf.recursive_sorted(self.G.edges()),
                             [['a', 'b'], ['d', 'e']])

    def test_add_edge_nonstring(self):
        self.G.add_edge(1, 2)

    def test_add_edge_selfloop(self):
        self.assertRaises(ValueError, self.G.add_edge, 'a', 'a')

    def test_add_edge_result_cycle(self):
        self.G.add_edges_from([('a', 'b'), ('a', 'c')])
        self.assertRaises(ValueError, self.G.add_edge, 'c', 'a')

    def test_add_edges_from_string(self):
        self.G.add_edges_from([('a', 'b'), ('b', 'c')])
        self.assertListEqual(sorted(self.G.nodes()), ['a', 'b', 'c'])
        self.assertListEqual(hf.recursive_sorted(self.G.edges()),
                             [['a', 'b'], ['b', 'c']])
        self.G.add_nodes_from(['d', 'e', 'f'])
        self.G.add_edges_from([('d', 'e'), ('e', 'f')])
        self.assertListEqual(sorted(self.G.nodes()),
                             ['a', 'b', 'c', 'd', 'e', 'f'])
        self.assertListEqual(hf.recursive_sorted(self.G.edges()),
                             hf.recursive_sorted([('a', 'b'), ('b', 'c'),
                                                  ('d', 'e'), ('e', 'f')]))

    def test_add_edges_from_nonstring(self):
        self.G.add_edges_from([(1, 2), (2, 3)])

    def test_add_edges_from_self_loop(self):
        self.assertRaises(ValueError, self.G.add_edges_from,
                          [('a', 'a')])

    def test_add_edges_from_result_cycle(self):
        self.assertRaises(ValueError, self.G.add_edges_from,
                          [('a', 'b'), ('b', 'c'), ('c', 'a')])

    def test_update_node_parents_bm_constructor(self):
        self.g = BayesianModel([('a', 'b'), ('b', 'c')])
        self.assertListEqual(self.g.predecessors('a'), [])
        self.assertListEqual(self.g.predecessors('b'), ['a'])
        self.assertListEqual(self.g.predecessors('c'), ['b'])

    def test_update_node_parents(self):
        self.G.add_nodes_from(['a', 'b', 'c'])
        self.G.add_edges_from([('a', 'b'), ('b', 'c')])
        self.assertListEqual(self.G.predecessors('a'), [])
        self.assertListEqual(self.G.predecessors('b'), ['a'])
        self.assertListEqual(self.G.predecessors('c'), ['b'])

    def tearDown(self):
        del self.G

class TestBayesianModelMethods(unittest.TestCase):

    def setUp(self):
        self.G = BayesianModel([('a', 'd'), ('b', 'd'),
                                ('d', 'e'), ('b', 'c')])
        self.G1 = BayesianModel([('diff', 'grade'), ('intel', 'grade')])
        diff_cpd = TabularCPD('diff', 2, values=[[0.2], [0.8]])
        intel_cpd = TabularCPD('intel', 3, values=[[0.5], [0.3], [0.2]])
        grade_cpd = TabularCPD('grade', 3, values=[[0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
                                                   [0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
                                                   [0.8, 0.8, 0.8, 0.8, 0.8, 0.8]],
                               evidence=['diff', 'intel'], evidence_card=[2, 3])
        self.G1.add_cpds(diff_cpd, intel_cpd, grade_cpd)

    def test_moral_graph(self):
        moral_graph = self.G.moralize()
        self.assertListEqual(sorted(moral_graph.nodes()), ['a', 'b', 'c', 'd', 'e'])
        for edge in moral_graph.edges():
            self.assertTrue(edge in [('a', 'b'), ('a', 'd'), ('b', 'c'), ('d', 'b'), ('e', 'd')] or
                            (edge[1], edge[0]) in [('a', 'b'), ('a', 'd'), ('b', 'c'), ('d', 'b'), ('e', 'd')])

    def test_moral_graph_with_edge_present_over_parents(self):
        G = BayesianModel([('a', 'd'), ('d', 'e'), ('b', 'd'), ('b', 'c'), ('a', 'b')])
        moral_graph = G.moralize()
        self.assertListEqual(sorted(moral_graph.nodes()), ['a', 'b', 'c', 'd', 'e'])
        for edge in moral_graph.edges():
            self.assertTrue(edge in [('a', 'b'), ('c', 'b'), ('d', 'a'), ('d', 'b'), ('d', 'e')] or
                            (edge[1], edge[0]) in [('a', 'b'), ('c', 'b'), ('d', 'a'), ('d', 'b'), ('d', 'e')])

    def test_local_independencies(self):
        self.assertEqual(self.G.local_independencies('a'), Independencies(['a', ['b', 'c']]))
        self.assertEqual(self.G.local_independencies('c'), Independencies(['c', ['a', 'd', 'e'], 'b']))
        self.assertEqual(self.G.local_independencies('d'), Independencies(['d', 'c', ['b', 'a']]))
        self.assertEqual(self.G.local_independencies('e'), Independencies(['e', ['c', 'b', 'a'], 'd']))
        self.assertEqual(self.G.local_independencies('b'), Independencies(['b', 'a']))
        self.assertEqual(self.G1.local_independencies('grade'), Independencies())

    def test_get_independencies(self):
        chain = BayesianModel([('X', 'Y'), ('Y', 'Z')])
        self.assertEqual(chain.get_independencies(), Independencies(('X', 'Z', 'Y'), ('Z', 'X', 'Y')))
        fork = BayesianModel([('Y', 'X'), ('Y', 'Z')])
        self.assertEqual(fork.get_independencies(), Independencies(('X', 'Z', 'Y'), ('Z', 'X', 'Y')))
        collider = BayesianModel([('X', 'Y'), ('Z', 'Y')])
        self.assertEqual(collider.get_independencies(), Independencies(('X', 'Z'), ('Z', 'X')))

    def test_is_imap(self):
        val = [0.01, 0.01, 0.08, 0.006, 0.006, 0.048, 0.004, 0.004, 0.032,
               0.04, 0.04, 0.32, 0.024, 0.024, 0.192, 0.016, 0.016, 0.128]
        JPD = JointProbabilityDistribution(['diff', 'intel', 'grade'], [2, 3, 3], val)
        fac = DiscreteFactor(['diff', 'intel', 'grade'], [2, 3, 3], val)
        self.assertTrue(self.G1.is_imap(JPD))
        self.assertRaises(TypeError, self.G1.is_imap, fac)

    def test_get_immoralities(self):
        G = BayesianModel([('x', 'y'), ('z', 'y'), ('x', 'z'), ('w', 'y')])
        self.assertEqual(G.get_immoralities(), {('w', 'x'), ('w', 'z')})
        G1 = BayesianModel([('x', 'y'), ('z', 'y'), ('z', 'x'), ('w', 'y')])
        self.assertEqual(G1.get_immoralities(), {('w', 'x'), ('w', 'z')})
        G2 = BayesianModel([('x', 'y'), ('z', 'y'), ('x', 'z'), ('w', 'y'), ('w', 'x')])
        self.assertEqual(G2.get_immoralities(), {('w', 'z')})

    def test_is_iequivalent(self):
        G = BayesianModel([('x', 'y'), ('z', 'y'), ('x', 'z'), ('w', 'y')])
        self.assertRaises(TypeError, G.is_iequivalent, MarkovModel())
        G1 = BayesianModel([('V', 'W'), ('W', 'X'), ('X', 'Y'), ('Z', 'Y')])
        G2 = BayesianModel([('W', 'V'), ('X', 'W'), ('X', 'Y'), ('Z', 'Y')])
        self.assertTrue(G1.is_iequivalent(G2))
        G3 = BayesianModel([('W', 'V'), ('W', 'X'), ('Y', 'X'), ('Z', 'Y')])
        self.assertFalse(G3.is_iequivalent(G2))

    def tearDown(self):
        del self.G
        del self.G1

class TestBayesianModelCPD(unittest.TestCase):

    def setUp(self):
        self.G = BayesianModel([('d', 'g'), ('i', 'g'), ('g', 'l'),
                                ('i', 's')])

    def test_active_trail_nodes(self):
        self.assertEqual(sorted(self.G.active_trail_nodes('d')), ['d', 'g', 'l'])
        self.assertEqual(sorted(self.G.active_trail_nodes('i')), ['g', 'i', 'l', 's'])

    def test_active_trail_nodes_args(self):
        self.assertEqual(sorted(self.G.active_trail_nodes('d', observed='g')), ['d', 'i', 's'])
        self.assertEqual(sorted(self.G.active_trail_nodes('l', observed='g')), ['l'])
        self.assertEqual(sorted(self.G.active_trail_nodes('s', observed=['i', 'l'])), ['s'])
        self.assertEqual(sorted(self.G.active_trail_nodes('s', observed=['d', 'l'])), ['g', 'i', 's'])

    def test_is_active_trail_triplets(self):
        self.assertTrue(self.G.is_active_trail('d', 'l'))
        self.assertTrue(self.G.is_active_trail('g', 's'))
        self.assertFalse(self.G.is_active_trail('d', 'i'))
        self.assertTrue(self.G.is_active_trail('d', 'i', observed='g'))
        self.assertFalse(self.G.is_active_trail('d', 'l', observed='g'))
        self.assertFalse(self.G.is_active_trail('i', 'l', observed='g'))
        self.assertTrue(self.G.is_active_trail('d', 'i', observed='l'))
        self.assertFalse(self.G.is_active_trail('g', 's', observed='i'))

    def test_is_active_trail(self):
        self.assertFalse(self.G.is_active_trail('d', 's'))
        self.assertTrue(self.G.is_active_trail('s', 'l'))
        self.assertTrue(self.G.is_active_trail('d', 's', observed='g'))
        self.assertFalse(self.G.is_active_trail('s', 'l', observed='g'))

    def test_is_active_trail_args(self):
        self.assertFalse(self.G.is_active_trail('s', 'l', 'i'))
        self.assertFalse(self.G.is_active_trail('s', 'l', 'g'))
        self.assertTrue(self.G.is_active_trail('d', 's', 'l'))
        self.assertFalse(self.G.is_active_trail('d', 's', ['i', 'l']))

    def test_get_cpds(self):
        cpd_d = TabularCPD('d', 2, values=np.random.rand(2, 1))
        cpd_i = TabularCPD('i', 2, values=np.random.rand(2, 1))
        cpd_g = TabularCPD('g', 2, values=np.random.rand(2, 4),
                           evidence=['d', 'i'], evidence_card=[2, 2])
        cpd_l = TabularCPD('l', 2, values=np.random.rand(2, 2),
                           evidence=['g'], evidence_card=[2])
        cpd_s = TabularCPD('s', 2, values=np.random.rand(2, 2),
                           evidence=['i'], evidence_card=[2])
        self.G.add_cpds(cpd_d, cpd_i, cpd_g, cpd_l, cpd_s)

        self.assertEqual(self.G.get_cpds('d').variable, 'd')

    def test_get_cpds1(self):
        self.model = BayesianModel([('A', 'AB')])
        cpd_a = TabularCPD('A', 2, values=np.random.rand(2, 1))
        cpd_ab = TabularCPD('AB', 2, values=np.random.rand(2, 2),
                            evidence=['A'], evidence_card=[2])

        self.model.add_cpds(cpd_a, cpd_ab)
        self.assertEqual(self.model.get_cpds('A').variable, 'A')
        self.assertEqual(self.model.get_cpds('AB').variable, 'AB')
        self.assertRaises(ValueError, self.model.get_cpds, 'B')

        self.model.add_node('B')
        self.assertRaises(ValueError, self.model.get_cpds, 'B')

    def test_add_single_cpd(self):
        cpd_s = TabularCPD('s', 2, np.random.rand(2, 2), ['i'], [2])
        self.G.add_cpds(cpd_s)
        self.assertListEqual(self.G.get_cpds(), [cpd_s])

    def test_add_multiple_cpds(self):
        cpd_d = TabularCPD('d', 2, values=np.random.rand(2, 1))
        cpd_i = TabularCPD('i', 2, values=np.random.rand(2, 1))
        cpd_g = TabularCPD('g', 2, values=np.random.rand(2, 4),
                           evidence=['d', 'i'], evidence_card=[2, 2])
        cpd_l = TabularCPD('l', 2, values=np.random.rand(2, 2),
                           evidence=['g'], evidence_card=[2])
        cpd_s = TabularCPD('s', 2, values=np.random.rand(2, 2),
                           evidence=['i'], evidence_card=[2])

        self.G.add_cpds(cpd_d, cpd_i, cpd_g, cpd_l, cpd_s)
        self.assertEqual(self.G.get_cpds('d'), cpd_d)
        self.assertEqual(self.G.get_cpds('i'), cpd_i)
        self.assertEqual(self.G.get_cpds('g'), cpd_g)
        self.assertEqual(self.G.get_cpds('l'), cpd_l)
        self.assertEqual(self.G.get_cpds('s'), cpd_s)

    def test_check_model(self):
        cpd_g = TabularCPD('g', 2, values=np.array([[0.2, 0.3, 0.4, 0.6],
                                                    [0.8, 0.7, 0.6, 0.4]]),
                           evidence=['d', 'i'], evidence_card=[2, 2])

        cpd_s = TabularCPD('s', 2, values=np.array([[0.2, 0.3],
                                                    [0.8, 0.7]]),
                           evidence=['i'], evidence_card=[2])

        cpd_l = TabularCPD('l', 2, values=np.array([[0.2, 0.3],
                                                    [0.8, 0.7]]),
                           evidence=['g'], evidence_card=[2])

        self.G.add_cpds(cpd_g, cpd_s, cpd_l)
        self.assertRaises(ValueError, self.G.check_model)

        cpd_d = TabularCPD('d', 2, values=[[0.8, 0.2]])
        cpd_i = TabularCPD('i', 2, values=[[0.7, 0.3]])
        self.G.add_cpds(cpd_d, cpd_i)

        self.assertTrue(self.G.check_model())

    def test_check_model1(self):
        cpd_g = TabularCPD('g', 2, values=np.array([[0.2, 0.3],
                                                    [0.8, 0.7]]),
                           evidence=['i'], evidence_card=[2])
        self.G.add_cpds(cpd_g)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_g)

        cpd_g = TabularCPD('g', 2, values=np.array([[0.2, 0.3, 0.4, 0.6],
                                                    [0.8, 0.7, 0.6, 0.4]]),
                           evidence=['d', 's'], evidence_card=[2, 2])
        self.G.add_cpds(cpd_g)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_g)

        cpd_g = TabularCPD('g', 2, values=np.array([[0.2, 0.3],
                                                    [0.8, 0.7]]),
                           evidence=['l'], evidence_card=[2])
        self.G.add_cpds(cpd_g)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_g)

        cpd_l = TabularCPD('l', 2, values=np.array([[0.2, 0.3],
                                                    [0.8, 0.7]]),
                           evidence=['d'], evidence_card=[2])
        self.G.add_cpds(cpd_l)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_l)

        cpd_l = TabularCPD('l', 2, values=np.array([[0.2, 0.3, 0.4, 0.6],
                                                    [0.8, 0.7, 0.6, 0.4]]),
                           evidence=['d', 'i'], evidence_card=[2, 2])
        self.G.add_cpds(cpd_l)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_l)

        cpd_l = TabularCPD('l', 2, values=np.array([[0.2, 0.3, 0.4, 0.6, 0.2, 0.3, 0.4, 0.6],
                                                    [0.8, 0.7, 0.6, 0.4, 0.8, 0.7, 0.6, 0.4]]),
                           evidence=['g', 'd', 'i'], evidence_card=[2, 2, 2])
        self.G.add_cpds(cpd_l)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_l)

    def test_check_model2(self):
        cpd_s = TabularCPD('s', 2, values=np.array([[0.5, 0.3],
                                                    [0.8, 0.7]]),
                           evidence=['i'], evidence_card=[2])
        self.G.add_cpds(cpd_s)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_s)

        cpd_g = TabularCPD('g', 2, values=np.array([[0.2, 0.3, 0.4, 0.6],
                                                    [0.3, 0.7, 0.6, 0.4]]),
                           evidence=['d', 'i'], evidence_card=[2, 2])
        self.G.add_cpds(cpd_g)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_g)

        cpd_l = TabularCPD('l', 2, values=np.array([[0.2, 0.3],
                                                    [0.1, 0.7]]),
                           evidence=['g'], evidence_card=[2])
        self.G.add_cpds(cpd_l)
        self.assertRaises(ValueError, self.G.check_model)
        self.G.remove_cpds(cpd_l)

    def tearDown(self):
        del self.G

class TestBayesianModelFitPredict(unittest.TestCase):

    def setUp(self):
        self.model_disconnected = BayesianModel()
        self.model_disconnected.add_nodes_from(['A', 'B', 'C', 'D', 'E'])
        self.model_connected = BayesianModel([('A', 'B'), ('C', 'B'), ('C', 'D'), ('B', 'E')])

        self.model2 = BayesianModel([('A', 'C'), ('B', 'C')])
        self.data1 = pd.DataFrame(data={'A': [0, 0, 1], 'B': [0, 1, 0], 'C': [1, 1, 0]})
        self.data2 = pd.DataFrame(data={'A': [0, np.NaN, 1],
                                        'B': [0, 1, 0],
                                        'C': [1, 1, np.NaN],
                                        'D': [np.NaN, 'Y', np.NaN]})

    def test_bayesian_fit(self):
        print(isinstance(BayesianEstimator, BaseEstimator))
        print(isinstance(MaximumLikelihoodEstimator, BaseEstimator))
        self.model2.fit(self.data1, estimator_type=BayesianEstimator, prior_type="dirichlet", pseudo_counts=[9, 3])
        self.assertEqual(self.model2.get_cpds('B'), TabularCPD('B', 2, [[11.0 / 15], [4.0 / 15]]))

    def test_fit_missing_data(self):
        self.model2.fit(self.data2, state_names={'C': [0, 1]}, complete_samples_only=False)
        cpds = set([TabularCPD('A', 2, [[0.5], [0.5]]),
                    TabularCPD('B', 2, [[2. / 3], [1. / 3]]),
                    TabularCPD('C', 2, [[0, 0.5, 0.5, 0.5], [1, 0.5, 0.5, 0.5]],
                               evidence=['A', 'B'], evidence_card=[2, 2])])
        self.assertSetEqual(cpds, set(self.model2.get_cpds()))

    def test_disconnected_fit(self):
        values = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 5)),
                              columns=['A', 'B', 'C', 'D', 'E'])
        self.model_disconnected.fit(values)

        for node in ['A', 'B', 'C', 'D', 'E']:
            cpd = self.model_disconnected.get_cpds(node)
            self.assertEqual(cpd.variable, node)
            np_test.assert_array_equal(cpd.cardinality, np.array([2]))
            value = (values.ix[:, node].value_counts() /
                     values.ix[:, node].value_counts().sum())
            value = value.reindex(sorted(value.index)).values
            np_test.assert_array_equal(cpd.values, value)

    def test_connected_predict(self):
        np.random.seed(42)
        values = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 5)),
                              columns=['A', 'B', 'C', 'D', 'E'])
        fit_data = values[:800]
        predict_data = values[800:].copy()
        self.model_connected.fit(fit_data)
        self.assertRaises(ValueError, self.model_connected.predict, predict_data)
        predict_data.drop('E', axis=1, inplace=True)
        e_predict = self.model_connected.predict(predict_data)
        np_test.assert_array_equal(e_predict.values.ravel(),
                                   np.array([1, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1,
                                             1, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0,
                                             0, 1, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1, 0, 0,
                                             0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 0, 1,
                                             0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 1,
                                             1, 0, 1, 1, 0, 1, 1, 1, 0, 0, 1, 1, 0, 1,
                                             1, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0,
                                             1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1,
                                             0, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1,
                                             1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1,
                                             1, 0, 1, 0, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1,
                                             0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0,
                                             1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 1,
                                             1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 1,
                                             1, 1, 1, 0]))

    def tearDown(self):
        del self.model_connected
        del self.model_disconnected

class TestDirectedGraphCPDOperations(unittest.TestCase):

    def setUp(self):
        self.graph = BayesianModel()

    def test_add_single_cpd(self):
        cpd = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                         evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd)
        self.assertListEqual(self.graph.get_cpds(), [cpd])

    def test_add_multiple_cpds(self):
        cpd1 = TabularCPD('diff', 2, values=np.random.rand(2, 1))
        cpd2 = TabularCPD('intel', 2, values=np.random.rand(2, 1))
        cpd3 = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                          evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd1, cpd2, cpd3)
        self.assertListEqual(self.graph.get_cpds(), [cpd1, cpd2, cpd3])

    def test_remove_single_cpd(self):
        cpd1 = TabularCPD('diff', 2, values=np.random.rand(2, 1))
        cpd2 = TabularCPD('intel', 2, values=np.random.rand(2, 1))
        cpd3 = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                          evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd1, cpd2, cpd3)
        self.graph.remove_cpds(cpd1)
        self.assertListEqual(self.graph.get_cpds(), [cpd2, cpd3])

    def test_remove_multiple_cpds(self):
        cpd1 = TabularCPD('diff', 2, values=np.random.rand(2, 1))
        cpd2 = TabularCPD('intel', 2, values=np.random.rand(2, 1))
        cpd3 = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                          evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd1, cpd2, cpd3)
        self.graph.remove_cpds(cpd1, cpd3)
        self.assertListEqual(self.graph.get_cpds(), [cpd2])

    def test_remove_single_cpd_string(self):
        cpd1 = TabularCPD('diff', 2, values=np.random.rand(2, 1))
        cpd2 = TabularCPD('intel', 2, values=np.random.rand(2, 1))
        cpd3 = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                          evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd1, cpd2, cpd3)
        self.graph.remove_cpds('diff')
        self.assertListEqual(self.graph.get_cpds(), [cpd2, cpd3])

    def test_remove_multiple_cpds_string(self):
        cpd1 = TabularCPD('diff', 2, values=np.random.rand(2, 1))
        cpd2 = TabularCPD('intel', 2, values=np.random.rand(2, 1))
        cpd3 = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                          evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd1, cpd2, cpd3)
        self.graph.remove_cpds('diff', 'grade')
        self.assertListEqual(self.graph.get_cpds(), [cpd2])

    def test_get_cpd_for_node(self):
        cpd1 = TabularCPD('diff', 2, values=np.random.rand(2, 1))
        cpd2 = TabularCPD('intel', 2, values=np.random.rand(2, 1))
        cpd3 = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                          evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd1, cpd2, cpd3)
        self.assertEqual(self.graph.get_cpds('diff'), cpd1)
        self.assertEqual(self.graph.get_cpds('intel'), cpd2)
        self.assertEqual(self.graph.get_cpds('grade'), cpd3)

    def test_get_cpd_raises_error(self):
        cpd1 = TabularCPD('diff', 2, values=np.random.rand(2, 1))
        cpd2 = TabularCPD('intel', 2, values=np.random.rand(2, 1))
        cpd3 = TabularCPD('grade', 2, values=np.random.rand(2, 4),
                          evidence=['diff', 'intel'], evidence_card=[2, 2])
        self.graph.add_edges_from([('diff', 'grade'), ('intel', 'grade')])
        self.graph.add_cpds(cpd1, cpd2, cpd3)
        self.assertRaises(ValueError, self.graph.get_cpds, 'sat')

    def tearDown(self):
        del self.graph


from matplotlib import pyplot as plt
import numpy as np
import iterative
import pascal
import power
plt.style.use('ggplot')
qr = []
lu = []

for i in range(2, 13):
    q = pascal.solve_qr_b(pascal.pascal_matrix(i), pascal.harmonic_vector(i))
    l = pascal.solve_lu_b(pascal.pascal_matrix(i), pascal.harmonic_vector(i))
    qr.append(q)
    lu.append(l)

plt.subplot(1, 1, 1)
x = range(2, 13)
y = [i[1] for i in qr]
z = [i[2] for i in qr]

plt.plot(x, y, color='blue')  # error from householder
plt.plot(x, z, color='green')  # solution error of qr
plt.yscale('log')
plt.savefig('./qr_err.png')

y = [i[1] for i in lu]
z = [i[2] for i in lu]
plt.clf()
plt.plot(x, y, color='blue')
plt.plot(x, z, color='green')
plt.yscale('log')
plt.savefig('./lu_err.png')
plt.clf()

jacobi, gs = iterative.generate_data()
j_vals = [i[1] for i in jacobi]
g_vals = [i[1] for i in gs]

jacobi_approx = sum(j_vals) / len(j_vals)  # 2c
gs_approx = sum(g_vals) / len(g_vals)

print("Averages, jacobi then gauss-seidel, then iterations")
print(jacobi_approx)
print(gs_approx)
print(float(sum(j[2] for j in jacobi))/sum(g[2] for g in gs))

exact = np.array([9.0/190, 28.0/475, 33.0/475]).reshape(3,1)
errs_jacobi = [pascal.norm_inf(j-exact) for j in j_vals]
errs_gs = [pascal.norm_inf(g-exact) for g in g_vals]

plt.plot([j[2] for j in jacobi], errs_jacobi, 'ko', [g[2] for g in gs], errs_jacobi, 'bo')
plt.savefig('./iterative_err.png')
plt.clf()

powers = power.generate_data()
ds = [p[0] for p in powers if p[0] is not None]
ts = [p[1] for p in powers if p[1] is not None]
tis = [p[2] for p in powers if p[2] is not None]
maxs = [p[3] for p in powers if p[3] is not None]
mins = [p[4] for p in powers if p[4] is not None]
big = max(maxs)
small = max(mins)
maxs = [float(m)/big for m in maxs]
mins = [float(m)/small for m in mins]

plt.scatter(ds, ts, c=maxs)
plt.savefig('./power_mat.png')
plt.clf()
plt.scatter([1.0/d for d in ds], tis, c=mins)
plt.savefig('./power_inv.png')
plt.clf()


"""
=================================
Compare BIRCH and MiniBatchKMeans
=================================

This example compares the timing of Birch (with and without the global
clustering step) and MiniBatchKMeans on a synthetic dataset having
100,000 samples and 2 features generated using make_blobs.

If ``n_clusters`` is set to None, the data is reduced from 100,000
samples to a set of 158 clusters. This can be viewed as a preprocessing
step before the final (global) clustering step that further reduces these
158 clusters to 100 clusters.
"""

# Authors: Manoj Kumar <manojkumarsivaraj334@gmail.com
#          Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
# License: BSD 3 clause

print(__doc__)

from itertools import cycle
from time import time
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors

from sklearn.preprocessing import StandardScaler
from sklearn.cluster import Birch, MiniBatchKMeans
from sklearn.datasets.samples_generator import make_blobs

# Generate centers for the blobs so that it forms a 10 X 10 grid.
xx = np.linspace(-22, 22, 10)
yy = np.linspace(-22, 22, 10)
xx, yy = np.meshgrid(xx, yy)
n_centres = np.hstack((np.ravel(xx)[:, np.newaxis],
                       np.ravel(yy)[:, np.newaxis]))

# Generate blobs to do a comparison between MiniBatchKMeans and Birch.
X, y = make_blobs(n_samples=100000, centers=n_centres, random_state=0)

# Use all colors that matplotlib provides by default.
colors_ = cycle(colors.cnames.keys())

fig = plt.figure(figsize=(12, 4))
fig.subplots_adjust(left=0.04, right=0.98, bottom=0.1, top=0.9)

# Compute clustering with Birch with and without the final clustering step
# and plot.
birch_models = [Birch(threshold=1.7, n_clusters=None),
                Birch(threshold=1.7, n_clusters=100)]
final_step = ['without global clustering', 'with global clustering']

for ind, (birch_model, info) in enumerate(zip(birch_models, final_step)):
    t = time()
    birch_model.fit(X)
    time_ = time() - t
    print("Birch %s as the final step took %0.2f seconds" % (
          info, (time() - t)))

    # Plot result
    labels = birch_model.labels_
    centroids = birch_model.subcluster_centers_
    n_clusters = np.unique(labels).size
    print("n_clusters : %d" % n_clusters)

    ax = fig.add_subplot(1, 3, ind + 1)
    for this_centroid, k, col in zip(centroids, range(n_clusters), colors_):
        mask = labels == k
        ax.plot(X[mask, 0], X[mask, 1], 'w',
                markerfacecolor=col, marker='.')
        if birch_model.n_clusters is None:
            ax.plot(this_centroid[0], this_centroid[1], '+', markerfacecolor=col,
                    markeredgecolor='k', markersize=5)
    ax.set_ylim([-25, 25])
    ax.set_xlim([-25, 25])
    ax.set_autoscaley_on(False)
    ax.set_title('Birch %s' % info)

# Compute clustering with MiniBatchKMeans.
mbk = MiniBatchKMeans(init='k-means++', n_clusters=100, batch_size=100,
                      n_init=10, max_no_improvement=10, verbose=0,
                      random_state=0)
t0 = time()
mbk.fit(X)
t_mini_batch = time() - t0
print("Time taken to run MiniBatchKMeans %0.2f seconds" % t_mini_batch)
mbk_means_labels_unique = np.unique(mbk.labels_)

ax = fig.add_subplot(1, 3, 3)
for this_centroid, k, col in zip(mbk.cluster_centers_,
                                 range(n_clusters), colors_):
    mask = mbk.labels_ == k
    ax.plot(X[mask, 0], X[mask, 1], 'w', markerfacecolor=col, marker='.')
    ax.plot(this_centroid[0], this_centroid[1], '+', markeredgecolor='k',
            markersize=5)
ax.set_xlim([-25, 25])
ax.set_ylim([-25, 25])
ax.set_title("MiniBatchKMeans")
ax.set_autoscaley_on(False)
plt.show()


"""
=====================
Lasso and Elastic Net
=====================

Lasso and elastic net (L1 and L2 penalisation) implemented using a
coordinate descent.

The coefficients can be forced to be positive.
"""
print(__doc__)

# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
# License: BSD 3 clause

import numpy as np
import matplotlib.pyplot as plt

from sklearn.linear_model import lasso_path, enet_path
from sklearn import datasets

diabetes = datasets.load_diabetes()
X = diabetes.data
y = diabetes.target

X /= X.std(axis=0)  # Standardize data (easier to set the l1_ratio parameter)

# Compute paths

eps = 5e-3  # the smaller it is the longer is the path

print("Computing regularization path using the lasso...")
alphas_lasso, coefs_lasso, _ = lasso_path(X, y, eps, fit_intercept=False)

print("Computing regularization path using the positive lasso...")
alphas_positive_lasso, coefs_positive_lasso, _ = lasso_path(
    X, y, eps, positive=True, fit_intercept=False)
print("Computing regularization path using the elastic net...")
alphas_enet, coefs_enet, _ = enet_path(
    X, y, eps=eps, l1_ratio=0.8, fit_intercept=False)

print("Computing regularization path using the positve elastic net...")
alphas_positive_enet, coefs_positive_enet, _ = enet_path(
    X, y, eps=eps, l1_ratio=0.8, positive=True, fit_intercept=False)

# Display results

plt.figure(1)
ax = plt.gca()
ax.set_color_cycle(2 * ['b', 'r', 'g', 'c', 'k'])
l1 = plt.plot(-np.log10(alphas_lasso), coefs_lasso.T)
l2 = plt.plot(-np.log10(alphas_enet), coefs_enet.T, linestyle='--')

plt.xlabel('-Log(alpha)')
plt.ylabel('coefficients')
plt.title('Lasso and Elastic-Net Paths')
plt.legend((l1[-1], l2[-1]), ('Lasso', 'Elastic-Net'), loc='lower left')
plt.axis('tight')

plt.figure(2)
ax = plt.gca()
ax.set_color_cycle(2 * ['b', 'r', 'g', 'c', 'k'])
l1 = plt.plot(-np.log10(alphas_lasso), coefs_lasso.T)
l2 = plt.plot(-np.log10(alphas_positive_lasso), coefs_positive_lasso.T,
              linestyle='--')

plt.xlabel('-Log(alpha)')
plt.ylabel('coefficients')
plt.title('Lasso and positive Lasso')
plt.legend((l1[-1], l2[-1]), ('Lasso', 'positive Lasso'), loc='lower left')
plt.axis('tight')

plt.figure(3)
ax = plt.gca()
ax.set_color_cycle(2 * ['b', 'r', 'g', 'c', 'k'])
l1 = plt.plot(-np.log10(alphas_enet), coefs_enet.T)
l2 = plt.plot(-np.log10(alphas_positive_enet), coefs_positive_enet.T,
              linestyle='--')

plt.xlabel('-Log(alpha)')
plt.ylabel('coefficients')
plt.title('Elastic-Net and positive Elastic-Net')
plt.legend((l1[-1], l2[-1]), ('Elastic-Net', 'positive Elastic-Net'),
           loc='lower left')
plt.axis('tight')
plt.show()


'''
receiver runs the ADC and photoresistor to receive an input signal.

USes MCP3008 ADC via the hardware SPI interface.

Connections are:

    MCP3008 VDD -> 3.3V (red)
    MCP3008 VREF -> 3.3V (red)
    MCP3008 AGND -> GND (orange)
    MCP3008 CLK -> SCLK (yellow)
    MCP3008 DOUT -> MISO (green)
    MCP3008 DIN -> MOSI (yellow)
    MCP3008 CS -> CE0 (red)
    MCP3008 DGND -> GND (orange)

The photoresistor goes from 4 kOhms (dark) to like 90 Ohms. (flashlight).
Output is 1024*Vin/Vref.

Build a voltage divider with like a 200 Ohm resistor in series w/ the photoR and measure
Vout between them. I put photoresistor between vout and ground.

The signal is intended to be processed using signal_processor

'''
import time

import numpy
import matplotlib
matplotlib.use('Agg')  # works headless (e.g. on Raspberry Pi)
import matplotlib.pyplot as plt
try:
    import spidev
except ImportError:
    print('no spidev')

GAP = 0.001

class ADC(object):
    """
    The Analog-to-digital converter
    """
    def __init__(self):
        self.adc = None

    def __enter__(self):
        self.adc = spidev.SpiDev()
        self.adc.open(0, 0)

    def __exit__(self, exc_type, exc_value, traceback):
        self.adc.close()

    def read(self, input_number):
        """
        read SPI data from MCP3008 chip

        There are 8 possible channels (0 through 7)

        Will return value between 0 and 1023
        """
        if ((input_number > 7) or (input_number < 0)):
            return -1

        r = self.adc.xfer2([1, (8 + input_number) << 4, 0])
        adcValue = ((r[1] & 3) << 8) + r[2]

        return adcValue

class Receiver(object):
    """
    Stream processor that uses adc
    """
    @property
    def times(self):
        return numpy.linspace(0, 10, len(self.vals))

    def receive(self, adc):
        self.vals = []
        # receive for 10 seconds
        print('Receiving')
        start = time.time()
        while time.time() - start < 30.0:
            self.vals.append(adc.read(0))
            time.sleep(GAP / 10)

    def plot(self, fname='adc.pdf'):
        print('Plotting')
        t = self.times
        plt.figure(figsize=(12, 10))
        plt.plot(t, self.vals, '-')
        plt.xlabel('Time (s)')
        plt.ylabel('ADC signal')
        plt.title('ADC Signal Trace')
        plt.grid(color='0.7')
        if fname:
            plt.savefig(fname)

    def save(self, fname='adc.txt'):
        """
        Save results to file
        """
        print('Saving')
        with open(fname, 'w') as f:
            f.writelines(['{0:04d}\n'.format(vi) for vi in self.vals])

    def load(self, fname='adc.txt'):
        print('Loading')
        with open(fname) as f:
            vals = f.readlines()
        self.vals = [float(vi) for vi in vals]

if __name__ == '__main__':
    adc = ADC()
    receiver = Receiver()
    with adc:
        vals = receiver.receive(adc)
    receiver.plot()
    receiver.save()


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for `DataFeeder`."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os.path
import numpy as np
import six
from six.moves import xrange  # pylint: disable=redefined-builtin

# pylint: disable=wildcard-import
from tensorflow.contrib.learn.python.learn.learn_io import *
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.lib.io import file_io
from tensorflow.python.platform import test

# pylint: enable=wildcard-import

class DataFeederTest(test.TestCase):
  # pylint: disable=undefined-variable
  """Tests for `DataFeeder`."""

  def setUp(self):
    self._base_dir = os.path.join(self.get_temp_dir(), 'base_dir')
    file_io.create_dir(self._base_dir)

  def tearDown(self):
    file_io.delete_recursively(self._base_dir)

  def _wrap_dict(self, data, prepend=''):
    return {prepend + '1': data, prepend + '2': data}

  def _assert_raises(self, input_data):
    with self.assertRaisesRegexp(TypeError, 'annot convert'):
      data_feeder.DataFeeder(input_data, None, n_classes=0, batch_size=1)

  def _assert_dtype(self, expected_np_dtype, expected_tf_dtype, input_data):
    feeder = data_feeder.DataFeeder(input_data, None, n_classes=0, batch_size=1)
    if isinstance(input_data, dict):
      for v in list(feeder.input_dtype.values()):
        self.assertEqual(expected_np_dtype, v)
    else:
      self.assertEqual(expected_np_dtype, feeder.input_dtype)
    with ops.Graph().as_default() as g, self.session(g):
      inp, _ = feeder.input_builder()
      if isinstance(inp, dict):
        for v in list(inp.values()):
          self.assertEqual(expected_tf_dtype, v.dtype)
      else:
        self.assertEqual(expected_tf_dtype, inp.dtype)

  def test_input_int8(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.int8)
    self._assert_dtype(np.int8, dtypes.int8, data)
    self._assert_dtype(np.int8, dtypes.int8, self._wrap_dict(data))

  def test_input_int16(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.int16)
    self._assert_dtype(np.int16, dtypes.int16, data)
    self._assert_dtype(np.int16, dtypes.int16, self._wrap_dict(data))

  def test_input_int32(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.int32)
    self._assert_dtype(np.int32, dtypes.int32, data)
    self._assert_dtype(np.int32, dtypes.int32, self._wrap_dict(data))

  def test_input_int64(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.int64)
    self._assert_dtype(np.int64, dtypes.int64, data)
    self._assert_dtype(np.int64, dtypes.int64, self._wrap_dict(data))

  def test_input_uint32(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.uint32)
    self._assert_dtype(np.uint32, dtypes.uint32, data)
    self._assert_dtype(np.uint32, dtypes.uint32, self._wrap_dict(data))

  def test_input_uint64(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.uint64)
    self._assert_dtype(np.uint64, dtypes.uint64, data)
    self._assert_dtype(np.uint64, dtypes.uint64, self._wrap_dict(data))

  def test_input_uint8(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.uint8)
    self._assert_dtype(np.uint8, dtypes.uint8, data)
    self._assert_dtype(np.uint8, dtypes.uint8, self._wrap_dict(data))

  def test_input_uint16(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.uint16)
    self._assert_dtype(np.uint16, dtypes.uint16, data)
    self._assert_dtype(np.uint16, dtypes.uint16, self._wrap_dict(data))

  def test_input_float16(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.float16)
    self._assert_dtype(np.float16, dtypes.float16, data)
    self._assert_dtype(np.float16, dtypes.float16, self._wrap_dict(data))

  def test_input_float32(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.float32)
    self._assert_dtype(np.float32, dtypes.float32, data)
    self._assert_dtype(np.float32, dtypes.float32, self._wrap_dict(data))

  def test_input_float64(self):
    data = np.matrix([[1, 2], [3, 4]], dtype=np.float64)
    self._assert_dtype(np.float64, dtypes.float64, data)
    self._assert_dtype(np.float64, dtypes.float64, self._wrap_dict(data))

  def test_input_bool(self):
    data = np.array([[False for _ in xrange(2)] for _ in xrange(2)])
    self._assert_dtype(np.bool, dtypes.bool, data)
    self._assert_dtype(np.bool, dtypes.bool, self._wrap_dict(data))

  def test_input_string(self):
    input_data = np.array([['str%d' % i for i in xrange(2)] for _ in xrange(2)])
    self._assert_dtype(input_data.dtype, dtypes.string, input_data)
    self._assert_dtype(input_data.dtype, dtypes.string,
                       self._wrap_dict(input_data))

  def _assertAllClose(self, src, dest, src_key_of=None, src_prop=None):

    def func(x):
      val = getattr(x, src_prop) if src_prop else x
      return val if src_key_of is None else src_key_of[val]

    if isinstance(src, dict):
      for k in list(src.keys()):
        self.assertAllClose(func(src[k]), dest)
    else:
      self.assertAllClose(func(src), dest)

  def test_unsupervised(self):

    def func(feeder):
      with self.cached_session():
        inp, _ = feeder.input_builder()
        feed_dict_fn = feeder.get_feed_dict_fn()
        feed_dict = feed_dict_fn()
        self._assertAllClose(inp, [[1, 2]], feed_dict, 'name')

    data = np.matrix([[1, 2], [2, 3], [3, 4]])
    func(data_feeder.DataFeeder(data, None, n_classes=0, batch_size=1))
    func(
        data_feeder.DataFeeder(
            self._wrap_dict(data), None, n_classes=0, batch_size=1))

  def test_data_feeder_regression(self):

    def func(df):
      inp, out = df.input_builder()
      feed_dict_fn = df.get_feed_dict_fn()
      feed_dict = feed_dict_fn()
      self._assertAllClose(inp, [[3, 4], [1, 2]], feed_dict, 'name')
      self._assertAllClose(out, [2, 1], feed_dict, 'name')

    x = np.matrix([[1, 2], [3, 4]])
    y = np.array([1, 2])
    func(data_feeder.DataFeeder(x, y, n_classes=0, batch_size=3))
    func(
        data_feeder.DataFeeder(
            self._wrap_dict(x, 'in'),
            self._wrap_dict(y, 'out'),
            n_classes=self._wrap_dict(0, 'out'),
            batch_size=3))

  def test_epoch(self):

    def func(feeder):
      with self.cached_session():
        feeder.input_builder()
        epoch = feeder.make_epoch_variable()
        feed_dict_fn = feeder.get_feed_dict_fn()
        # First input
        feed_dict = feed_dict_fn()
        self.assertAllClose(feed_dict[epoch.name], [0])
        # Second input
        feed_dict = feed_dict_fn()
        self.assertAllClose(feed_dict[epoch.name], [0])
        # Third input
        feed_dict = feed_dict_fn()
        self.assertAllClose(feed_dict[epoch.name], [0])
        # Back to the first input again, so new epoch.
        feed_dict = feed_dict_fn()
        self.assertAllClose(feed_dict[epoch.name], [1])

    data = np.matrix([[1, 2], [2, 3], [3, 4]])
    labels = np.array([0, 0, 1])
    func(data_feeder.DataFeeder(data, labels, n_classes=0, batch_size=1))
    func(
        data_feeder.DataFeeder(
            self._wrap_dict(data, 'in'),
            self._wrap_dict(labels, 'out'),
            n_classes=self._wrap_dict(0, 'out'),
            batch_size=1))

  def test_data_feeder_multioutput_regression(self):

    def func(df):
      inp, out = df.input_builder()
      feed_dict_fn = df.get_feed_dict_fn()
      feed_dict = feed_dict_fn()
      self._assertAllClose(inp, [[3, 4], [1, 2]], feed_dict, 'name')
      self._assertAllClose(out, [[3, 4], [1, 2]], feed_dict, 'name')

    x = np.matrix([[1, 2], [3, 4]])
    y = np.array([[1, 2], [3, 4]])
    func(data_feeder.DataFeeder(x, y, n_classes=0, batch_size=2))
    func(
        data_feeder.DataFeeder(
            self._wrap_dict(x, 'in'),
            self._wrap_dict(y, 'out'),
            n_classes=self._wrap_dict(0, 'out'),
            batch_size=2))

  def test_data_feeder_multioutput_classification(self):

    def func(df):
      inp, out = df.input_builder()
      feed_dict_fn = df.get_feed_dict_fn()
      feed_dict = feed_dict_fn()
      self._assertAllClose(inp, [[3, 4], [1, 2]], feed_dict, 'name')
      self._assertAllClose(
          out, [[[0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1]],
                [[1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0]]], feed_dict,
          'name')

    x = np.matrix([[1, 2], [3, 4]])
    y = np.array([[0, 1, 2], [2, 3, 4]])
    func(data_feeder.DataFeeder(x, y, n_classes=5, batch_size=2))
    func(
        data_feeder.DataFeeder(
            self._wrap_dict(x, 'in'),
            self._wrap_dict(y, 'out'),
            n_classes=self._wrap_dict(5, 'out'),
            batch_size=2))

  def test_streaming_data_feeder(self):

    def func(df):
      inp, out = df.input_builder()
      feed_dict_fn = df.get_feed_dict_fn()
      feed_dict = feed_dict_fn()
      self._assertAllClose(inp, [[[1, 2]], [[3, 4]]], feed_dict, 'name')
      self._assertAllClose(out, [[[1], [2]], [[2], [2]]], feed_dict, 'name')

    def x_iter(wrap_dict=False):
      yield np.array([[1, 2]]) if not wrap_dict else self._wrap_dict(
          np.array([[1, 2]]), 'in')
      yield np.array([[3, 4]]) if not wrap_dict else self._wrap_dict(
          np.array([[3, 4]]), 'in')

    def y_iter(wrap_dict=False):
      yield np.array([[1], [2]]) if not wrap_dict else self._wrap_dict(
          np.array([[1], [2]]), 'out')
      yield np.array([[2], [2]]) if not wrap_dict else self._wrap_dict(
          np.array([[2], [2]]), 'out')

    func(
        data_feeder.StreamingDataFeeder(
            x_iter(), y_iter(), n_classes=0, batch_size=2))
    func(
        data_feeder.StreamingDataFeeder(
            x_iter(True),
            y_iter(True),
            n_classes=self._wrap_dict(0, 'out'),
            batch_size=2))
    # Test non-full batches.
    func(
        data_feeder.StreamingDataFeeder(
            x_iter(), y_iter(), n_classes=0, batch_size=10))
    func(
        data_feeder.StreamingDataFeeder(
            x_iter(True),
            y_iter(True),
            n_classes=self._wrap_dict(0, 'out'),
            batch_size=10))

  def test_dask_data_feeder(self):
    if HAS_PANDAS and HAS_DASK:
      x = pd.DataFrame(
          dict(
              a=np.array([.1, .3, .4, .6, .2, .1, .6]),
              b=np.array([.7, .8, .1, .2, .5, .3, .9])))
      x = dd.from_pandas(x, npartitions=2)
      y = pd.DataFrame(dict(labels=np.array([1, 0, 2, 1, 0, 1, 2])))
      y = dd.from_pandas(y, npartitions=2)
      # TODO(ipolosukhin): Remove or restore this.
      # x = extract_dask_data(x)
      # y = extract_dask_labels(y)
      df = data_feeder.DaskDataFeeder(x, y, n_classes=2, batch_size=2)
      inp, out = df.input_builder()
      feed_dict_fn = df.get_feed_dict_fn()
      feed_dict = feed_dict_fn()
      self.assertAllClose(feed_dict[inp.name], [[0.40000001, 0.1],
                                                [0.60000002, 0.2]])
      self.assertAllClose(feed_dict[out.name], [[0., 0., 1.], [0., 1., 0.]])

  # TODO(rohanj): Fix this test by fixing data_feeder. Currently, h5py doesn't
  # support permutation based indexing lookups (More documentation at
  # http://docs.h5py.org/en/latest/high/dataset.html#fancy-indexing)
  def DISABLED_test_hdf5_data_feeder(self):

    def func(df):
      inp, out = df.input_builder()
      feed_dict_fn = df.get_feed_dict_fn()
      feed_dict = feed_dict_fn()
      self._assertAllClose(inp, [[3, 4], [1, 2]], feed_dict, 'name')
      self.assertAllClose(out, [2, 1], feed_dict, 'name')

    try:
      import h5py  # pylint: disable=g-import-not-at-top
      x = np.matrix([[1, 2], [3, 4]])
      y = np.array([1, 2])
      file_path = os.path.join(self._base_dir, 'test_hdf5.h5')
      h5f = h5py.File(file_path, 'w')
      h5f.create_dataset('x', data=x)
      h5f.create_dataset('y', data=y)
      h5f.close()
      h5f = h5py.File(file_path, 'r')
      x = h5f['x']
      y = h5f['y']
      func(data_feeder.DataFeeder(x, y, n_classes=0, batch_size=3))
      func(
          data_feeder.DataFeeder(
              self._wrap_dict(x, 'in'),
              self._wrap_dict(y, 'out'),
              n_classes=self._wrap_dict(0, 'out'),
              batch_size=3))
    except ImportError:
      print("Skipped test for hdf5 since it's not installed.")

class SetupPredictDataFeederTest(DataFeederTest):
  """Tests for `DataFeeder.setup_predict_data_feeder`."""

  def test_iterable_data(self):
    # pylint: disable=undefined-variable

    def func(df):
      self._assertAllClose(six.next(df), [[1, 2], [3, 4]])
      self._assertAllClose(six.next(df), [[5, 6]])

    data = [[1, 2], [3, 4], [5, 6]]
    x = iter(data)
    x_dict = iter([self._wrap_dict(v) for v in iter(data)])
    func(data_feeder.setup_predict_data_feeder(x, batch_size=2))
    func(data_feeder.setup_predict_data_feeder(x_dict, batch_size=2))

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


"""
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 hierarchical 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 hierarchical 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_wrong_arg_memory():
    # Test either if an error is raised when memory is not
    # either a str or a joblib.Memory instance
    rng = np.random.RandomState(0)
    n_samples = 100
    X = rng.randn(n_samples, 50)
    memory = 5
    clustering = AgglomerativeClustering(memory=memory)
    assert_raises(ValueError, clustering.fit, X)

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)


# coding: utf-8
import sys, os
sys.path.append(os.pardir)  # 親ディレクトリのファイルをインポートするための設定
import numpy as np
import matplotlib.pyplot as plt
from simple_convnet import SimpleConvNet
from matplotlib.image import imread
from common.layers import Convolution

def filter_show(filters, nx=4, show_num=16):
    """
    c.f. https://gist.github.com/aidiary/07d530d5e08011832b12#file-draw_weight-py
    """
    FN, C, FH, FW = filters.shape
    ny = int(np.ceil(show_num / nx))

    fig = plt.figure()
    fig.subplots_adjust(left=0, right=1, bottom=0, top=1, hspace=0.05, wspace=0.05)

    for i in range(show_num):
        ax = fig.add_subplot(4, 4, i+1, xticks=[], yticks=[])
        ax.imshow(filters[i, 0], cmap=plt.cm.gray_r, interpolation='nearest')

network = SimpleConvNet(input_dim=(1,28,28),
                        conv_param = {'filter_num':30, 'filter_size':5, 'pad':0, 'stride':1},
                        hidden_size=100, output_size=10, weight_init_std=0.01)

# 学習後の重み
network.load_params("params.pkl")

filter_show(network.params['W1'], 16)

img = imread('../dataset/lena_gray.png')
img = img.reshape(1, 1, *img.shape)

fig = plt.figure()

w_idx = 1

for i in range(16):
    w = network.params['W1'][i]
    b = 0  # network.params['b1'][i]

    w = w.reshape(1, *w.shape)
    #b = b.reshape(1, *b.shape)
    conv_layer = Convolution(w, b)
    out = conv_layer.forward(img)
    out = out.reshape(out.shape[2], out.shape[3])

    ax = fig.add_subplot(4, 4, i+1, xticks=[], yticks=[])
    ax.imshow(out, cmap=plt.cm.gray_r, interpolation='nearest')

plt.show()


import sys, os
from copy import deepcopy
import warnings
import pickle, json
import numpy as np
try:
    import h5py
except:
    pass
try:
    from sedpy.observate import load_filters
except:
    pass

"""Convenience functions for reading and reconstructing results from a fitting
run, including reconstruction of the model for making posterior samples
"""

__all__ = ["results_from", "emcee_restarter",
           "get_sps", "get_model",
           "traceplot", "subcorner",
           "compare_paramfile"]

def unpick(pickled):
    """create a serialized object that can go into hdf5 in py2 and py3, and can be read by both
    """
    try:
        obj = pickle.loads(pickled, encoding='bytes')
    except(TypeError):
        obj = pickle.loads(pickled)

    return obj

def results_from(filename, model_file=None, dangerous=True, **kwargs):
    """Read a results file with stored model and MCMC chains.

    :param filename:
        Name and path to the file holding the results.  If ``filename`` ends in
        "h5" then it is assumed that this is an HDF5 file, otherwise it is
        assumed to be a pickle.

    :param dangerous: (default, True)
        If True, use the stored paramfile text to import the parameter file and
        reconstitute the model object.  This executes code in the stored
        paramfile text during import, and is therefore dangerous.

    :returns results:
        A dictionary of various results including:
          + `"chain"`  - Samples from the posterior probability (ndarray).
          + `"lnprobability"` - The posterior probability of each sample.
          + `"weights"` -  The weight of each sample, if `dynesty` was used.
          + `"theta_labels"` - List of strings describing free parameters.
          + `"bestfit"` - The prediction of the data for the posterior sample with
            the highest `"lnprobability"`, as a dictionary.
          + `"run_params"` - A dictionary of arguments supplied to prospector at
            the time of the fit.
          + `"paramfile_text"` - Text of the file used to run prospector, string

    :returns obs:
        The obs dictionary

    :returns model:
        The models.SedModel() object, if it could be regenerated from the stored
        `"paramfile_text"`.  Otherwise, `None`.

    """
    # Read the basic chain, parameter, and run_params info
    if filename.split('.')[-1] == 'h5':
        res = read_hdf5(filename, **kwargs)
        if "_mcmc.h5" in filename:
            mf_default = filename.replace('_mcmc.h5', '_model')
        else:
            mf_default = "x"
    else:
        with open(filename, 'rb') as rf:
            res = pickle.load(rf)
        mf_default = filename.replace('_mcmc', '_model')

    # Now try to read the model object itself from a pickle
    if model_file is None:
        mname = mf_default
    else:
        mname = model_file
    param_file = (res['run_params'].get('param_file', ''),
                  res.get("paramfile_text", ''))
    model, powell_results = read_model(mname, param_file=param_file,
                                       dangerous=dangerous, **kwargs)
    if dangerous:
        try:
            model = get_model(res)
        except:
            model = None
    res['model'] = model
    if powell_results is not None:
        res["powell_results"] = powell_results

    return res, res["obs"], model

def emcee_restarter(restart_from="", niter=32, **kwargs):
    """Get the obs, model, and sps objects from a previous run, as well as the
    run_params and initial positions (which are determined from the end of the
    last run, and inserted into the run_params dictionary)

    :param restart_from:
        Name of the file to restart the sampling from.  An error is raised if
        this does not include an emcee style chain of shape (nwalker, niter,
        ndim)

    :param niter: (default: 32)
        Number of additional iterations to do (added toi run_params)

    :returns obs:
        The `obs` dictionary used in the last run.

    :returns model:
        The model object used in the last run.

    :returns sps:
        The `sps` object used in the last run.

    :returns noise:
        A tuple of (None, None), since it is assumed the noise model in the
        last run was trivial.

    :returns run_params:
        A dictionary of parameters controlling the operation.  This is the same
        as used in the last run, but with the "niter" key changed, and a new
        "initial_positions" key that gives the ending positions of the emcee
        walkers from the last run.  The filename from which the run is
        restarted is also stored in the "restart_from" key.
    """
    result, obs, model = results_from(restart_from)
    noise = (None, None)

    # check for emcee style outputs
    is_emcee = (len(result["chain"].shape) == 3) & (result["chain"].shape[0] > 1)
    msg = "Result file {} does not have a chain of the proper shape."
    assert is_emcee, msg.format(restart_from)

    sps = get_sps(result)
    run_params = deepcopy(result["run_params"])
    run_params["niter"] = niter
    run_params["restart_from"] = restart_from

    initial_positions = result["chain"][:, -1, :]
    run_params["initial_positions"] = initial_positions

    return obs, model, sps, noise, run_params

def read_model(model_file, param_file=('', ''), dangerous=False, **extras):
    """Read the model pickle.  This can be difficult if there are user defined
    functions that have to be loaded dynamically.  In that case, import the
    string version of the paramfile and *then* try to unpickle the model
    object.

    :param model_file:
        String, name and path to the model pickle.

    :param dangerous: (default: False)
        If True, try to import the given paramfile.

    :param param_file:
        2-element tuple.  The first element is the name of the paramfile, which
        will be used to set the name of the imported module.  The second
        element is the param_file contents as a string.  The code in this
        string will be imported.
    """
    model = powell_results = None
    if os.path.exists(model_file):
        try:
            with open(model_file, 'rb') as mf:
                mod = pickle.load(mf)
        except(AttributeError):
            # Here one can deal with module and class names that changed
            with open(model_file, 'rb') as mf:
                mod = load(mf)
        except(ImportError, KeyError):
            # here we load the parameter file as a module using the stored
            # source string.  Obviously this is dangerous as it will execute
            # whatever is in the stored source string.  But it can be used to
            # recover functions (especially dependcy functions) that are user
            # defined
            path, filename = os.path.split(param_file[0])
            modname = filename.replace('.py', '')
            if dangerous:
                user_module = import_module_from_string(param_file[1], modname)
            with open(model_file, 'rb') as mf:
                mod = pickle.load(mf)

        model = mod['model']

        for k, v in list(model.theta_index.items()):
            if type(v) is tuple:
                model.theta_index[k] = slice(*v)
        powell_results = mod['powell']

    return model, powell_results

def read_hdf5(filename, **extras):
    """Read an HDF5 file (with a specific format) into a dictionary of results.

    This HDF5 file is assumed to have the groups ``sampling`` and ``obs`` which
    respectively contain the sampling chain and the observational data used in
    the inference.

    All attributes of these groups as well as top-level attributes are loaded
    into the top-level of the dictionary using ``json.loads``, and therefore
    must have been written with ``json.dumps``.  This should probably use
    JSONDecoders, but who has time to learn that.

    :param filename:
        Name of the HDF5 file.
    """
    groups = {"sampling": {}, "obs": {},
              "bestfit": {}, "optimization": {}}
    res = {}
    with h5py.File(filename, "r") as hf:
        # loop over the groups
        for group, d in groups.items():
            # check the group exists
            if group not in hf:
                continue
            # read the arrays in that group into the dictionary for that group
            for k, v in hf[group].items():
                d[k] = np.array(v)
            # unserialize the attributes and put them in the dictionary
            for k, v in hf[group].attrs.items():
                try:
                    d[k] = json.loads(v)
                except:
                    try:
                        d[k] = unpick(v)
                    except:
                        d[k] = v
        # do top-level attributes.
        for k, v in hf.attrs.items():
            try:
                res[k] = json.loads(v)
            except:
                try:
                    res[k] = unpick(v)
                except:
                    res[k] = v
        res.update(groups['sampling'])
        res["bestfit"] = groups["bestfit"]
        res["optimization"] = groups["optimization"]
        res['obs'] = groups['obs']
        try:
            res['obs']['filters'] = load_filters([str(f) for f in res['obs']['filters']])
        except:
            pass
        try:
            res['rstate'] = unpick(res['rstate'])
        except:
            pass
        #try:
        #    mp = [names_to_functions(p.copy()) for p in res['model_params']]
        #    res['model_params'] = mp
        #except:
        #    pass

    return res

def read_pickles(filename, **kwargs):
    """Alias for backwards compatability. Calls `results_from()`.
    """
    return results_from(filename, **kwargs)

def get_sps(res):
    """This gets exactly the SPS object used in the fiting (modulo any
    changes to FSPS itself).

    It (scarily) imports the paramfile (stored as text in the results
    dictionary) as a module and then uses the `load_sps` method defined in the
    paramfile module.

    :param res:
        A results dictionary (the output of `results_from()`)

    :returns sps:
        An sps object (i.e. from prospect.sources)
    """
    import os
    param_file = (res['run_params'].get('param_file', ''),
                  res.get("paramfile_text", ''))
    path, filename = os.path.split(param_file[0])
    modname = filename.replace('.py', '')
    user_module = import_module_from_string(param_file[1], modname)
    try:
        sps = user_module.load_sps(**res['run_params'])
    except(AttributeError):
        sps = user_module.build_sps(**res['run_params'])

    # Now check that the SSP libraries are consistent
    flib = res['run_params'].get('sps_libraries', None)
    try:
        rlib = sps.ssp.libraries
    except(AttributeError):
        rlib = None
    if (flib is None) or (rlib is None):
        warnings.warn("Could not check SSP library versions.")
    else:
        liberr = ("The FSPS libraries used in fitting({}) are not the "
                  "same as the FSPS libraries that you are using now ({})".format(flib, rlib))
        # If fitting and reading in are happening in different python versions,
        # ensure string comparison doesn't throw error:
        if type(flib[0]) == 'bytes':
            flib = [i.decode() for i in flib]
        if type(rlib[0]) == 'bytes':
            rlib = [i.decode() for i in rlib]
        assert (flib[0] == rlib[0]) and (flib[1] == rlib[1]), liberr

    return sps

def get_model(res):
    """This gets exactly the model object used in the fiting.

    It (scarily) imports the paramfile (stored as text in the results
    dictionary) as a module and then uses the `load_model` method defined in the
    paramfile module, with `run_params` dictionary passed to it.

    :param res:
        A results dictionary (the output of `results_from()`)

    :returns model:
        A prospect.models.SedModel object
    """
    import os
    param_file = (res['run_params'].get('param_file', ''),
                  res.get("paramfile_text", ''))
    path, filename = os.path.split(param_file[0])
    modname = filename.replace('.py', '')
    user_module = import_module_from_string(param_file[1], modname)
    try:
        model = user_module.load_model(**res['run_params'])
    except(AttributeError):
        model = user_module.build_model(**res['run_params'])
    return model

def import_module_from_string(source, name, add_to_sys_modules=True):
    """Well this seems dangerous.
    """
    import imp
    user_module = imp.new_module(name)
    exec(source, user_module.__dict__)
    if add_to_sys_modules:
        sys.modules[name] = user_module

    return user_module

def traceplot(results, showpars=None, start=0, chains=slice(None),
              figsize=None, truths=None, **plot_kwargs):
    """Plot the evolution of each parameter value with iteration #, for each
    walker in the chain.

    :param results:
        A Prospector results dictionary, usually the output of
        ``results_from('resultfile')``.

    :param showpars: (optional)
        A list of strings of the parameters to show.  Defaults to all
        parameters in the ``"theta_labels"`` key of the ``sample_results``
        dictionary.

    :param chains:
        If results are from an ensemble sampler, setting `chain` to an integer
        array of walker indices will cause only those walkers to be used in
        generating the plot.  Useful for to keep the plot from getting too cluttered.

    :param start: (optional, default: 0)
        Integer giving the iteration number from which to start plotting.

    :param **plot_kwargs:
        Extra keywords are passed to the
        ``matplotlib.axes._subplots.AxesSubplot.plot()`` method.

    :returns tracefig:
        A multipaneled Figure object that shows the evolution of walker
        positions in the parameters given by ``showpars``, as well as
        ln(posterior probability)
    """
    import matplotlib.pyplot as pl

    # Get parameter names
    try:
        parnames = np.array(results['theta_labels'])
    except(KeyError):
        parnames = np.array(results['model'].theta_labels())
    # Restrict to desired parameters
    if showpars is not None:
        ind_show = np.array([p in showpars for p in parnames], dtype=bool)
        parnames = parnames[ind_show]
    else:
        ind_show = slice(None)

    # Get the arrays we need (trace, lnp, wghts)
    trace = results['chain'][..., ind_show]
    if trace.ndim == 2:
        trace = trace[None, :]
    trace = trace[chains, start:, :]
    lnp = np.atleast_2d(results['lnprobability'])[chains, start:]
    wghts = results.get('weights', None)
    if wghts is not None:
        wghts = wghts[start:]
    nwalk = trace.shape[0]

    # Set up plot windows
    ndim = len(parnames) + 1
    nx = int(np.floor(np.sqrt(ndim)))
    ny = int(np.ceil(ndim * 1.0 / nx))
    sz = np.array([nx, ny])
    factor = 3.0           # size of one side of one panel
    lbdim = 0.2 * factor   # size of left/bottom margin
    trdim = 0.2 * factor   # size of top/right margin
    whspace = 0.05 * factor         # w/hspace size
    plotdim = factor * sz + factor * (sz - 1) * whspace
    dim = lbdim + plotdim + trdim

    if figsize is None:
        fig, axes = pl.subplots(nx, ny, figsize=(dim[1], dim[0]), sharex=True)
    else:
        fig, axes = pl.subplots(nx, ny, figsize=figsize, sharex=True)
    axes = np.atleast_2d(axes)
    #lb = lbdim / dim
    #tr = (lbdim + plotdim) / dim
    #fig.subplots_adjust(left=lb[1], bottom=lb[0], right=tr[1], top=tr[0],
    #                    wspace=whspace, hspace=whspace)

    # Sequentially plot the chains in each parameter
    for i in range(ndim - 1):
        ax = axes.flat[i]
        for j in range(nwalk):
            ax.plot(trace[j, :, i], **plot_kwargs)
        ax.set_title(parnames[i], y=1.02)
    # Plot lnprob
    ax = axes.flat[-1]
    for j in range(nwalk):
        ax.plot(lnp[j, :], **plot_kwargs)
    ax.set_title('lnP', y=1.02)

    [ax.set_xlabel("iteration") for ax in axes[-1, :]]
    #[ax.set_xticklabels('') for ax in axes[:-1, :].flat]

    if truths is not None:
        for i, t in enumerate(truths[ind_show]):
            axes.flat[i].axhline(t, color='k', linestyle=':')

    pl.tight_layout()
    return fig

def param_evol(results, **kwargs):
    """Backwards compatability
    """
    return traceplot(results, **kwargs)

def subcorner(results, showpars=None, truths=None,
              start=0, thin=1, chains=slice(None),
              logify=["mass", "tau"], **kwargs):
    """Make a triangle plot of the (thinned, latter) samples of the posterior
    parameter space.  Optionally make the plot only for a supplied subset of
    the parameters.

    :param showpars: (optional)
        List of string names of parameters to include in the corner plot.

    :param truths: (optional)
        List of truth values for the chosen parameters.

    :param start: (optional, default: 0)
        The iteration number to start with when drawing samples to plot.

    :param thin: (optional, default: 1)
        The thinning of each chain to perform when drawing samples to plot.

    :param chains: (optional)
        If results are from an ensemble sampler, setting `chain` to an integer
        array of walker indices will cause only those walkers to be used in
        generating the plot.  Useful for emoving stuck walkers.

    :param kwargs:
        Remaining keywords are passed to the ``corner`` plotting package.

    :param logify:
        A list of parameter names to plot in `log10(parameter)` instead of
        `parameter`
    """
    try:
        import corner as triangle
    except(ImportError):
        import triangle
    except:
        raise ImportError("Please install the `corner` package.")

    # pull out the parameter names and flatten the thinned chains
    # Get parameter names
    try:
        parnames = np.array(results['theta_labels'], dtype='U20')
    except(KeyError):
        parnames = np.array(results['model'].theta_labels())
    # Restrict to desired parameters
    if showpars is not None:
        ind_show = np.array([parnames.tolist().index(p) for p in showpars])
        parnames = parnames[ind_show]
    else:
        ind_show = slice(None)

    # Get the arrays we need (trace, wghts)
    trace = results['chain'][..., ind_show]
    if trace.ndim == 2:
        trace = trace[None, :]
    trace = trace[chains, start::thin, :]
    wghts = results.get('weights', None)
    if wghts is not None:
        wghts = wghts[start::thin]
    samples = trace.reshape(trace.shape[0] * trace.shape[1], trace.shape[2])

    # logify some parameters
    xx = samples.copy()
    if truths is not None:
        xx_truth = np.array(truths).copy()
    else:
        xx_truth = None
    for p in logify:
        if p in parnames:
            idx = parnames.tolist().index(p)
            xx[:, idx] = np.log10(xx[:, idx])
            parnames[idx] = "log({})".format(parnames[idx])
            if truths is not None:
                xx_truth[idx] = np.log10(xx_truth[idx])

    # mess with corner defaults
    corner_kwargs = {"plot_datapoints": False, "plot_density": False,
                     "fill_contours": True, "show_titles": True}
    corner_kwargs.update(kwargs)

    fig = triangle.corner(xx, labels=parnames, truths=xx_truth,
                          quantiles=[0.16, 0.5, 0.84], weights=wghts, **corner_kwargs)

    return fig

def subtriangle(results, **kwargs):
    """Backwards compatability
    """
    return subcorner(results, **kwargs)

def compare_paramfile(res, filename):
    """Compare the runtime parameter file text stored in the `res` dictionary
    to the text of some existing file with fully qualified path `filename`.
    """
    from pprint import pprint
    from difflib import unified_diff

    a = res["paramfile_text"]
    aa = a.split('\n')
    with open(filename, "r") as f:
        b = json.dumps(f.read())
    bbl = json.loads(b)
    bb = bbl.split('\n')
    pprint([l for l in unified_diff(aa, bb)])

def names_to_functions(p):
    """Replace names of functions (or pickles of objects) in a parameter
    description with the actual functions (or pickles).
    """
    from importlib import import_module
    for k, v in list(p.items()):
        try:
            m = import_module(v[1])
            f = m.__dict__[v[0]]
        except:
            try:
                f = pickle.loads(v)
            except:
                f = v

        p[k] = f

    return p