"""Tests for the Kamil ADC reference-channel signal processing.

References are built from an explicit *unwrapped* phase ramp whose first sample
lies in ``(-pi, pi]`` and whose steps are below ``pi``, so that
``np.unwrap(np.angle(ref))`` recovers exactly the phase we specify. This mirrors
how a real swept reference behaves (absolute phase anchored at sample 0,
accumulating forward) and lets us reason precisely about the frequency mapping.
"""

from __future__ import annotations

import unittest

import numpy as np

from python_app.hardware_full.kamil_adc.processing import (
    KamilAdcProcessingParams,
    KamilAdcSweepProcessor,
)


def _reference(phase: np.ndarray, amplitude: float = 1000.0) -> np.ndarray:
    """A reference signal whose unwrapped phase equals ``phase`` (radians)."""
    return amplitude * np.exp(1j * np.asarray(phase, dtype=np.float64))


# Calibration shared by the processor tests: phase 0 rad -> 2 GHz, phase 100 rad
# -> 4 GHz (slope 2e7 Hz/rad). Band [2.5, 3.5] GHz corresponds to phase [25, 75].
_CALIBRATION = dict(
    phase0_rad=0.0,
    freq0_hz=2_000_000_000.0,
    phase1_rad=100.0,
    freq1_hz=4_000_000_000.0,
    band_start_hz=2_500_000_000.0,
    band_stop_hz=3_500_000_000.0,
    band_points=11,
)


class ProcessingParamsTest(unittest.TestCase):
    def _params(self, **overrides: float) -> KamilAdcProcessingParams:
        return KamilAdcProcessingParams(**{**_CALIBRATION, **overrides})  # type: ignore[arg-type]

    def test_slope_is_constant_from_anchors(self) -> None:
        self.assertAlmostEqual(self._params().hz_per_rad, 2.0e7)

    def test_rejects_equal_phase_anchors(self) -> None:
        with self.assertRaisesRegex(ValueError, "distinct phases"):
            self._params(phase0_rad=5.0, phase1_rad=5.0)

    def test_rejects_equal_frequency_anchors(self) -> None:
        with self.assertRaisesRegex(ValueError, "distinct frequencies"):
            self._params(freq0_hz=3.0e9, freq1_hz=3.0e9)

    def test_rejects_inverted_band(self) -> None:
        with self.assertRaisesRegex(ValueError, "stop_hz must be greater"):
            self._params(band_start_hz=4.0e9, band_stop_hz=3.0e9)

    def test_rejects_degenerate_point_count(self) -> None:
        with self.assertRaisesRegex(ValueError, "points must be >= 2"):
            self._params(band_points=1)


