import unittest

import numpy as np

from ballistic_trajectory import predict_ballistic
from helpers import box_wh


def trajectory_observations(with_outlier=False):
    observations = []
    for index, timestamp in enumerate(np.linspace(-0.4, 0.0, 9)):
        center = np.array([
            100.0 + 50.0 * timestamp + 10.0 * timestamp * timestamp,
            80.0 - 10.0 * timestamp + 3.0 * timestamp * timestamp,
        ])
        if with_outlier and index == 3:
            center += np.array([90.0, -70.0])
        size = np.array([
            20.0 * np.exp(0.4 * timestamp),
            10.0 * np.exp(0.4 * timestamp),
        ])
        observations.append({
            "ts": float(timestamp),
            "center": center,
            "box": [
                center[0] - 0.5 * size[0],
                center[1] - 0.5 * size[1],
                center[0] + 0.5 * size[0],
                center[1] + 0.5 * size[1],
            ],
        })
    return observations


class BallisticTrajectoryTests(unittest.TestCase):
    def test_robust_fit_ignores_single_bad_observation(self):
        prediction = predict_ballistic(
            trajectory_observations(with_outlier=True),
            0.2,
            320,
            240,
        )

        np.testing.assert_allclose(prediction["center"], [110.4, 78.12], atol=0.2)
        np.testing.assert_allclose(prediction["velocity"], [50.0, -10.0], atol=0.5)
        np.testing.assert_allclose(prediction["acceleration"], [20.0, 6.0], atol=1.0)

    def test_approach_growth_predicts_larger_box(self):
        observations = trajectory_observations()
        prediction = predict_ballistic(observations, 0.2, 320, 240)
        last_size = box_wh(observations[-1]["box"])

        self.assertGreater(box_wh(prediction["box"])[0], last_size[0])
        self.assertGreater(box_wh(prediction["box"])[1], last_size[1])

    def test_prediction_horizon_is_limited(self):
        prediction = predict_ballistic(
            trajectory_observations(),
            2.0,
            320,
            240,
            max_horizon_sec=0.55,
        )

        self.assertAlmostEqual(prediction["horizon"], 0.55)

    def test_too_short_history_returns_no_prediction(self):
        prediction = predict_ballistic(
            trajectory_observations()[:3],
            0.2,
            320,
            240,
        )

        self.assertIsNone(prediction)


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