"""Small executable model of the RTL code-tracking observables.

This is deliberately scoped to the parts that matter to the navigation
measurement study:

* the normalized early-minus-late magnitude discriminator;
* its signed Q2.13 output;
* the fixed C1/C2 code-loop filter and Q20 NCO command;
* the 20 one-millisecond discriminator samples available per NAV bit.

It is not an RF waveform simulator.  The C/N0-to-discriminator-noise mapping
is an explicit calibration assumption until real correlator/DLL telemetry is
recorded from the board.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


DLL_OUTPUT_FRACTION_BITS = 13
DLL_OUTPUT_WORD_BITS = 15
DLL_SPACING_CHIPS = 0.5
CODE_LOOP_C1 = 21.25
CODE_LOOP_C2 = 0.05712890625
CODE_LOOP_OUTPUT_FRACTION_BITS = 20
CODE_LOOP_OUTPUT_WORD_BITS = 28
DLL_SAMPLES_PER_NAV_BIT = 20
CODE_FREQUENCY_BASIS_HZ = 1.023e6
ABS_SAMPLE_COUNTER_MODULUS = 1 << 32

# Synthetic calibration point.  It only fixes the scale of the Monte Carlo
# discriminator telemetry.  Dynamic weighting uses a measured/expected ratio,
# so the absolute choice cancels in the correctly modelled case.
DLL_STD_AT_45_DBHZ = 0.020


def _signed_saturate(value: int, word_bits: int) -> int:
    minimum = -(1 << (word_bits - 1))
    maximum = (1 << (word_bits - 1)) - 1
    return min(maximum, max(minimum, value))


def quantize_signed(value: float, word_bits: int, fraction_bits: int) -> int:
    scaled = int(np.rint(value * (1 << fraction_bits)))
    return _signed_saturate(scaled, word_bits)


def dequantize_signed(value: int, fraction_bits: int) -> float:
    return float(value) / float(1 << fraction_bits)


def triangular_correlation(code_offset_chips: float) -> float:
    """Ideal one-chip triangular GPS C/A autocorrelation main lobe."""

    return max(0.0, 1.0 - abs(code_offset_chips))


def ideal_early_late_magnitudes(
    code_error_chips: float,
    spacing_chips: float = DLL_SPACING_CHIPS,
) -> tuple[float, float]:
    half_spacing = 0.5 * spacing_chips
    early = triangular_correlation(code_error_chips + half_spacing)
    late = triangular_correlation(code_error_chips - half_spacing)
    return early, late


def normalized_early_minus_late(early_magnitude: float, late_magnitude: float) -> float:
    denominator = early_magnitude + late_magnitude
    if denominator <= 0.0:
        return 0.0
    return (early_magnitude - late_magnitude) / denominator


def rtl_discriminator(code_error_chips: float) -> tuple[int, float]:
    """Return the RTL-like Q2.13 discriminator code and dequantized value."""

    early, late = ideal_early_late_magnitudes(code_error_chips)
    value = normalized_early_minus_late(early, late)
    code = quantize_signed(
        value,
        DLL_OUTPUT_WORD_BITS,
        DLL_OUTPUT_FRACTION_BITS,
    )
    return code, dequantize_signed(code, DLL_OUTPUT_FRACTION_BITS)


def expected_discriminator_std(cn0_dbhz: float) -> float:
    """Expected normalized DLL jitter used by the synthetic calibration.

    The 20 dB amplitude-domain law mirrors the firmware's code-sigma model.
    The result is bounded because the normalized discriminator is finite.
    """

    value = DLL_STD_AT_45_DBHZ * 10.0 ** ((45.0 - cn0_dbhz) / 20.0)
    return float(np.clip(value, 0.002, 0.35))


def quantized_code_epoch_measurement(
    true_abs_sample: float,
) -> tuple[int, int, float]:
    """Model ABS_SAMPLE + REM_CODE_PHASE and firmware reconstruction.

    ``true_abs_sample`` is the ideal code boundary in RF-sample units.  The
    hardware captures the following integer sample and reports the overshoot
    with the remainder scale used by the C correction formula.
    """

    raw_unwrapped = int(np.ceil(true_abs_sample))
    overshoot_samples = raw_unwrapped - true_abs_sample
    remainder = int(np.rint(overshoot_samples * CODE_FREQUENCY_BASIS_HZ))
    raw_register = raw_unwrapped % ABS_SAMPLE_COUNTER_MODULUS
    corrected = raw_unwrapped - remainder / CODE_FREQUENCY_BASIS_HZ
    return raw_register, remainder, corrected


@dataclass(frozen=True)
class DllWindow:
    samples: np.ndarray
    mean: float
    variance: float
    expected_variance: float
    variance_ratio: float
    mean_z_score: float


def synthesize_dll_window(
    rng: np.random.Generator,
    *,
    cn0_dbhz: float,
    jitter_multiplier: float = 1.0,
    discriminator_bias: float = 0.0,
    sample_count: int = DLL_SAMPLES_PER_NAV_BIT,
) -> DllWindow:
    """Generate and Q2.13-quantize one NAV-bit DLL telemetry window."""

    if sample_count < 2:
        raise ValueError("DLL variance needs at least two samples")
    expected_std = expected_discriminator_std(cn0_dbhz)
    actual_std = expected_std * max(0.0, jitter_multiplier)
    raw = discriminator_bias + rng.normal(0.0, actual_std, sample_count)
    minimum = -(1 << (DLL_OUTPUT_WORD_BITS - 1))
    maximum = (1 << (DLL_OUTPUT_WORD_BITS - 1)) - 1
    quantized_codes = np.clip(
        np.rint(raw * (1 << DLL_OUTPUT_FRACTION_BITS)),
        minimum,
        maximum,
    ).astype(np.int32)
    samples = quantized_codes.astype(float) / float(1 << DLL_OUTPUT_FRACTION_BITS)
    mean = float(np.mean(samples))
    variance = float(np.var(samples, ddof=1))
    expected_variance = expected_std * expected_std
    ratio = variance / max(expected_variance, 1.0e-15)
    standard_error = np.sqrt(max(variance, expected_variance) / sample_count)
    mean_z = abs(mean) / max(float(standard_error), 1.0e-12)
    return DllWindow(
        samples=samples,
        mean=mean,
        variance=variance,
        expected_variance=expected_variance,
        variance_ratio=float(ratio),
        mean_z_score=float(mean_z),
    )


@dataclass
class CodeLoopFilterRtl:
    """Update-rate model of LF_code.sv, including coefficient quantization."""

    integrator_q20: int = 0

    c1_integer: int = int(CODE_LOOP_C1 * (1 << 2))
    c2_integer: int = int(CODE_LOOP_C2 * (1 << 11))

    def update_code(self, discriminator_q13: int) -> int:
        # C1 product has Q15 and is zero-padded to Q20 in LF_code.sv.
        proportional_q20 = (self.c1_integer * discriminator_q13) << 5

        # C2 product has Q24 and LF_code.sv slices it back to Q20.
        integral_increment_q20 = (self.c2_integer * discriminator_q13) >> 4
        self.integrator_q20 = _signed_saturate(
            self.integrator_q20 + integral_increment_q20,
            CODE_LOOP_OUTPUT_WORD_BITS,
        )
        return _signed_saturate(
            self.integrator_q20 + proportional_q20,
            CODE_LOOP_OUTPUT_WORD_BITS,
        )

    def update(self, discriminator: float) -> float:
        discriminator_q13 = quantize_signed(
            discriminator,
            DLL_OUTPUT_WORD_BITS,
            DLL_OUTPUT_FRACTION_BITS,
        )
        command_q20 = self.update_code(discriminator_q13)
        return dequantize_signed(command_q20, CODE_LOOP_OUTPUT_FRACTION_BITS)