class SweepProcessorTest(unittest.TestCase):
    def _processor(self, **overrides: float) -> KamilAdcSweepProcessor:
        return KamilAdcSweepProcessor(KamilAdcProcessingParams(**{**_CALIBRATION, **overrides}))  # type: ignore[arg-type]

    def test_grid_is_fixed_and_identical_across_calls(self) -> None:
        processor = self._processor()
        grid = processor.grid_hz
        self.assertEqual(grid.shape, (11,))
        self.assertAlmostEqual(float(grid[0]), 2.5e9)
        self.assertAlmostEqual(float(grid[-1]), 3.5e9)
        np.testing.assert_array_equal(grid, processor.grid_hz)

    def test_frequency_axis_follows_calibration_law(self) -> None:
        processor = self._processor()
        # A dense ramp 0 -> 50 rad: endpoints map to 2 GHz and 3 GHz.
        ref = _reference(np.linspace(0.0, 50.0, 201))
        freqs = processor.reference_frequency_axis(ref)
        self.assertAlmostEqual(float(freqs[0]), 2.0e9, delta=1.0)
        self.assertAlmostEqual(float(freqs[-1]), 3.0e9, delta=1.0)

    def test_uses_config_constants_not_sweep_endpoints(self) -> None:
        """Two sweeps with different phase spans map a given absolute phase to the
        SAME frequency — proving fixed config anchors, not endpoint normalization."""
        processor = self._processor()
        ref_short = _reference(np.linspace(0.0, 100.0, 401))   # spans phase [0, 100]
        ref_long = _reference(np.linspace(0.0, 130.0, 521))    # spans phase [0, 130]
        # Encode S(f) = (f - 3 GHz) / 1 GHz, a line in TRUE frequency.
        line = lambda ref: ((processor.reference_frequency_axis(ref) - 3.0e9) / 1.0e9) * np.abs(ref)
        s_short = processor.process(line(ref_short), ref_short)
        s_long = processor.process(line(ref_long), ref_long)
        self.assertIsNotNone(s_short)
        self.assertIsNotNone(s_long)
        expected = (processor.grid_hz.astype(np.float64) - 3.0e9) / 1.0e9
        # Endpoint normalization would compress the longer sweep's axis and break this.
        np.testing.assert_allclose(s_short.real, expected, atol=1e-3)
        np.testing.assert_allclose(s_long.real, expected, atol=1e-3)
        np.testing.assert_allclose(s_short.imag, 0.0, atol=1e-3)

    def test_amplitude_normalization_divides_by_reference_magnitude(self) -> None:
        processor = self._processor()
        ref = _reference(np.linspace(0.0, 100.0, 401), amplitude=4.0)
        # |main| = 8 everywhere -> |S| = 8 / 4 = 2.
        main = 8.0 * np.exp(1j * np.angle(ref))
        result = processor.process(main, ref)
        self.assertIsNotNone(result)
        np.testing.assert_allclose(np.abs(result), 2.0, atol=1e-3)

    def test_rejects_sweep_that_does_not_cover_band(self) -> None:
        processor = self._processor()
        # Phase [0, 40] -> freq [2.0, 2.8] GHz, short of the 3.5 GHz band stop.
        ref = _reference(np.linspace(0.0, 40.0, 201))
        self.assertIsNone(processor.process(np.ones(201, dtype=np.complex128), ref))

    def test_accepts_sweep_that_covers_band(self) -> None:
        processor = self._processor()
        ref = _reference(np.linspace(0.0, 100.0, 401))  # freq [2.0, 4.0] GHz
        result = processor.process(np.abs(ref).astype(np.complex128), ref)
        self.assertIsNotNone(result)
        self.assertEqual(result.shape, (11,))
        self.assertEqual(result.dtype, np.complex64)
        np.testing.assert_allclose(np.abs(result), 1.0, atol=1e-3)  # main=|ref| -> |S|=1

    def test_interpolates_linear_trace_onto_grid(self) -> None:
        processor = self._processor()
        ref = _reference(np.linspace(0.0, 100.0, 401))
        freqs = processor.reference_frequency_axis(ref)
        # S(f) = (f - 2.5 GHz) / 1 GHz  ->  must resample to that same line on the grid.
        main = ((freqs - 2.5e9) / 1.0e9) * np.abs(ref)
        result = processor.process(main, ref)
        self.assertIsNotNone(result)
        expected = (processor.grid_hz.astype(np.float64) - 2.5e9) / 1.0e9
        np.testing.assert_allclose(result.real, expected, atol=1e-3)
        np.testing.assert_allclose(result.imag, 0.0, atol=1e-3)

    def test_recovers_sweep_offset_by_a_full_branch(self) -> None:
        """A stray +/-2*pi branch jump on the unwrap anchor must be snapped back.

        The same physical sweep, offset by one full turn (as happens when trigger
        jitter carries sample 0 across the +/-pi cut), must map to the SAME frequency
        axis as the unshifted sweep instead of sliding ~one branch off the band.
        """
        processor = self._processor()
        phase = np.linspace(0.0, 100.0, 401)  # freq [2.0, 4.0] GHz, spans the band
        baseline = processor.reference_frequency_axis(_reference(phase))
        for turns in (+1, -1, +2):
            shifted = processor.reference_frequency_axis(_reference(phase + turns * 2.0 * np.pi))
            np.testing.assert_allclose(shifted, baseline, atol=1e-3)

    def test_branch_recovery_keeps_a_glitched_sweep_usable(self) -> None:
        """A branch-jumped sweep is pulled back onto the band, not rejected."""
        processor = self._processor()
        phase = np.linspace(0.0, 100.0, 401) + 2.0 * np.pi  # one full turn off
        ref = _reference(phase)
        result = processor.process(np.abs(ref).astype(np.complex128), ref)
        self.assertIsNotNone(result)
        np.testing.assert_allclose(np.abs(result), 1.0, atol=1e-3)

    def test_leaves_genuine_sub_branch_float_untouched(self) -> None:
        """A real <pi sweep-to-sweep float must NOT be mistaken for a branch jump."""
        processor = self._processor()
        phase = np.linspace(0.0, 100.0, 401)
        baseline = processor.reference_frequency_axis(_reference(phase))
        for float_rad in (0.5, -0.5, 2.0, -2.0):
            floated = processor.reference_frequency_axis(_reference(phase + float_rad))
            # The float shifts the axis by float_rad * hz_per_rad and is preserved,
            # i.e. it is not snapped away as if it were a 2*pi branch error.
            expected = baseline + float_rad * processor.params.hz_per_rad
            np.testing.assert_allclose(floated, expected, atol=1e-3)

    def test_handles_descending_phase_direction(self) -> None:
        # phase 0 -> 2 GHz, phase -100 -> 4 GHz (negative slope). Phase ramp
        # 0 -> -100 therefore sweeps frequency UP across the band.
        processor = self._processor(phase1_rad=-100.0)
        ref = _reference(np.linspace(0.0, -100.0, 401))
        result = processor.process(np.abs(ref).astype(np.complex128), ref)
        self.assertIsNotNone(result)
        self.assertEqual(result.shape, (11,))
        np.testing.assert_allclose(np.abs(result), 1.0, atol=1e-3)

    def test_rejects_too_few_points(self) -> None:
        processor = self._processor()
        one = np.ones(1, dtype=np.complex128)
        self.assertIsNone(processor.process(one, one))

    def test_rejects_length_mismatch(self) -> None:
        processor = self._processor()
        self.assertIsNone(
            processor.process(np.ones(10, dtype=np.complex128), np.ones(9, dtype=np.complex128))
        )

    def test_drops_zero_amplitude_reference_points(self) -> None:
        processor = self._processor()
        ref = _reference(np.linspace(0.0, 100.0, 401))
        main = np.abs(ref).astype(np.complex128)
        ref[10] = 0.0  # vanished reference samples must be ignored, not crash
        ref[200] = 0.0
        result = processor.process(main, ref)
        self.assertIsNotNone(result)
        self.assertTrue(np.all(np.isfinite(result.real)))
        self.assertTrue(np.all(np.isfinite(result.imag)))


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