import unittest

import cv2
import numpy as np

from helpers import acquisition_step_is_plausible
from target_physics import analyze_motion_group
from helpers import box_is_edge_osd, box_is_osd_candidate, in_osd_zone


def frame_with_points(points, shape=(240, 320)):
    frame = np.zeros(shape, dtype=np.uint8)
    for x, y in points:
        cv2.circle(frame, (int(round(x)), int(round(y))), 2, 255, -1)
    return frame


class TargetPhysicsTests(unittest.TestCase):
    def test_provisional_target_cannot_jump_across_the_frame(self):
        previous = np.array([350, 290, 380, 310], dtype=np.float32)
        jumped = np.array([395, 220, 430, 240], dtype=np.float32)
        self.assertFalse(
            acquisition_step_is_plausible(
                previous,
                jumped,
                np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32),
                0.08,
                720,
                576,
            )
        )

    def test_provisional_target_allows_camera_compensated_motion(self):
        previous = np.array([350, 290, 380, 310], dtype=np.float32)
        current = np.array([370, 300, 400, 320], dtype=np.float32)
        affine = np.array([[1, 0, 20], [0, 1, 10]], dtype=np.float32)
        self.assertTrue(
            acquisition_step_is_plausible(
                previous,
                current,
                affine,
                0.08,
                720,
                576,
            )
        )

    def test_edge_anchored_osd_box_is_rejected(self):
        self.assertTrue(box_is_edge_osd([0, 4, 80, 60], 720, 576))
        self.assertFalse(box_is_edge_osd([300, 200, 380, 260], 720, 576))
        self.assertTrue(in_osd_zone(150, 468, 720, 576))
        self.assertTrue(box_is_osd_candidate([294, 538, 332, 560], 720, 576))

    def test_camera_motion_is_not_independent_target_motion(self):
        points = np.array(
            [(125, 85), (135, 85), (145, 85), (125, 95), (135, 95),
             (145, 95), (125, 105), (135, 105), (145, 105)],
            dtype=np.float32,
        )
        affine = np.array([[1, 0, 3], [0, 1, 2]], dtype=np.float32)
        previous = frame_with_points(points)
        current = cv2.warpAffine(previous, affine, (320, 240))

        evidence = analyze_motion_group(
            previous, current, [115, 75, 155, 115], affine=affine, dt=0.04
        )

        self.assertTrue(evidence.reliable)
        self.assertFalse(evidence.valid)
        self.assertLess(evidence.residual_px, 0.25)

    def test_coherent_expanding_group_is_valid_target_motion(self):
        points = np.array(
            [(125, 85), (135, 85), (145, 85), (125, 95), (135, 95),
             (145, 95), (125, 105), (135, 105), (145, 105)],
            dtype=np.float32,
        )
        center = np.array([135, 95], dtype=np.float32)
        moved = (points - center) * 1.08 + center + np.array([4, 1], dtype=np.float32)

        evidence = analyze_motion_group(
            frame_with_points(points),
            frame_with_points(moved),
            [115, 75, 155, 115],
            dt=0.04,
        )

        self.assertTrue(evidence.reliable)
        self.assertTrue(evidence.valid)
        self.assertGreaterEqual(evidence.coherent_count, 3)
        self.assertGreater(evidence.scale_ratio, 1.01)

    def test_screen_fixed_osd_is_not_a_target_after_camera_compensation(self):
        points = np.array(
            [(25, 15), (35, 15), (45, 15), (25, 25), (35, 25),
             (45, 25), (25, 35), (35, 35), (45, 35)],
            dtype=np.float32,
        )
        frame = frame_with_points(points)
        affine = np.array([[1, 0, 4], [0, 1, 2]], dtype=np.float32)

        evidence = analyze_motion_group(
            frame, frame, [15, 5, 55, 45], affine=affine, dt=0.04
        )

        self.assertTrue(evidence.reliable)
        self.assertTrue(evidence.screen_static)
        self.assertFalse(evidence.valid)

    def test_distant_target_moving_out_of_frame_is_rejected(self):
        points = np.array(
            [(3, 90), (8, 90), (13, 90), (3, 98), (8, 98),
             (13, 98), (3, 106), (8, 106), (13, 106)],
            dtype=np.float32,
        )

        evidence = analyze_motion_group(
            frame_with_points(points),
            frame_with_points(points + np.array([-2, 0], dtype=np.float32)),
            [0, 82, 18, 112],
            dt=0.04,
        )

        self.assertTrue(evidence.reliable)
        self.assertTrue(evidence.edge_violation)
        self.assertFalse(evidence.valid)


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