import unittest

import numpy as np

from brainbox import video


class TestVideo(unittest.TestCase):
    def setUp(self) -> None:
        """Test frames are 12 2x2x3 arrays, where
            frames[0][...,0] =
                [[1, 1],
                 [1, 1]]

            frames[1][...,0] =
                [[2, 2],
                 [2, 2]]

            [...]

            frames[-1][...,0] =
                [[12, 12],
                 [12, 12]]
            """
        self.frames = np.cumsum(np.ones((12, 2, 2, 3)), axis=0).astype(np.uint8)

    def test_frame_diff(self):
        # Test with three colour channels (2x2x3)
        frame1, frame2 = [self.frames[i] for i in range(2)]
        df = video.frame_diff(frame1, frame2)
        expected = np.ones(frame1.shape[:2], dtype=np.uint8)
        np.testing.assert_equal(df, expected)

        # Test with single channel (2x2)
        df = video.frame_diff(frame1[..., 0], frame2[..., 0])
        np.testing.assert_equal(df, expected)

        # Test shape validation
        with self.assertRaises(ValueError):
            video.frame_diff(frame1[..., 0], frame2)

    def test_frame_diffs(self):
        # Test frame diffs defaults
        df = video.frame_diffs(self.frames)
        expected = np.ones((self.frames.shape[0] - 1, *self.frames.shape[1:-1]), dtype=np.uint8)
        np.testing.assert_equal(df, expected)

        # Test shape validation
        with self.assertRaises(ValueError):
            video.frame_diffs(self.frames, diff=20)

        # Test frames diff every 2nd frame with intensity frames
        d = 2  # Take difference every two frames
        df = video.frame_diffs(self.frames[..., 0], d)
        expected_shape = (self.frames.shape[0] - d, *self.frames.shape[1:-1])
        expected = np.ones(expected_shape, dtype=np.uint8) * 2
        np.testing.assert_equal(df, expected)


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