"""Synthetic code-tracking scenarios shared by the analysis and tests."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np

from .firmware_wls_model import (
    SigmaInputs,
    SigmaMethod,
    URA_SIGMA_M,
    cn0_code_sigma_m,
    elevation_sigma_m,
    enu_axes,
    firmware_sigma_m,
    firmware_wls,
    lla_to_ecef,
    position_error_enu_m,
    satellites_from_azel,
)
from .rtl_dll_model import expected_discriminator_std, synthesize_dll_window


TRACKING_PERIOD_S = 0.020
WLS_PERIOD_S = 0.200
TRACKING_UPDATES_PER_WLS = int(round(WLS_PERIOD_S / TRACKING_PERIOD_S))
FIRMWARE_CN0_ALPHA = 0.20
DLL_RATIO_ALPHA = 0.10
DLL_MEAN_ALPHA = 0.20

GOOD_AZIMUTH_DEG = np.asarray([8.0, 55.0, 104.0, 151.0, 205.0, 253.0, 301.0, 342.0])
GOOD_ELEVATION_DEG = np.asarray([73.0, 48.0, 31.0, 18.0, 64.0, 39.0, 24.0, 14.0])
POOR_AZIMUTH_DEG = np.asarray([25.0, 46.0, 68.0, 89.0, 112.0, 136.0, 159.0, 182.0])
POOR_ELEVATION_DEG = np.asarray([18.0, 22.0, 25.0, 29.0, 32.0, 27.0, 21.0, 16.0])


Profile = Callable[[float, int, float], float]


def _identity(_time_s: float, _satellite: int, baseline: float) -> float:
    return baseline


def _one(_time_s: float, _satellite: int, _baseline: float) -> float:
    return 1.0


def _zero(_time_s: float, _satellite: int, _baseline: float) -> float:
    return 0.0


@dataclass(frozen=True)
class Scenario:
    name: str
    description: str
    duration_s: float = 80.0
    poor_geometry: bool = False
    cn0_profile: Profile = _identity
    jitter_profile: Profile = _one
    discriminator_bias_profile: Profile = _zero
    pseudorange_bias_profile: Profile = _zero


@dataclass
class ChannelQuality:
    cn0_dbhz: float | None = None
    dll_ratio: float = 1.0
    dll_mean_z: float = 0.0

    def update(
        self,
        raw_cn0_dbhz: float,
        dll_samples: np.ndarray,
    ) -> None:
        if self.cn0_dbhz is None:
            self.cn0_dbhz = raw_cn0_dbhz
        else:
            self.cn0_dbhz = (
                FIRMWARE_CN0_ALPHA * raw_cn0_dbhz
                + (1.0 - FIRMWARE_CN0_ALPHA) * self.cn0_dbhz
            )

        measured_variance = float(np.var(dll_samples, ddof=1))
        expected_std = expected_discriminator_std(self.cn0_dbhz)
        raw_ratio = measured_variance / max(expected_std * expected_std, 1.0e-15)
        raw_ratio = float(np.clip(raw_ratio, 0.05, 25.0))
        self.dll_ratio = (
            DLL_RATIO_ALPHA * raw_ratio
            + (1.0 - DLL_RATIO_ALPHA) * self.dll_ratio
        )

        standard_error = np.sqrt(
            max(measured_variance, expected_std * expected_std)
            / len(dll_samples)
        )
        raw_mean_z = abs(float(np.mean(dll_samples))) / max(standard_error, 1.0e-12)
        raw_mean_z = float(np.clip(raw_mean_z, 0.0, 12.0))
        self.dll_mean_z = (
            DLL_MEAN_ALPHA * raw_mean_z
            + (1.0 - DLL_MEAN_ALPHA) * self.dll_mean_z
        )


@dataclass(frozen=True)
class EpochSample:
    method: SigmaMethod
    time_s: float
    enu_error_m: np.ndarray
    horizontal_error_m: float
    position_error_3d_m: float
    horizontal_sigma_m: float
    valid: bool
    chi_square_ratio: float
    fde_applied: bool
    mean_measurement_sigma_m: float
    maximum_dll_ratio: float


def scenario_matrix() -> tuple[Scenario, ...]:
    return (
        Scenario(
            "clean_open_sky",
            "C/N0 model and actual DLL jitter agree.",
        ),
        Scenario(
            "weak_signal",
            "All channels lose 11 dB; current C/N0 weighting already sees it.",
            cn0_profile=lambda _t, _s, baseline: baseline - 11.0,
        ),
        Scenario(
            "one_channel_hidden_jitter",
            "One high-C/N0 channel has 4x DLL jitter not visible in C/N0.",
            jitter_profile=lambda _t, satellite, _b: 4.0 if satellite == 1 else 1.0,
        ),
        Scenario(
            "common_hidden_jitter",
            "All channels have the same unmodelled 3x DLL jitter.",
            jitter_profile=lambda _t, _satellite, _b: 3.0,
        ),
        Scenario(
            "noise_burst_recovery",
            "Two channels have 5x jitter from 22 s to 45 s, then recover.",
            jitter_profile=lambda time_s, satellite, _b: (
                5.0
                if 22.0 <= time_s < 45.0 and satellite in (1, 5)
                else 1.0
            ),
        ),
        Scenario(
            "dll_transient_stress",
            "A channel has discriminator mean without excess variance during transient stress.",
            discriminator_bias_profile=lambda time_s, satellite, _b: (
                0.055
                if 25.0 <= time_s < 34.0 and satellite == 3
                else 0.0
            ),
            pseudorange_bias_profile=lambda time_s, satellite, _b: (
                9.0
                if 25.0 <= time_s < 34.0 and satellite == 3
                else 0.0
            ),
        ),
        Scenario(
            "stable_multipath_bias",
            "Low-jitter, high-C/N0 channel has a stable 14 m code bias.",
            pseudorange_bias_profile=lambda _t, satellite, _b: (
                14.0 if satellite == 6 else 0.0
            ),
        ),
        Scenario(
            "mixed_multipath",
            "One channel has both a 10 m bias and 3x excess jitter.",
            jitter_profile=lambda _t, satellite, _b: 3.0 if satellite == 2 else 1.0,
            pseudorange_bias_profile=lambda _t, satellite, _b: (
                10.0 if satellite == 2 else 0.0
            ),
        ),
        Scenario(
            "poor_geometry_hidden_jitter",
            "Clustered geometry plus one channel with unmodelled 4x jitter.",
            poor_geometry=True,
            jitter_profile=lambda _t, satellite, _b: 4.0 if satellite == 4 else 1.0,
        ),
        Scenario(
            "cn0_step_recovery",
            "C/N0 is initially weak and improves; no hidden DLL noise exists.",
            cn0_profile=lambda time_s, _satellite, baseline: (
                baseline - 12.0 if time_s < 30.0 else baseline
            ),
        ),
    )


def _geometry(scenario: Scenario) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    receiver = lla_to_ecef(39.9208, 32.8541, 930.0)
    if scenario.poor_geometry:
        azimuth = POOR_AZIMUTH_DEG
        elevation = POOR_ELEVATION_DEG
    else:
        azimuth = GOOD_AZIMUTH_DEG
        elevation = GOOD_ELEVATION_DEG
    satellites = satellites_from_azel(receiver, azimuth, elevation)
    return receiver, satellites, elevation.copy()


def run_scenario(
    scenario: Scenario,
    seed: int,
    methods: tuple[SigmaMethod, ...] = tuple(SigmaMethod),
) -> list[EpochSample]:
    rng = np.random.default_rng(seed)
    receiver, satellites, elevations = _geometry(scenario)
    satellite_count = len(satellites)
    geometric_ranges = np.linalg.norm(satellites - receiver, axis=1)
    baseline_cn0 = np.clip(34.0 + 0.16 * elevations, 30.0, 48.0)
    ura_indices = np.zeros(satellite_count, dtype=np.uint8)
    quality = [ChannelQuality() for _ in range(satellite_count)]

    # These are deliberately slower/non-DLL terms.  They prevent a synthetic
    # white-noise-only experiment from overstating what DLL weighting can fix.
    ura_error = rng.normal(0.0, 0.45 * URA_SIGMA_M[0], satellite_count)
    atmospheric_state = rng.normal(
        0.0,
        0.35 * elevation_sigma_m(elevations),
        satellite_count,
    )
    code_noise_state = np.zeros(satellite_count)
    initial_state = np.concatenate(
        [
            receiver + np.asarray([450.0, -300.0, 220.0]),
            [0.0],
        ]
    )
    method_state = {method: initial_state.copy() for method in methods}
    samples: list[EpochSample] = []
    clock_bias_m = 82.0
    receiver_enu_axes = enu_axes(receiver)

    tracking_steps = int(round(scenario.duration_s / TRACKING_PERIOD_S))
    latest_actual_jitter = np.ones(satellite_count)
    for tracking_step in range(tracking_steps):
        time_s = tracking_step * TRACKING_PERIOD_S
        for satellite in range(satellite_count):
            true_cn0 = scenario.cn0_profile(
                time_s,
                satellite,
                float(baseline_cn0[satellite]),
            )
            actual_jitter = scenario.jitter_profile(
                time_s,
                satellite,
                1.0,
            )
            discriminator_bias = scenario.discriminator_bias_profile(
                time_s,
                satellite,
                0.0,
            )
            latest_actual_jitter[satellite] = actual_jitter
            dll_window = synthesize_dll_window(
                rng,
                cn0_dbhz=true_cn0,
                jitter_multiplier=actual_jitter,
                discriminator_bias=discriminator_bias,
            )
            # The hardware NWPR C/N0 estimate is noisy; firmware applies alpha=0.2.
            raw_cn0 = float(np.clip(true_cn0 + rng.normal(0.0, 1.3), 10.0, 60.0))
            quality[satellite].update(raw_cn0, dll_window.samples)

        if (tracking_step + 1) % TRACKING_UPDATES_PER_WLS != 0:
            continue

        true_cn0 = np.asarray(
            [
                scenario.cn0_profile(
                    time_s,
                    satellite,
                    float(baseline_cn0[satellite]),
                )
                for satellite in range(satellite_count)
            ]
        )
        actual_code_sigma = cn0_code_sigma_m(true_cn0) * latest_actual_jitter
        innovation = rng.normal(0.0, actual_code_sigma)
        code_noise_state = (
            0.25 * code_noise_state
            + np.sqrt(1.0 - 0.25**2) * innovation
        )
        atmospheric_innovation = rng.normal(
            0.0,
            0.35 * elevation_sigma_m(elevations),
        )
        atmospheric_state = (
            0.97 * atmospheric_state
            + np.sqrt(1.0 - 0.97**2) * atmospheric_innovation
        )
        pseudorange_bias = np.asarray(
            [
                scenario.pseudorange_bias_profile(
                    time_s,
                    satellite,
                    0.0,
                )
                for satellite in range(satellite_count)
            ]
        )
        observations = (
            geometric_ranges
            + clock_bias_m
            + ura_error
            + atmospheric_state
            + code_noise_state
            + pseudorange_bias
        )
        sigma_inputs = SigmaInputs(
            ura_index=ura_indices,
            elevation_deg=elevations,
            cn0_dbhz=np.asarray([item.cn0_dbhz for item in quality], dtype=float),
            dll_variance_ratio=np.asarray([item.dll_ratio for item in quality]),
            dll_mean_z=np.asarray([item.dll_mean_z for item in quality]),
            actual_jitter_multiplier=latest_actual_jitter.copy(),
        )

        for method in methods:
            sigma = firmware_sigma_m(sigma_inputs, method)
            solution = firmware_wls(
                satellites,
                observations,
                sigma,
                method_state[method],
            )
            # NavEquations.c updates its persistent WLS seed only after an
            # integrity-valid solution.
            if solution.valid:
                method_state[method] = np.concatenate(
                    [solution.position_ecef_m, [solution.clock_bias_m]]
                )
            enu_error = position_error_enu_m(solution.position_ecef_m, receiver)
            enu_covariance = (
                # The top-left block is ECEF position covariance.
                # R * Q * R' maps it to ENU.
                receiver_enu_axes
                @ solution.covariance[:3, :3]
                @ receiver_enu_axes.T
            )
            horizontal_sigma = float(
                np.sqrt(max(0.0, enu_covariance[0, 0] + enu_covariance[1, 1]))
            )
            samples.append(
                EpochSample(
                    method=method,
                    time_s=time_s,
                    enu_error_m=enu_error,
                    horizontal_error_m=float(np.linalg.norm(enu_error[:2])),
                    position_error_3d_m=float(np.linalg.norm(enu_error)),
                    horizontal_sigma_m=horizontal_sigma,
                    valid=solution.valid,
                    chi_square_ratio=solution.chi_square_ratio,
                    fde_applied=solution.excluded_index is not None,
                    mean_measurement_sigma_m=float(np.mean(sigma)),
                    maximum_dll_ratio=float(
                        max(item.dll_ratio for item in quality)
                    ),
                )
            )
    return samples