import dv_processing as dv
import unittest
import numpy as np


class StereoMatcherTest(unittest.TestCase):

    @staticmethod
    def read_test_data():
        store = dv.EventStore()
        reader = dv.io.MonoCameraRecording("./data/sample_data.aedat4")
        packet = reader.getNextEventBatch()
        while packet is not None:
            store.add(packet)
            packet = reader.getNextEventBatch()
        return store

    def test_disparity(self):
        resolution = (640, 480)
        matcher = dv.SemiDenseStereoMatcher(resolution, resolution)

        store = StereoMatcherTest.read_test_data()

        left_events = dv.EventStore()
        left_events.add(store)
        left_events.add(store)
        left_events.add(store)
        left_events.add(store)

        true_disparity_value = 10
        right_events = dv.EventStore()
        for event in store:
            right_events.push_back(event.timestamp(), event.x() - true_disparity_value, event.y(), event.polarity())

        right_events.add(right_events)
        right_events.add(right_events)
        right_events.add(right_events)

        disparity = matcher.computeDisparity(left_events, right_events)
        for event in store:
            disparity_value = disparity[event.y(), event.x()] / 16.0
            self.assertEqual(round(disparity_value), true_disparity_value)

    def test_depth_estimation(self):
        resolution = (640, 480)
        store = StereoMatcherTest.read_test_data()

        left_events = dv.EventStore()
        left_events.add(store)
        left_events.add(store)
        left_events.add(store)
        left_events.add(store)

        true_disparity_value = 10
        right_events = dv.EventStore()
        for event in store:
            right_events.push_back(event.timestamp(), event.x() - true_disparity_value, event.y(), event.polarity())

        right_events.add(right_events)
        right_events.add(right_events)
        right_events.add(right_events)

        ideal_geometry = dv.camera.CameraGeometry(resolution[0] * 0.5, resolution[0] * 0.5, resolution[0] * 0.5,
                                                  resolution[1] * 0.5, resolution)

        #Cameras 1.4m apart
        T = np.identity(4)
        T[0, -1] = 1.4
        stereo_geometry = dv.camera.StereoGeometry(ideal_geometry, ideal_geometry, dv.kinematics.Transformationf(0, T))
        matcher = dv.SemiDenseStereoMatcher(stereo_geometry)
        disparity = matcher.computeDisparity(left_events, right_events)

        depth_frame = matcher.estimateDepthFrame(disparity)
        depth_events = matcher.estimateDepth(disparity, store)

        # Check estimated depth is within tolerance (+- 0.1 px disparity)
        minDepth = stereo_geometry.convertDisparityToDepth(true_disparity_value + 0.1)
        maxDepth = stereo_geometry.convertDisparityToDepth(true_disparity_value - 0.1)

        # Not to run too long, test on small sample
        for event in depth_events.slice(0, 500):
            self.assertGreaterEqual(event.depth(), int(minDepth))
            self.assertLessEqual(event.depth(), int(maxDepth))

            # Check if estimated depth frame yields the same result as for the events
            rectifiedEventPos = stereo_geometry.remapPoint(dv.camera.StereoGeometry.CameraPosition.LEFT,
                                                           (event.x(), event.y()))
            if rectifiedEventPos is None:
                continue

            rectX, rectY = rectifiedEventPos
            index = rectY * resolution[0] + rectX
            depth = depth_frame.depth[index]

            self.assertEqual(depth, event.depth())
            self.assertGreaterEqual(depth, minDepth)
            self.assertLessEqual(depth, maxDepth)


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