import unittest

import scipy as sp

from connectome.model.util import diag_to_zero
from connectome.shuffling import shuffle_graph

N_DRAWS = 100

class TestUniformShufflingCython(unittest.TestCase):
    def test_shuffle_small_target_equal_source(self):
        a = (sp.rand(100, 100) < .3).astype(float)
        diag_to_zero(a)
        first = a.sum(0)
        second = a.sum(1)
        shuffle_graph(a, True, nr_draws=N_DRAWS)
        first_2 = a.sum(0)
        second_2 = a.sum(1)
        self.assertTrue((first == first_2).all())
        self.assertTrue((second == second_2).all())
        for k in range(a.shape[0]):
            self.assertFalse(a[k,k])

    def test_shuffle_small_target_unequal_source(self):
        a = (sp.rand(100, 200) < .3).astype(float)
        first = a.sum(0)
        second = a.sum(1)
        shuffle_graph(a, False, nr_draws=N_DRAWS)
        first_2 = a.sum(0)
        second_2 = a.sum(1)
        self.assertTrue((first == first_2).all())
        self.assertTrue((second == second_2).all())


    def test_shuffle_big_target_equal_source(self):
        N = 2000
        p = .2
        a = (sp.rand(N, N) < p).astype(float)
        diag_to_zero(a)
        a_copy = a.copy()
        first = a.sum(0)
        second = a.sum(1)
        a, nr_square_moves, nr_triangle_moves = shuffle_graph(a, True, nr_draws=N_DRAWS)
        first_2 = a.sum(0)
        second_2 = a.sum(1)
        self.assertTrue((first == first_2).all())
        self.assertTrue((second == second_2).all())
        for k in range(a.shape[0]):
            self.assertFalse(a[k,k])

    def test_shuffle_big_target_unequal_source(self):
        N = 2000
        p = .2
        a = (sp.rand(N, N//2) < p).astype(float)
        a_copy = a.copy()
        first = a.sum(0)
        second = a.sum(1)
        a, nr_square_moves, nr_triangle_moves = shuffle_graph(a, False, nr_draws=N_DRAWS)
        first_2 = a.sum(0)
        second_2 = a.sum(1)
        self.assertTrue((first == first_2).all())
        self.assertTrue((second == second_2).all())
        self.assertGreater(nr_square_moves, 10)
        off_diag = False
        for k in range(min(a.shape[0],a.shape[1])):
            off_diag |= bool(a[k,k].astype(bool))
        self.assertTrue(off_diag)

    def test_shuffle_big_target_unequal_source_float32(self):
        N = 2000
        p = .2
        a = (sp.rand(N, N//2) < p).astype(float)
        a_copy = a.copy()
        first = a.sum(0)
        second = a.sum(1)
        a, nr_square_moves, nr_triangle_moves = shuffle_graph(a.astype(sp.float32), False, nr_draws=N_DRAWS)
        first_2 = a.sum(0)
        second_2 = a.sum(1)
        self.assertTrue((first == first_2).all())
        self.assertTrue((second == second_2).all())
        off_diag = False
        for k in range(min(a.shape[0],a.shape[1])):
            off_diag |= bool(a[k,k].astype(bool))
        self.assertTrue(off_diag)


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