from __future__ import annotations

import unittest

import numpy as np

from pps_python.pps_loop import PIConfig, simulate_controller
from pps_python.pps_model import (
    C,
    RF_HZ,
    TRACKING_HZ,
    ReceiverEpochSimulator,
    build_clock_trajectory,
    lla_to_ecef,
    make_satellite_constellation,
    signed_u32_delta,
)


def make_simulation(duration_s: float = 70.0, seed: int = 11):
    rng = np.random.default_rng(seed)
    shared = build_clock_trajectory(
        duration_s + 2.0,
        1.0,
        rng,
        base_ppm=1.5,
        temperature_ppm=0.3,
        temperature_period_s=100.0,
        random_walk_ppm_sqrt_s=0.0,
    )
    rf = build_clock_trajectory(
        duration_s + 2.0,
        RF_HZ,
        rng,
        base_ppm=0.1,
        temperature_ppm=0.0,
        temperature_period_s=100.0,
        random_walk_ppm_sqrt_s=0.0,
        shared_fractional=shared.fractional_frequency,
        initial_ticks=float((1 << 32) - 100_000),
    )
    tracking = build_clock_trajectory(
        duration_s + 2.0,
        TRACKING_HZ,
        rng,
        base_ppm=0.2,
        temperature_ppm=0.1,
        temperature_period_s=130.0,
        random_walk_ppm_sqrt_s=0.0,
        shared_fractional=shared.fractional_frequency,
    )
    receiver = lla_to_ecef(39.9208, 32.8541, 890.0)
    simulator = ReceiverEpochSimulator(
        receiver,
        make_satellite_constellation(receiver),
        rf,
        tracking,
        rng,
    )
    return receiver, tracking, simulator.simulate(duration_s)


class PpsModelTests(unittest.TestCase):
    def test_u32_wrap_delta_matches_firmware_rule(self):
        self.assertEqual(signed_u32_delta(4, 0xFFFF_FFFC), 8)
        self.assertEqual(signed_u32_delta(0xFFFF_FFFC, 4), -8)

    def test_tow_and_bias_reconstruct_reference_epoch(self):
        _, _, epochs = make_simulation(70.0)
        mature = [epoch for epoch in epochs if epoch.transmit_true_s >= 50.0]
        timing_error = np.asarray(
            [epoch.timing_measurement_error_s for epoch in mature]
        )
        raw_bias = np.asarray([epoch.raw_bias_phase_s for epoch in mature])
        self.assertLess(np.sqrt(np.mean(timing_error**2)), 250.0e-9)
        self.assertGreater(np.median(np.abs(raw_bias)), 100.0e-6)
        self.assertGreater(
            np.median(np.abs(raw_bias)),
            1000.0 * np.median(np.abs(timing_error)),
        )

    def test_fractional_code_phase_is_applied_to_capture_counter(self):
        _, _, epochs = make_simulation(70.0)
        mature = [epoch for epoch in epochs if epoch.transmit_true_s >= 50.0]
        applied_s = np.asarray(
            [
                (
                    epoch.reference_tracking_tick_raw
                    - epoch.reference_tracking_tick
                )
                / TRACKING_HZ
                for epoch in mature
            ]
        )
        self.assertGreater(np.max(applied_s), 40.0e-9)
        self.assertTrue(np.all(applied_s >= 0.0))
        self.assertTrue(np.all(applied_s <= 64.0e-9))
        # The integer mailbox can represent the correction only on the 8 ns
        # PPS grid; the model must use the same nearest-tick quantization.
        applied_ticks = applied_s * TRACKING_HZ
        self.assertLess(np.max(np.abs(applied_ticks - np.rint(applied_ticks))), 1.0e-9)

    def test_timestamped_pi_rejects_apb_latency_as_phase(self):
        _, tracking, epochs = make_simulation(70.0)
        config = PIConfig(natural_frequency_hz=0.035, damping=0.9, max_slew_ppm=20.0)
        timestamped = simulate_controller(
            "timestamped",
            "timestamped",
            epochs,
            tracking,
            345_600.0,
            70.0,
            config,
        )
        untagged = simulate_controller(
            "untagged",
            "untagged",
            epochs,
            tracking,
            345_600.0,
            70.0,
            config,
        )
        timestamped_tail = timestamped.phase_error_s[len(timestamped.phase_error_s) // 2 :]
        untagged_tail = untagged.phase_error_s[len(untagged.phase_error_s) // 2 :]
        timestamped_rms = np.sqrt(np.nanmean(timestamped_tail**2))
        untagged_rms = np.sqrt(np.nanmean(untagged_tail**2))
        self.assertLess(timestamped_rms, 10.0e-6)
        self.assertGreater(untagged_rms, 1.0e-3)
        self.assertGreater(untagged_rms, 100.0 * timestamped_rms)


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