from __future__ import annotations

import numpy as np
import unittest

from rfg_adc_plotter.state.background_buffer import BackgroundMedianBuffer


class BackgroundMedianBufferTests(unittest.TestCase):
    def test_buffer_returns_median_for_partial_fill(self):
        buffer = BackgroundMedianBuffer(max_rows=4)
        buffer.push(np.asarray([1.0, 5.0, 9.0], dtype=np.float32))
        buffer.push(np.asarray([3.0, 7.0, 11.0], dtype=np.float32))

        median = buffer.median()

        self.assertIsNotNone(median)
        self.assertTrue(np.allclose(median, np.asarray([2.0, 6.0, 10.0], dtype=np.float32)))

    def test_buffer_wraparound_keeps_latest_rows(self):
        buffer = BackgroundMedianBuffer(max_rows=2)
        buffer.push(np.asarray([1.0, 5.0], dtype=np.float32))
        buffer.push(np.asarray([3.0, 7.0], dtype=np.float32))
        buffer.push(np.asarray([9.0, 11.0], dtype=np.float32))

        median = buffer.median()

        self.assertIsNotNone(median)
        self.assertTrue(np.allclose(median, np.asarray([6.0, 9.0], dtype=np.float32)))

    def test_buffer_reset_clears_state(self):
        buffer = BackgroundMedianBuffer(max_rows=2)
        buffer.push(np.asarray([1.0, 2.0], dtype=np.float32))

        buffer.reset()

        self.assertIsNone(buffer.rows)
        self.assertIsNone(buffer.median())
        self.assertEqual(buffer.count, 0)
        self.assertEqual(buffer.head, 0)


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