import unittest

import scipy as sp

from connectome.function.task.statistics import pair_correlation


class TestAutocorr(unittest.TestCase):
    def test_autocorr_columns(self):
        """
        Check I got the axis right.
        """
        arr = sp.zeros((100, 10))
        arr[:, 0] = 1
        arr[:, 2] = 1
        arr[:,-1] = 1
        cv = pair_correlation(arr.T)
        self.assertLess(abs(1-cv), 1e-5)

    def test_autocorr_rows(self):
        """
        Check I got the axis right.
        """
        arr = sp.zeros((100, 10))
        val = sp.rand(100) > .5
        for k in range(10):
            arr[:,k] = val
        cv = pair_correlation(arr)
        self.assertLess(abs(1 - cv), 1e-5)


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