import unittest

import scipy as sp
from connectome.model.block import Block

from connectome.model.network import Network, PopulationMismatch, EmbeddedNetwork


class TestNetworkInitialization(unittest.TestCase):
        def test_nr_exc(self):
            net = Network(adjacency_matrix=[[0, 2], [3, 0]])
            self.assertEqual(2, net.nr_exc)

        def test_nr_inh(self):
            net = Network(adjacency_matrix=[[0, 2], [3, 0]])
            self.assertEqual(0, net.nr_inh)

        def test_nr_exc_zero_graph(self):
            net = Network(adjacency_matrix=[[0 , 0], [0, 0]])
            self.assertEqual(2, net.nr_exc)

        def test_nr_inh_zero_graph(self):
            net = Network(adjacency_matrix=[[0, 0], [0, 0]])
            self.assertEqual(0, net.nr_inh)

        def test_nr_inh_exc_graph_explicitly_give_number_exc(self):
            net = Network(adjacency_matrix=[[0, 0], [0, 0]], nr_exc=1)
            self.assertEqual(1, net.nr_exc)

        def test_nr_neurons_explicitly_give_number_exc(self):
            net = Network(adjacency_matrix=[[0, 0], [0, 0]], nr_exc=1)
            self.assertEqual(2, net.nr_neurons)

        def test_all_exc(self):
            net = Network(adjacency_matrix=[[0, 0], [0, 0]])
            self.assertEqual(2, net.nr_neurons)

        def test_set_nr_inh(self):
            net = Network(adjacency_matrix=[[0, 0], [0, 0]])
            net.nr_inh = 1
            self.assertEqual(1, net.nr_exc)

        def test_nr_inh_exc_graph_explicitly_give_number_inh(self):
            net = Network(adjacency_matrix=[[0, 0], [0, 0]], nr_exc=1)
            self.assertEqual(1, net.nr_inh)

        def test_nr_neurons(self):
            net = Network(adjacency_matrix=[[0 , 2], [3, 0]])
            self.assertEqual(2, net.nr_neurons)

        def test_assymetric(self):
            net = Network(adjacency_matrix=[[1 , 2], [3, 4], [5, 6]])
            self.assertEqual(2, net.nr_neurons)

        def test_nr_exc_too_large_exception(self):
            with self.assertRaises(PopulationMismatch):
                Network(adjacency_matrix=[[1 , 2], [3, 4], [5, 6]], nr_exc=3)

        def test_nr_exc_too_large_exception_due_to_inh(self):
            with self.assertRaises(Exception):
                net = Network(adjacency_matrix=[[1 , -2], [3, -4], [5, -6]], nr_exc=1)
                net.nr_exc = 2

        def test_negative_nr_exc(self):
            with self.assertRaises(PopulationMismatch):
                net = Network(adjacency_matrix=[[1 , -2], [3, -4], [5, -6]], nr_exc=1)
                net.nr_exc = -1

        def test_initialize_zero(self):
            net = Network(nr_exc=5, nr_inh=2)
            self.assertTrue((sp.zeros((7, 7)) == net.adjacency_matrix).all())

        def test_max_nr_inh(self):
            a = sp.asarray([[1, -1], [1,-1]])
            self.assertEqual(1, Network._max_nr_inh(a))

        def test_max_nr_exc(self):
            a = sp.asarray([[1, -1], [1,-1]])
            self.assertEqual(1, Network._max_nr_exc(a))

        def test_set_adjacency_matrix(self):
            a = sp.asarray([[1, -1], [1,-1]])
            b = sp.asarray([[1, -1], [1,-2]])
            net = Network(adjacency_matrix=a)
            net.adjacency_matrix = b
            self.assertTrue((net.adjacency_matrix==b).all())

        def test_set_adjacency_matrix_bad_inh_nr(self):
            a = sp.asarray([[1, -1], [1,-1]])
            b = sp.asarray([[1, 1], [1, 2]])
            net = Network(adjacency_matrix=a)
            with self.assertRaises(PopulationMismatch):
                net.adjacency_matrix = b

        def test_set_adjacency_matrix_bad_exc_nr(self):
            a = sp.asarray([[1, -1], [1,-1]])
            b = sp.asarray([[-1, 1], [1, 2]])
            net = Network(adjacency_matrix=a)
            with self.assertRaises(PopulationMismatch):
                net.adjacency_matrix = b

        def test_set_adjacency_matrix_bad_nr(self):
            a = sp.asarray([[1, 1], [1, 1]])
            b = sp.asarray([[1, 1, 2], [1, 2, 3]])
            net = Network(adjacency_matrix=a)
            with self.assertRaises(PopulationMismatch):
                net.adjacency_matrix = b

        def test_binarize(self):
            a = sp.asarray([[.4, -.1], [.1,-.1]])
            b = sp.asarray([[1, -1], [1, -1]])
            net = Network(adjacency_matrix=a)
            net.binarize()
            self.assertTrue((net.adjacency_matrix==b).all())

        def test_diag_to_zero(self):
            a = sp.asarray([[.4, -.1], [.1,-.1]])
            net = Network(adjacency_matrix=a)
            net.diag_to_zero()
            self.assertFalse(net._adjacency_matrix.diagonal().any())

        def test_max_nr_exc_all(self):
            a = sp.asarray([[0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0]])
            self.assertEqual(4, Network._max_nr_exc(a))

        def test_max_nr_exc_2(self):
            a = sp.asarray([[0,0,0,-1],
                            [0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0]])
            self.assertEqual(3, Network._max_nr_exc(a))

        def test_max_nr_inh_all(self):
            a = sp.asarray([[0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0]])
            self.assertEqual(4, Network._max_nr_inh(a))

        def test_max_nr_inh_2(self):
            a = sp.asarray([[0,0,0,-1],
                            [0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0]])
            self.assertEqual(4, Network._max_nr_inh(a))

        def test_max_nr_inh_3(self):
            a = sp.asarray([[1,0,0,-1],
                            [0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0]])
            self.assertEqual(3, Network._max_nr_inh(a))

        def test_max_nr_inh_4(self):
            a = sp.asarray([[1,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0],
                            [0,0,0,0]])
            self.assertEqual(3, Network._max_nr_inh(a))


