import unittest

import scipy as sp
from theano import shared

from connectome.function.kerasextension import BlockVector, BlockMatrix, BT


class TestBlockVector(unittest.TestCase):
    def test_dot(self):
        a = sp.randn(10).astype("float32")
        b = sp.randn(15).astype("float32")
        a_t = shared(a)
        b_t = shared(b)
        v1 = BlockVector([a_t, b_t])
        v2 = BlockVector([a_t, b_t])
        res = v1.dot(v2).eval()
        expected = a.dot(a) + b.dot(b)
        self.assertLess(abs(res-expected), 1e-5)

    def test_add(self):
        a = sp.randn(10).astype("float32")
        b = sp.randn(10).astype("float32")
        a_t = shared(a)
        b_t = shared(b)
        res = (a_t + b_t).eval()
        expected = a + b
        self.assertLess(sp.absolute(res-expected).sum(), 1e-5)

    def test_elemwise_apply(self):
        a = sp.randn(10).astype("float32")
        b = sp.randn(15).astype("float32")
        v = BlockVector([shared(a), shared(b)])
        res = BT.nnet.relu(v).eval().to_dense_array()
        expected = sp.maximum(sp.concatenate((a, b)), 0)
        self.assertLess(sp.absolute(res-expected).sum(), 1e-5)


class TestBlockMatrix(unittest.TestCase):
    def test_dot(self):
        M = [[sp.randn(10, 20), sp.randn(10, 30)],
             [sp.randn(14, 20), sp.randn(14, 30)]]
        V = [sp.randn(20), sp.randn(30)]
        block_vec = BlockVector([shared(v) for v in V])
        block_mat = BlockMatrix([[shared(col) for col in row] for row in M])
        M_SP = sp.concatenate([sp.concatenate([col for col in row], axis=1) for row in M], axis=0)
        V_SP = sp.concatenate(V, axis=0)
        expected = M_SP.dot(V_SP)
        res = block_mat.dot(block_vec).eval().to_dense_array()
        self.assertLess(sp.absolute(expected-res).sum(), 1e-5)

    def test_matmul(self):
        M = [[sp.randn(10, 20), sp.randn(10, 30)],
             [sp.randn(14, 20), sp.randn(14, 30)]]
        V = [sp.randn(20), sp.randn(30)]
        block_vec = BlockVector([shared(v) for v in V])
        block_mat = BlockMatrix([[shared(col) for col in row] for row in M])
        M_SP = sp.concatenate([sp.concatenate([col for col in row], axis=1) for row in M], axis=0)
        V_SP = sp.concatenate(V, axis=0)
        expected = M_SP.dot(V_SP)
        res = (block_mat @ block_vec).eval().to_dense_array()
        self.assertLess(sp.absolute(expected-res).sum(), 1e-5)

    def test_rmatmul(self):
        M = [[sp.randn(10, 20), sp.randn(10, 30)],
             [sp.randn(14, 20), sp.randn(14, 30)]]
        V = [sp.randn(20), sp.randn(30)]
        block_vec = BlockVector([shared(v) for v in V])
        block_mat = BlockMatrix([[shared(col) for col in row] for row in M])
        M_SP = sp.concatenate([sp.concatenate([col for col in row], axis=1) for row in M], axis=0)
        V_SP = sp.concatenate(V, axis=0)
        expected = M_SP.dot(V_SP)
        res = (block_vec @ block_mat.T).eval().to_dense_array()
        self.assertLess(sp.absolute(expected-res).sum(), 1e-5)

    def test_get_diag(self):
        M = [[1, 2, 3],
             [4, 5, 6],
             [7, 8, 9]]
        block_mat = BlockMatrix(M)
        diag = block_mat.get_diag()
        self.assertEqual(list(diag), [1, 5, 9])

    def test_set_diag(self):
        M = [[0, 2, 3],
             [4, 0, 6],
             [7, 8, 0]]
        block_mat = BlockMatrix(M)
        block_mat.set_diag([1, 5, 9])
        self.assertEqual(list(block_mat.get_diag()), [1, 5, 9])

    def test_get_upper_diag(self):
        M = [[1, 2, 3],
             [4, 5, 6],
             [7, 8, 9]]
        block_mat = BlockMatrix(M)
        diag = block_mat.get_diag(1)
        self.assertEqual(list(diag), [2, 6])

    def test_get_diag_single_block(self):
        M = [[42]]
        block_mat = BlockMatrix(M)
        diag = block_mat.get_diag(0)
        self.assertEqual([42], list(diag))

    def test_get_diag_single_block_empty(self):
        M = [[42]]
        block_mat = BlockMatrix(M)
        diag = block_mat.get_diag(1)
        self.assertEqual([], list(diag))

    def test_comparison(self):
        shape = (10, 20)
        a = shared(sp.rand(*shape)).astype('float32')
        b = shared(sp.rand(*shape)).astype('float32')
        c = shared(sp.rand(*shape)).astype('float32')
        d = shared(sp.rand(*shape)).astype('float32')
        m = BlockMatrix([[a, b],
                         [c, d]])
        m_eval = m.eval()
        self.assertTrue(m_eval == m_eval)

    def test_comparison_unequal(self):
        shape = (10, 20)
        a = shared(sp.rand(*shape)).astype('float32')
        b = shared(sp.rand(*shape)).astype('float32')
        c = shared(sp.rand(*shape)).astype('float32')
        d = shared(sp.rand(*shape)).astype('float32')
        m = BlockMatrix([[a, b],
                         [c, d]])
        m1 = BlockMatrix([[a, a],
                         [c, d]])
        self.assertFalse(m.eval() == m1.eval())

    def test_transpose_after_eval(self):
        shape = (5, 7)

        def make():
            return shared((sp.rand(*shape) * 100).astype(int))

        a = make()
        b = make()
        c = make()
        d = make()
        m = BlockMatrix([[a, b],
                         [c, d]])
        self.assertTrue(m.eval().T.T == m.eval())

    def test_transpose_with_int(self):
        m = BlockMatrix([[1, 2],
                         [3, 4]])
        mt = BlockMatrix([[1, 3],
                          [2, 4]])
        self.assertTrue(m.T, m)

    def test_transpose_before_eval(self):
        shape = (5, 7)

        def make():
            return shared((sp.rand(*shape) * 100).astype(int))

        a = make()
        b = make()
        c = make()
        d = make()
        m = BlockMatrix([[a, b],
                         [c, d]])
        self.assertTrue(m.T.T.eval() == m.eval())

    def test_set_upper_diag(self):
        M = [[0, 0, 3],
             [4, 0, 0],
             [7, 8, 0]]
        block_mat = BlockMatrix(M)
        block_mat.set_diag([2, 6], 1)
        self.assertEqual(list(block_mat.get_diag(1)), [2, 6])


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