"""Frozen curve arithmetic/filter comparisons using current model-shaped data."""

from types import SimpleNamespace
import math
import unittest

from openpilot.starpilot.conditional_mode.curve import CurveDetector, frozen_raw_curve, frozen_road_curvature


def model(accel=2.0):
  # The frozen helper picks the largest absolute orientationRate*velocity.
  rates = [0.0] * 33
  rates[5] = accel / 15.0
  return SimpleNamespace(orientationRate=SimpleNamespace(z=rates, t=[i * 0.05 for i in range(33)]), velocity=SimpleNamespace(x=[15.0] * 33))


class TestFrozenCurve(unittest.TestCase):
  def test_frozen_max_lateral_accel_and_blinker(self):
    self.assertEqual(frozen_road_curvature(model(), 15.0), (2.0 / 225.0, 1.0))
    self.assertEqual(frozen_raw_curve(model(), 15.0, 0.0, False, False), (True, False))
    self.assertEqual(frozen_raw_curve(model(), 15.0, 0.008, True, False), (False, True))
    self.assertEqual(frozen_raw_curve(model(0.0), 15.0, 0.0, False, False), (False, False))
    broken = model()
    broken.orientationRate.z[1] = math.nan
    self.assertIsNone(frozen_road_curvature(broken, 15.0))

  def test_executed_frozen_curvature_reference_cases(self):
    # Generated by executing only calculate_road_curvature from frozen
    # 678af783/starpilot/common/starpilot_utilities.py against these arrays.
    for speed, acceleration, expected in ((0.0, 0.0, 0.0), (5.0, 0.2, 0.008), (15.0, 2.25, 0.01), (30.0, -1.76, -0.0019555555555555554)):
      result = frozen_road_curvature(model(acceleration), speed)
      self.assertIsNotNone(result)
      if result is not None:
        self.assertAlmostEqual(result[0], expected, places=12)
        self.assertEqual(result[1], 1.0)

  def test_frozen_filter_and_no_duplicate_source_tick(self):
    detector = CurveDetector()
    states = []
    for tick in range(30):
      instant = 100_000_000_000 + tick * 50_000_000
      states.append(detector.step(observed_mono_ns=instant, speed_mps=20.0, raw_curve=(True, False), traffic_mode=False))
      self.assertIsNone(detector.step(observed_mono_ns=instant, speed_mps=20.0, raw_curve=(True, False), traffic_mode=False))
    self.assertFalse(states[0])
    self.assertTrue(states[-1])
    self.assertIsNone(detector.step(observed_mono_ns=103_000_000_000, speed_mps=20.0, raw_curve=(True, False), traffic_mode=False))
    self.assertFalse(detector.step(observed_mono_ns=103_050_000_000, speed_mps=20.0, raw_curve=(False, False), traffic_mode=False))

  def test_exact_and_adjacent_nanosecond_gaps(self):
    for base in (100_000_000_000, 100_100_000_000):
      for gap, expected in ((99_999_999, True), (100_000_000, True), (100_000_001, True), (250_000_000, True), (250_000_001, None)):
        with self.subTest(base_ns=base, gap_ns=gap):
          detector = CurveDetector()
          for tick in range(40):
            instant = base + tick * 50_000_000
            detector.step(observed_mono_ns=instant, speed_mps=30.0, raw_curve=(True, False), traffic_mode=False)
          result = detector.step(observed_mono_ns=instant + gap, speed_mps=30.0, raw_curve=(True, False), traffic_mode=False)
          self.assertIs(result, expected)
          if expected is None:
            self.assertEqual(detector.value, 0.0)
            self.assertIsNone(detector.last_observed_mono_ns)

  def test_backward_and_invalid_observations_do_not_renew_consensus(self):
    detector = CurveDetector()
    instant = 100_000_000_000
    detector.step(observed_mono_ns=instant, speed_mps=30.0, raw_curve=(True, False), traffic_mode=False)
    value = detector.value
    for stamp in (instant, instant - 1):
      self.assertIsNone(detector.step(observed_mono_ns=stamp, speed_mps=30.0, raw_curve=(True, False), traffic_mode=False))
      self.assertEqual(detector.value, value)
      self.assertEqual(detector.last_observed_mono_ns, instant)
    for stamp in (-1, 10**21 + 1, True, 100.0, math.nan, math.inf):
      with self.subTest(stamp=stamp):
        detector.step(observed_mono_ns=instant, speed_mps=30.0, raw_curve=(True, False), traffic_mode=False)
        self.assertIsNone(detector.step(observed_mono_ns=stamp, speed_mps=30.0, raw_curve=(True, False), traffic_mode=False))  # ty: ignore[invalid-argument-type]
        self.assertEqual(detector.value, 0.0)
        self.assertIsNone(detector.last_observed_mono_ns)

  def test_unknown_traffic_mode_resets_instead_of_inventing_false(self):
    detector = CurveDetector()
    self.assertIsNone(detector.step(observed_mono_ns=100_000_000_000, speed_mps=15.0, raw_curve=(True, False), traffic_mode=None))
    self.assertEqual(detector.value, 0.0)
    self.assertFalse(detector.step(observed_mono_ns=100_050_000_000, speed_mps=15.0, raw_curve=(True, False), traffic_mode=False))


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