class TestNetworkViews(unittest.TestCase):
    def test_view_all(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        excpected = sp.asarray([[1, 2, -3],
                        [4, 11,  -6],
                        [7, 8,  -9]])
        net = Network(adjacency_matrix=a)
        W = net[:,:]
        W[1,1] = 11
        self.assertEqual((3,3), W.shape)
        self.assertTrue((net.adjacency_matrix == excpected).all())

    def test_view_EE(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        excpected = sp.asarray([[11, 2, -3],
                        [4, 5,  -6],
                        [7, 8,  -9]])
        net = Network(adjacency_matrix=a)
        W = net["E", "E"]
        W[0,0] = 11
        self.assertEqual((2,2), W.shape)
        self.assertTrue((net.adjacency_matrix == excpected).all())

    def test_view_from_e_to_all(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        excpected = sp.asarray([[1, 2],
                                [4, 5],
                                [7, 8]])
        net = Network(adjacency_matrix=a)
        W = net["E",:]
        self.assertTrue((W == excpected).all())

    def test_view_from_i_to_all(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        excpected = sp.asarray([[-3],
                                [-6],
                                [-9]])
        net = Network(adjacency_matrix=a)
        W = net["I",:]
        self.assertTrue((W == excpected).all())

    def test_view_from_all_to_e(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        excpected = sp.asarray([[1, 2, -3],
                                [4, 5, -6]])
        net = Network(adjacency_matrix=a)
        W = net[:, "E"]
        self.assertTrue((W == excpected).all())

    def test_view_from_all_to_i(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        excpected = sp.asarray([[7, 8, -9]])
        net = Network(adjacency_matrix=a)
        W = net[:, "I"]
        self.assertTrue((W == excpected).all())

    def test_view_II(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        expected = sp.asarray([[1, 2, -3],
                        [4, 5,  -6],
                        [7, 8,  -42]])
        net = Network(adjacency_matrix=a)
        W = net["I", "I"]
        W[0,0] = -42
        self.assertEqual((1,1), W.shape)
        self.assertTrue((net.adjacency_matrix == expected).all())

    def test_view_EI(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        expected = sp.asarray([[1, 2,  -3],
                        [4, 5,  -42],
                        [7, 8, -9]])
        net = Network(a)
        W = net["I", "E"]
        W[1,0] = -42
        self.assertEqual((2,1), W.shape)
        self.assertTrue((net.adjacency_matrix == expected).all())

    def test_view_IE(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        expected = sp.asarray([[1, 2,  -3],
                        [4, 5,  -6],
                        [7, 42, -9]])
        net = Network(a)
        W = net["E", "I"]
        W[0, 1] = 42
        self.assertEqual((1,2), W.shape)
        self.assertTrue((net.adjacency_matrix == expected).all())

    def test_set_view_E_to_all(self):
        a = sp.asarray([[1, 2, -3],
                        [4, 5, -6],
                        [7, 8, -9]])
        expected = sp.asarray([[11, 12, -3],
                               [14, 15, -6],
                               [17, 18, -9]])
        net = Network(a)
        net["E", :] = sp.asarray([[11, 12],
                                  [14, 15],
                                  [17, 18]])
        self.assertTrue((net.adjacency_matrix == expected).all())


class TestRelabelNonEmbedded(unittest.TestCase):
    def test_random_relabel_preserves_inh_and_exc(self):
        a = sp.rand(100, 100)
        a[:,-20:] *= -1
        net = Network(a)
        before = net.adjacency_matrix.copy()
        net.random_relabel_neurons()
        after = net.adjacency_matrix
        self.assertFalse((before == after).all())
        self.assertTrue((after[:,:-20] >= 0).all())
        self.assertTrue((after[:,-20:] <= 0).all())


class TestEmbeddedNetwork(unittest.TestCase):
    def test_random_relabel_preserves_inh_and_exc(self):
        a = sp.rand(100, 100)
        a[:,-20:] *= -1
        net = EmbeddedNetwork(a)
        before = net.adjacency_matrix.copy()
        net.random_relabel_neurons()
        after = net.adjacency_matrix
        self.assertFalse((before == after).all())
        self.assertTrue((after[:,:-20] >= 0).all())
        self.assertTrue((after[:,-20:] <= 0).all())

    def test_random_relabel_shuffles_positions_and_adjacency_simultaneously(self):
        n = 10
        n_half = int(n/2)
        a = sp.arange(n*n).reshape((n, n))
        a[:,-n_half:] *= -1
        net = EmbeddedNetwork(a)
        net.positions = sp.arange(n*3).reshape((n, 3))

        old_post, old_pre = sp.where(a == 1234)   # 1234 is just a random number
        old_location = net.positions[ old_pre]
        net.random_relabel_neurons()
        new_post, new_pre = sp.where(net.adjacency_matrix == 1234)
        new_location = net.positions[ new_pre]
        self.assertTrue((old_location==new_location).all())

    def test_wrong_number_positions(self):
        with self.assertRaises(PopulationMismatch):
            net = EmbeddedNetwork(sp.rand(3,3))
            net.positions = sp.rand(4,4)


class TestNetworkComparison(unittest.TestCase):
    def test_compare_self(self):
        net1 = Network([[1,2,-3],
                        [1,2,-3],
                        [1,2,-3]], nr_exc=2, nr_inh=1)
        self.assertEqual(net1, net1)

    def test_compare_to_copy(self):
        net1 = Network([[1,2,-3],
                        [1,2,-3],
                        [1,2,-3]], nr_exc=2, nr_inh=1)
        net2 = net1.copy()
        self.assertEqual(net1, net2)

    def test_compare_embedded_to_unenbedded(self):
        net1 = Network([[1,2,-3],
                        [1,2,-3],
                        [1,2,-3]], nr_exc=2, nr_inh=1)
        net2 = EmbeddedNetwork(net1.adjacency_matrix, nr_exc=2, nr_inh=1)
        self.assertNotEqual(net1, net2)

    def test_compare_to_modified(self):
        net1 = Network([[1,2,-3],
                        [1,2,-3],
                        [1,2,-3]], nr_exc=2, nr_inh=1)
        net2 = net1.copy()
        net2.adjacency_matrix[2,2] *= -1
        self.assertNotEqual(net1, net2)

    def test_compare_to_different_exc_inh_number(self):
        net1 = Network([[0,0,0],
                        [0,0,0],
                        [0,0,0]], nr_exc=2, nr_inh=1)
        net2 = Network([[0,0,0],
                        [0,0,0],
                        [0,0,0]], nr_exc=1, nr_inh=2)
        self.assertNotEqual(net1, net2)


class TestRemoveNeurons(unittest.TestCase):
    def test_remove_neurons_not_embedded(self):
        net = Network([[0,1,0],
                       [1,1,1],
                       [0,1,0]], nr_exc=2, nr_inh=1)
        net.remove_neurons([1])
        self.assertTrue((net.adjacency_matrix == sp.asarray([[0,0],
                                                             [0,0]])).all())
        self.assertEqual(net.nr_exc, 1)
        self.assertEqual(net.nr_inh, 1)

    def test_remove_neurons_only_inh_non_embedded(self):
        net = Network([[0,0,1],
                       [0,0,1],
                       [1,1,1]], nr_exc=2, nr_inh=1)
        net.remove_neurons([2])
        self.assertEqual(net.nr_exc, 2)
        self.assertEqual(net.nr_inh, 0)
        self.assertTrue((net.adjacency_matrix == sp.asarray([[0,0],
                                                             [0,0]])).all())

    def test_remove_neurons_embedded(self):
        net = EmbeddedNetwork([[0,0,1],
                               [0,0,1],
                               [1,1,1]], nr_exc=2, nr_inh=1)
        positions = net.positions
        positions[2] += 1
        self.assertTrue((positions==sp.asarray([[0,0,0],
                                                [0,0,0],
                                                [1,1,1]])).all())
        net.remove_neurons([2])
        self.assertTrue((net.adjacency_matrix == sp.asarray([[0,0],
                                                             [0,0]])).all())
        self.assertTrue((net.positions==sp.asarray([[0,0,0],
                                                    [0,0,0]])).all())
        self.assertEqual(net.nr_exc, 2)
        self.assertEqual(net.nr_inh, 0)


class TestNormalize(unittest.TestCase):
    def test_normalize_cols(self):
        net = Network([[1,1, -2, -3],
                       [1,13, -2, -30],
                       [2,1, -1, -3],
                       [1,11, -2, -3]])
        net.normalize("out")
        self.assertLess(sp.absolute(net.adjacency_matrix[:, :2].sum(axis=0) - 1).sum(), 1e-5)
        self.assertLess(sp.absolute(net.adjacency_matrix[:, 2:].sum(axis=0) - (-1)).sum() , 1e-5)

    def test_normalize_rows(self):
        net = Network([[1,2, -2, -3],
                       [1,1, -2, -13],
                       [12,1, -1, -3],
                       [1,2, -2, -3]])
        net.normalize("in")
        self.assertLess(sp.absolute(net.adjacency_matrix[:, :2].sum(axis=1) - 1).sum(), 1e-5)
        self.assertLess(sp.absolute(net.adjacency_matrix[:, 2:].sum(axis=1) - (-1)).sum(), 1e-5)


class TestBlocks(unittest.TestCase):
    def test_blocks(self):
        net = Network([[1, 2, -2, -3],
                       [1, 1, -2, -13],
                       [12, 1, -1, -3],
                       [1, 2, -2, -3]], blocks=[Block(1, "E"), Block(1, "E"), Block(2, "I")])
        self.assertEqual([Block(1, "E"), Block(1, "E"), Block(2, "I")], net.blocks)

    def test_blocks_inferred(self):
        net = Network([[1, 2, -2, -3],
                       [1, 1, -2, -13],
                       [12, 1, -1, -3],
                       [1, 2, -2, -3]])
        self.assertEqual([Block(2, "E"), Block(2, "I")], net.blocks)

    def test_blocks_relabel(self):
        net = Network([[1, 2, -2, -3],
                       [1, 1, -2, -13],
                       [12, 1, -1, -3],
                       [1, 2, -2, -3]], blocks=[Block(1, "E"), Block(1, "E"), Block(2, "I")])
        net.random_relabel_neurons()
        self.assertEqual([Block(2, "E"), Block(2, "I")], net.blocks)

    def test_start_end_1(self):
        net = Network([[1, 2, -3, -4],
                       [5, 6, -7, -8],
                       [9, 10, -11, -12],
                       [13, 14, -15, -16]], blocks=[Block(1, "E"), Block(1, "E"), Block(2, "I")])
        block = net.get_block(source_ix=0, target_ix=0)
        self.assertTrue((block==1).all())

    def test_start_end_2(self):
        net = Network([[1, 2, -3, -4],
                       [5, 6, -7, -8],
                       [9, 10, -11, -12],
                       [13, 14, -15, -16]], blocks=[Block(1, "E"), Block(1, "E"), Block(2, "I")])
        block = net.get_block(source_ix=1, target_ix=0)
        self.assertTrue((block==2).all())

    def test_start_end_3(self):
        net = Network([[1, 2, -3, -4],
                       [5, 6, -7, -8],
                       [9, 10, -11, -12],
                       [13, 14, -15, -16]], blocks=[Block(2, "E"),Block(2, "I")])
        block = net.get_block(source_ix=0, target_ix=1)
        self.assertTrue((block==sp.array([[9, 10],
                                          [13, 14]])).all())

    def test_start_end_4(self):
        net = Network([[1, 2, -3, -4],
                       [5, 6, -7, -8],
                       [9, 10, -11, -12],
                       [13, 14, -15, -16]], blocks=[Block(2, "E"),Block(2, "I")])
        block = net.get_block(source_ix=1, target_ix=1)
        self.assertTrue((block==sp.array([[-11, -12],
                                          [-15, -16]])).all())



if __name__ == "__main__":
    unittest.main()
