import unittest

import scipy as sp

import connectome.model.blockmodel as blmdl


def connectivity(W):
    return sp.count_nonzero(W) / (W.shape[0] * W.shape[1])


def reciprocity(W):
    return sp.count_nonzero(W * W.T) / sp.count_nonzero(W)


class TestBlockModel(unittest.TestCase):
    def test_connectivity_diagonal_only(self):
        conn = .15
        rec = 0.3
        nr_exc = 600
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=600,
                                  nr_inh=200,
                                  blocks=[[1, 0],
                                          [0, 1]])
        adjacency_matrix = block().adjacency_matrix[:nr_exc, :nr_exc]
        calculated_conn = connectivity(adjacency_matrix)
        self.assertLess(abs((calculated_conn - conn) / conn), 0.01)

    def test_connectivity_diagonal_only_unsatisfiable(self):
        conn = .2  # does only work for conn=1.5. so should get an error here
        rec = .3
        nr_nodes = 600
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=200,
                                  blocks=[[1, 0],
                                         [0, 1]])
        with self.assertRaises(blmdl.UnsatisfiableConstraints):
            block()

    def test_reciprocity_diagonal_only(self):
        conn = .3
        rec = .3 * 2
        nr_nodes = 603
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=100,
                                  blocks=[[1, 0],
                                         [0, 1]])
        adjacency_matrix = block().adjacency_matrix[:nr_nodes, :nr_nodes]
        self.assertLess(abs((reciprocity(adjacency_matrix) - rec) / rec), 0.01)

    def test_reciprocity_unsatisfiable(self):
        conn = .3
        rec = .34
        nr_nodes = 603
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=100,
                                  blocks=[[1, 0],
                                         [0, 1]])
        with self.assertRaises(blmdl.UnsatisfiableConstraints):
            block()

    def test_connectivity_off_diagonal_only(self):
        conn = .05
        rec = 0
        nr_nodes = 2000
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=100,
                                  blocks=[[0, 1, 0],
                                         [0, 0, 0],
                                         [0, 0, 0]])
        adjacency_matrix = block().adjacency_matrix[:nr_nodes, :nr_nodes]
        self.assertLess(abs((connectivity(adjacency_matrix) - conn) / conn), 0.01)

    def test_connectivity_off_diagonal_and_diagonal(self):
        conn = .05
        rec = 0
        nr_nodes = 1600
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=100,
                                  blocks=[[1, 1, 0],
                                         [0, 0, 0],
                                         [0, 0, 0]])
        adjacency_matrix = block().adjacency_matrix[:nr_nodes, :nr_nodes]
        self.assertLess(abs((connectivity(adjacency_matrix) - conn) / conn), 0.01)

    def test_reciprocity_off_diagonal_and_diagonal(self):
        conn = .3
        rec = .3
        nr_nodes = 2500
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=100,
                                  blocks=[[1, 1, 0],
                                         [0, 1, 1],
                                         [0, 0, 0]])
        adjacency_matrix = block().adjacency_matrix[:nr_nodes, :nr_nodes]
        self.assertLess(abs((reciprocity(adjacency_matrix) - rec) / rec), 0.01)

    def test_not_upper_triangular_exception(self):
        conn = .3
        rec = .3
        nr_nodes = 600
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=100,
                                  blocks=[[1, 1, 0],
                                                 [0, 1, 0],
                                                 [1, 0, 0]])
        with self.assertRaises(blmdl.NotUpperTriangular):
            block()

    def test_wrong_block_shape(self):
        conn = .3
        rec = .3
        nr_nodes = 600
        block = blmdl.BlockMatrix(exc_connectivity=conn,
                                  exc_reciprocity=rec,
                                  nr_exc=nr_nodes,
                                  nr_inh=100,
                                  blocks=[1])
        with self.assertRaises(ValueError):
            block()


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