from __future__ import annotations

import unittest

import numpy as np

from gps_measurement_python.code_smoothing_model import run_smoothing_trial
from gps_measurement_python.firmware_wls_model import (
    SigmaInputs,
    SigmaMethod,
    firmware_sigma_m,
)
from gps_measurement_python.measurement_scenarios import (
    run_scenario,
    scenario_matrix,
)
from gps_measurement_python.rtl_dll_model import (
    CodeLoopFilterRtl,
    quantized_code_epoch_measurement,
    rtl_discriminator,
    synthesize_dll_window,
)


class RtlDllModelTests(unittest.TestCase):
    def test_discriminator_zero_and_q13_symmetry(self):
        self.assertEqual(rtl_discriminator(0.0)[0], 0)
        positive = rtl_discriminator(0.05)[1]
        negative = rtl_discriminator(-0.05)[1]
        self.assertAlmostEqual(positive, -negative, places=6)
        self.assertLess(positive, 0.0)

    def test_code_loop_integrator_accumulates_constant_error(self):
        loop = CodeLoopFilterRtl()
        first = loop.update(0.01)
        twentieth = first
        for _ in range(19):
            twentieth = loop.update(0.01)
        self.assertGreater(first, 0.0)
        self.assertGreater(twentieth, first)

    def test_dll_window_detects_excess_variance(self):
        rng = np.random.default_rng(42)
        nominal = synthesize_dll_window(
            rng,
            cn0_dbhz=45.0,
            jitter_multiplier=1.0,
            sample_count=4000,
        )
        noisy = synthesize_dll_window(
            rng,
            cn0_dbhz=45.0,
            jitter_multiplier=3.0,
            sample_count=4000,
        )
        self.assertAlmostEqual(nominal.variance_ratio, 1.0, delta=0.10)
        self.assertAlmostEqual(noisy.variance_ratio, 9.0, delta=0.70)

    def test_abs_sample_remainder_reconstructs_subsample_epoch(self):
        for true_sample in (
            12_345_678.314159,
            12_345_678.0000001,
            12_345_678.9999999,
            12_345_678.0,
        ):
            raw, remainder, corrected = quantized_code_epoch_measurement(true_sample)
            self.assertEqual(raw, int(np.ceil(true_sample)))
            self.assertLessEqual(
                abs(corrected - true_sample),
                0.5 / 1.023e6 + 1.0e-12,
            )
            self.assertGreaterEqual(remainder, 0)


class FirmwareSigmaTests(unittest.TestCase):
    @staticmethod
    def inputs(ratio: float = 1.0, mean_z: float = 0.0) -> SigmaInputs:
        return SigmaInputs(
            ura_index=np.asarray([0]),
            elevation_deg=np.asarray([90.0]),
            cn0_dbhz=np.asarray([45.0]),
            dll_variance_ratio=np.asarray([ratio]),
            dll_mean_z=np.asarray([mean_z]),
            actual_jitter_multiplier=np.asarray([np.sqrt(ratio)]),
        )

    def test_current_matches_c_formula(self):
        sigma = firmware_sigma_m(self.inputs(), SigmaMethod.CURRENT)[0]
        expected = np.sqrt(2.4**2 + 3.0**2 + 3.0**2)
        self.assertAlmostEqual(sigma, expected, places=12)

    def test_excess_method_does_not_double_count_nominal_dll_noise(self):
        current = firmware_sigma_m(self.inputs(), SigmaMethod.CURRENT)[0]
        excess = firmware_sigma_m(self.inputs(), SigmaMethod.DLL_EXCESS)[0]
        naive = firmware_sigma_m(self.inputs(), SigmaMethod.NAIVE_ADD_DLL)[0]
        self.assertAlmostEqual(excess, current, places=12)
        self.assertGreater(naive, current)

    def test_excess_method_inflates_hidden_jitter(self):
        current = firmware_sigma_m(self.inputs(9.0), SigmaMethod.CURRENT)[0]
        excess = firmware_sigma_m(self.inputs(9.0), SigmaMethod.DLL_EXCESS)[0]
        self.assertGreater(excess, 1.8 * current)


class ScenarioTests(unittest.TestCase):
    def test_hidden_jitter_dynamic_sigma_improves_seeded_case(self):
        scenario = next(
            item
            for item in scenario_matrix()
            if item.name == "one_channel_hidden_jitter"
        )
        samples = run_scenario(
            scenario,
            17,
            (SigmaMethod.CURRENT, SigmaMethod.DLL_EXCESS),
        )
        current = np.asarray(
            [
                item.position_error_3d_m
                for item in samples
                if item.method == SigmaMethod.CURRENT and item.time_s >= 5.0
            ]
        )
        dynamic = np.asarray(
            [
                item.position_error_3d_m
                for item in samples
                if item.method == SigmaMethod.DLL_EXCESS and item.time_s >= 5.0
            ]
        )
        self.assertLess(
            np.sqrt(np.mean(dynamic * dynamic)),
            0.90 * np.sqrt(np.mean(current * current)),
        )

    def test_stable_multipath_is_not_solved_by_variance_only(self):
        scenario = next(
            item
            for item in scenario_matrix()
            if item.name == "stable_multipath_bias"
        )
        samples = run_scenario(
            scenario,
            23,
            (SigmaMethod.CURRENT, SigmaMethod.DLL_EXCESS),
        )
        current = np.asarray(
            [
                item.position_error_3d_m
                for item in samples
                if item.method == SigmaMethod.CURRENT and item.time_s >= 5.0
            ]
        )
        dynamic = np.asarray(
            [
                item.position_error_3d_m
                for item in samples
                if item.method == SigmaMethod.DLL_EXCESS and item.time_s >= 5.0
            ]
        )
        relative_difference = abs(
            np.sqrt(np.mean(dynamic * dynamic))
            / np.sqrt(np.mean(current * current))
            - 1.0
        )
        self.assertLess(relative_difference, 0.08)

    def test_doppler_prediction_avoids_blind_iir_dynamic_lag(self):
        result = run_smoothing_trial(31, doppler_cfo_mps=0.0)
        self.assertLess(result.doppler_smoothed_rmse_m, result.raw_rmse_m)
        self.assertLess(
            result.doppler_peak_dynamic_error_m,
            0.10 * result.scalar_iir_peak_dynamic_error_m,
        )
        self.assertGreater(
            result.firmware_iir_doppler_peak_dynamic_error_m,
            result.doppler_peak_dynamic_error_m,
        )


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