import unittest

import numpy as np
import tensorflow as tf

import warp


class MyTestCase(unittest.TestCase):
    def test_euler_to_rotation_matrix(self):
        # quarter rotation in every
        x, y, z = tf.expand_dims(tf.expand_dims(tf.constant(np.pi / 2), 0), 0)

        # TODO: Construct expected final rotation matrix, just 3x3 using numpy, so that we can do an
        #   elementwise comparison later. Probably also want to check the

        rotation_matrices = warp.euler_to_rotation_matrix(x, y, z)

        self.assertEqual(rotation_matrices.shape, [1, 3, 3])
        rot_mat = rotation_matrices[0]

        # TODO: Element-wise checks...

    def test_coordinates(self):
        height = 1000
        width = 2000
        coords = warp.image_coordinate(8, height, width)

        self.assertEqual(coords.shape, [8, height, width, 3])
        self.assertEqual(coords[0, 0, 0, 0], 0)
        self.assertEqual(coords[0, 0, 0, 1], 0)
        self.assertEqual(coords[0, 0, 0, 2], 1)

        self.assertEqual(coords[0, height - 1, 0, 0], 0)
        self.assertEqual(coords[0, height - 1, 0, 1], height - 1)
        self.assertEqual(coords[0, height - 1, 0, 2], 1)

        self.assertEqual(coords[0, height - 1, width - 1, 0], width - 1)
        self.assertEqual(coords[0, height - 1, width - 1, 1], height - 1)
        self.assertEqual(coords[0, height - 1, width - 1, 2], 1)


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