"""Project-specific Doppler WLS to ESKF timing/noise model.

The model compares four paths using one shared 50 Hz tracking stream and one
shared 100 Hz IMU stream:

* the legacy fixed-alpha Doppler IIR before 5 Hz velocity WLS;
* raw Doppler into WLS, as a nominal-model comparison rather than production;
* a global IMU-driven adaptive-alpha IIR;
* raw Doppler with the per-satellite robust WLS gate used before the ESKF.

The ESKF is reduced to the velocity and accelerometer-bias error states that
matter to this experiment.  Its process noise, 3-D measurement covariance and
99.9-percent velocity NIS gate mirror the project firmware.  Attitude and
position are intentionally outside this model's scope.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Callable

import numpy as np


IMU_HZ = 100.0
TRACKING_HZ = 50.0
VELOCITY_WLS_HZ = 5.0
IMU_DT_S = 1.0 / IMU_HZ
TRACKING_STEPS = int(round(IMU_HZ / TRACKING_HZ))
WLS_STEPS = int(round(IMU_HZ / VELOCITY_WLS_HZ))
FIXED_DOPPLER_ALPHA = 0.05
CN0_ALPHA = 0.20
VELOCITY_NIS_GATE = 16.26623619623813
PROCESS_ACCEL_NOISE_DENSITY = 0.025
ACCEL_BIAS_RANDOM_WALK = 0.0015
MINIMUM_HORIZONTAL_VELOCITY_STD_MPS = 2.0
MINIMUM_VERTICAL_VELOCITY_STD_MPS = 3.0


class DopplerPath(str, Enum):
    FIXED_IIR = "fixed_iir_wls_eskf"
    RAW = "raw_wls_eskf"
    ADAPTIVE_IIR = "adaptive_iir_wls_eskf"
    RAW_ROBUST = "raw_robust_wls_eskf"


PATH_ORDER = (
    DopplerPath.FIXED_IIR,
    DopplerPath.RAW,
    DopplerPath.ADAPTIVE_IIR,
    DopplerPath.RAW_ROBUST,
)


@dataclass(frozen=True)
class FlightScenario:
    name: str
    description: str
    duration_s: float
    initial_velocity_mps: np.ndarray
    acceleration: Callable[[float], np.ndarray]
    cn0_offset: Callable[[float], float]
    noise_multiplier: Callable[[float, int], float]
    outlier_mps: Callable[[float, int], float]
    receiver_clock_drift_mps: Callable[[float], float]
    gps_gaps: tuple[tuple[float, float], ...] = ()


@dataclass(frozen=True)
class VelocityWlsResult:
    velocity_mps: np.ndarray
    clock_drift_mps: float
    covariance_mps2: np.ndarray
    residual_mps: np.ndarray
    residual_rms_mps: float
    covariance_scale: float
    excluded_satellite: int | None


@dataclass
class AdaptiveAlpha:
    alpha: float = FIXED_DOPPLER_ALPHA
    hold_samples: int = 0
    filtered_acceleration_mps2: np.ndarray = field(
        default_factory=lambda: np.zeros(3, dtype=float)
    )

    def update(
        self,
        measured_acceleration_mps2: np.ndarray,
        median_cn0_dbhz: float,
    ) -> float:
        # A raw 100 Hz accelerometer sample is too noisy for coefficient
        # switching.  This 50 Hz vector EWMA has an approximately 0.16 s time
        # constant and, unlike smoothing |a|, does not retain the positive
        # magnitude bias of zero-mean sensor noise.
        motion_alpha = 0.12
        self.filtered_acceleration_mps2 += motion_alpha * (
            measured_acceleration_mps2 - self.filtered_acceleration_mps2
        )
        magnitude = float(np.linalg.norm(self.filtered_acceleration_mps2))

        if median_cn0_dbhz >= 46.0:
            quiet_alpha = 0.08
        elif median_cn0_dbhz >= 37.0:
            quiet_alpha = FIXED_DOPPLER_ALPHA
        else:
            quiet_alpha = 0.03

        if magnitude >= 2.5:
            self.alpha = 0.40
            self.hold_samples = 25
        elif magnitude >= 0.55:
            self.alpha = 0.20
            self.hold_samples = 25
        elif self.hold_samples > 0:
            self.hold_samples -= 1
        else:
            self.alpha = quiet_alpha
        return self.alpha


@dataclass
class VelocityBiasEskf:
    velocity_mps: np.ndarray
    accel_bias_mps2: np.ndarray
    covariance: np.ndarray
    accepted: int = 0
    rejected_nis: int = 0
    last_nis: float = float("nan")

    @classmethod
    def initialize(
        cls,
        initial_velocity_mps: np.ndarray,
        initial_velocity_error_mps: np.ndarray,
    ) -> "VelocityBiasEskf":
        covariance = np.zeros((6, 6), dtype=float)
        covariance[:3, :3] = np.eye(3)
        covariance[3:, 3:] = np.eye(3) * 0.1**2
        return cls(
            velocity_mps=initial_velocity_mps + initial_velocity_error_mps,
            accel_bias_mps2=np.zeros(3),
            covariance=covariance,
        )

    def propagate(self, measured_acceleration_mps2: np.ndarray, dt_s: float) -> None:
        self.velocity_mps += (
            measured_acceleration_mps2 - self.accel_bias_mps2
        ) * dt_s

        transition = np.eye(6)
        transition[:3, 3:] = -np.eye(3) * dt_s
        accel_variance = PROCESS_ACCEL_NOISE_DENSITY**2
        bias_variance = ACCEL_BIAS_RANDOM_WALK**2
        process = np.zeros((6, 6), dtype=float)
        # Match eskf.c's first-order discrete injection: velocity and bias
        # driving noise are added after F P F^T, without exact-discretization
        # dt^2/dt^3 cross terms.
        process[:3, :3] = np.eye(3) * accel_variance * dt_s
        process[3:, 3:] = np.eye(3) * bias_variance * dt_s
        self.covariance = (
            transition @ self.covariance @ transition.T + process
        )
        self.covariance = 0.5 * (self.covariance + self.covariance.T)

    def update_velocity(
        self,
        measurement_mps: np.ndarray,
        covariance_mps2: np.ndarray,
    ) -> bool:
        innovation = measurement_mps - self.velocity_mps
        innovation_covariance = self.covariance[:3, :3] + covariance_mps2
        inverse = np.linalg.inv(innovation_covariance)
        nis = float(innovation @ inverse @ innovation)
        self.last_nis = nis
        if not np.isfinite(nis) or nis > VELOCITY_NIS_GATE:
            self.rejected_nis += 1
            return False

        measurement_matrix = np.zeros((3, 6), dtype=float)
        measurement_matrix[:, :3] = np.eye(3)
        gain = self.covariance @ measurement_matrix.T @ inverse
        correction = gain @ innovation
        self.velocity_mps += correction[:3]
        self.accel_bias_mps2 += correction[3:]
        identity = np.eye(6)
        residual_transform = identity - gain @ measurement_matrix
        self.covariance = (
            residual_transform
            @ self.covariance
            @ residual_transform.T
            + gain @ covariance_mps2 @ gain.T
        )
        self.covariance = 0.5 * (self.covariance + self.covariance.T)
        self.accepted += 1
        return True


@dataclass(frozen=True)
class DopplerEskfTrace:
    scenario: str
    path: DopplerPath
    time_s: np.ndarray
    truth_velocity_mps: np.ndarray
    fused_velocity_mps: np.ndarray
    fused_velocity_covariance_mps2: np.ndarray
    acceleration_mps2: np.ndarray
    wls_time_s: np.ndarray
    wls_velocity_mps: np.ndarray
    wls_clock_drift_mps: np.ndarray
    wls_covariance_mps2: np.ndarray
    wls_covariance_scale: np.ndarray
    wls_nis: np.ndarray
    wls_accepted: np.ndarray
    adaptive_alpha: np.ndarray
    excluded_satellites: np.ndarray


def _zero_acceleration(_time_s: float) -> np.ndarray:
    return np.zeros(3)


def _airliner_acceleration(time_s: float) -> np.ndarray:
    if 15.0 <= time_s < 25.0:
        return np.asarray([2.0, 0.0, 0.0])
    if 38.0 <= time_s < 52.0:
        return np.asarray([0.0, 1.5, 0.25])
    if 68.0 <= time_s < 76.0:
        return np.asarray([-2.5, 0.0, -0.15])
    return np.zeros(3)


def _aggressive_acceleration(time_s: float) -> np.ndarray:
    if 10.0 <= time_s < 16.0:
        return np.asarray([6.0, 0.0, 1.0])
    if 26.0 <= time_s < 33.0:
        return np.asarray([0.0, 7.0, 0.0])
    if 46.0 <= time_s < 52.0:
        return np.asarray([-5.0, -4.0, -1.0])
    if 64.0 <= time_s < 76.0:
        phase = 2.0 * np.pi * (time_s - 64.0) / 3.0
        return np.asarray([3.0 * np.sin(phase), 4.0 * np.cos(phase), 0.5 * np.sin(phase)])
    return np.zeros(3)


def _constant_zero(_time_s: float) -> float:
    return 0.0


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


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


def flight_scenarios() -> tuple[FlightScenario, ...]:
    return (
        FlightScenario(
            "steady_cruise",
            "Constant aircraft velocity with nominal open-sky Doppler noise.",
            90.0,
            np.asarray([180.0, 20.0, 0.0]),
            _zero_acceleration,
            _constant_zero,
            _constant_one,
            _no_outlier,
            _constant_zero,
        ),
        FlightScenario(
            "airliner_maneuver",
            "Moderate longitudinal, turn and climb-rate accelerations.",
            90.0,
            np.asarray([120.0, 0.0, 0.0]),
            _airliner_acceleration,
            _constant_zero,
            _constant_one,
            _no_outlier,
            _constant_zero,
        ),
        FlightScenario(
            "aggressive_maneuver",
            "High acceleration and oscillatory manoeuvre segments.",
            90.0,
            np.asarray([80.0, 0.0, 0.0]),
            _aggressive_acceleration,
            _constant_zero,
            _constant_one,
            _no_outlier,
            _constant_zero,
        ),
        FlightScenario(
            "low_cn0_airliner",
            "Airliner manoeuvre with a 10 dB signal-quality loss.",
            90.0,
            np.asarray([120.0, 0.0, 0.0]),
            _airliner_acceleration,
            lambda _time_s: -10.0,
            _constant_one,
            _no_outlier,
            _constant_zero,
        ),
        FlightScenario(
            "single_satellite_outliers",
            "One channel has repeatable 12 m/s Doppler bursts.",
            90.0,
            np.asarray([120.0, 0.0, 0.0]),
            _airliner_acceleration,
            _constant_zero,
            _constant_one,
            lambda time_s, satellite: (
                12.0
                if satellite == 2
                and (
                    30.0 <= time_s < 31.0
                    or 56.0 <= time_s < 57.0
                    or 72.0 <= time_s < 72.6
                )
                else 0.0
            ),
            _constant_zero,
        ),
        FlightScenario(
            "clock_drift_steps",
            "Common receiver clock-frequency steps coexist with manoeuvres.",
            90.0,
            np.asarray([120.0, 0.0, 0.0]),
            _airliner_acceleration,
            _constant_zero,
            _constant_one,
            _no_outlier,
            lambda time_s: (
                18.0 if 28.0 <= time_s < 55.0 else (-9.0 if time_s >= 55.0 else 0.0)
            ),
        ),
        FlightScenario(
            "gps_gap_recovery",
            "Fifteen-second GNSS gap during an aircraft manoeuvre.",
            90.0,
            np.asarray([120.0, 0.0, 0.0]),
            _airliner_acceleration,
            _constant_zero,
            _constant_one,
            _no_outlier,
            _constant_zero,
            gps_gaps=((42.0, 57.0),),
        ),
        FlightScenario(
            "mixed_flight_faults",
            "Aggressive motion, temporary weak signal, channel bursts and clock step.",
            90.0,
            np.asarray([80.0, 0.0, 0.0]),
            _aggressive_acceleration,
            lambda time_s: -9.0 if 34.0 <= time_s < 60.0 else 0.0,
            lambda time_s, satellite: (
                2.5 if 34.0 <= time_s < 60.0 and satellite in (1, 6) else 1.0
            ),
            lambda time_s, satellite: (
                -10.0 if satellite == 4 and 54.0 <= time_s < 55.0 else 0.0
            ),
            lambda time_s: 12.0 if time_s >= 48.0 else 0.0,
            gps_gaps=((78.0, 83.0),),
        ),
    )


def local_los_matrix() -> np.ndarray:
    azimuth_deg = np.asarray([8.0, 55.0, 104.0, 151.0, 205.0, 253.0, 301.0, 342.0])
    elevation_deg = np.asarray([73.0, 48.0, 31.0, 18.0, 64.0, 39.0, 24.0, 14.0])
    azimuth = np.deg2rad(azimuth_deg)
    elevation = np.deg2rad(elevation_deg)
    return np.column_stack(
        [
            np.cos(elevation) * np.sin(azimuth),
            np.cos(elevation) * np.cos(azimuth),
            np.sin(elevation),
        ]
    )


def velocity_sigma_mps(cn0_dbhz: np.ndarray) -> np.ndarray:
    return np.clip(30.0 * np.power(10.0, -cn0_dbhz / 20.0), 0.05, 5.0)


def _solve_velocity_candidate(
    design: np.ndarray,
    observations_mps: np.ndarray,
    sigma_mps: np.ndarray,
    indices: np.ndarray,
) -> VelocityWlsResult:
    selected_design = design[indices]
    selected_observations = observations_mps[indices]
    selected_sigma = sigma_mps[indices]
    weights = 1.0 / np.square(selected_sigma)
    normal = selected_design.T @ (weights[:, None] * selected_design)
    inverse = np.linalg.inv(normal)
    solution = inverse @ (selected_design.T @ (weights * selected_observations))
    residual = selected_observations - selected_design @ solution
    weighted_residual_square_sum = float(np.sum(weights * np.square(residual)))
    degrees_of_freedom = len(indices) - 4
    covariance_scale = (
        max(1.0, weighted_residual_square_sum / degrees_of_freedom)
        if degrees_of_freedom > 0
        else 1.0
    )
    return VelocityWlsResult(
        velocity_mps=solution[:3].copy(),
        clock_drift_mps=float(solution[3]),
        covariance_mps2=covariance_scale * inverse[:3, :3],
        residual_mps=residual,
        residual_rms_mps=float(np.sqrt(np.mean(np.square(residual)))),
        covariance_scale=covariance_scale,
        excluded_satellite=None,
    )


def velocity_wls(
    observations_mps: np.ndarray,
    cn0_dbhz: np.ndarray,
    *,
    robust: bool = False,
) -> VelocityWlsResult:
    los = local_los_matrix()
    design = np.column_stack([-los, np.ones(len(los))])
    sigma = velocity_sigma_mps(cn0_dbhz)
    indices = np.arange(len(observations_mps), dtype=int)
    full = _solve_velocity_candidate(design, observations_mps, sigma, indices)
    if not robust or len(indices) < 6:
        return full

    normalized = np.abs(full.residual_mps) / sigma
    if float(np.max(normalized)) <= 3.5:
        return full

    # The largest full-fit residual is not necessarily the bad satellite:
    # WLS leverage can distribute one channel's error over the other rows.
    # Compare every leave-one-out hypothesis and use the one with the lowest
    # normalized post-fit residual energy.  The firmware robust velocity WLS
    # uses the same trigger and candidate selection.
    candidates: list[tuple[float, int, VelocityWlsResult]] = []
    for excluded in indices:
        used = indices[indices != excluded]
        candidate = _solve_velocity_candidate(design, observations_mps, sigma, used)
        score = float(
            np.sum(np.square(candidate.residual_mps / sigma[used]))
        )
        candidates.append((score, int(excluded), candidate))
    _score, excluded, candidate = min(candidates, key=lambda item: item[0])
    return VelocityWlsResult(
        **{
            **candidate.__dict__,
            "excluded_satellite": excluded,
        }
    )


def fusion_velocity_covariance(covariance_mps2: np.ndarray) -> np.ndarray:
    """Apply fusion_gps_adapter.c's ENU velocity covariance sanitizer floors."""

    output = np.asarray(covariance_mps2, dtype=float).copy()
    variance_floor = np.square(
        np.asarray(
            [
                MINIMUM_HORIZONTAL_VELOCITY_STD_MPS,
                MINIMUM_HORIZONTAL_VELOCITY_STD_MPS,
                MINIMUM_VERTICAL_VELOCITY_STD_MPS,
            ]
        )
    )
    for axis in range(3):
        output[axis, axis] = max(output[axis, axis], variance_floor[axis])
    for row in range(3):
        for column in range(row + 1, 3):
            covariance = 0.5 * (output[row, column] + output[column, row])
            limit = 0.95 * np.sqrt(output[row, row] * output[column, column])
            covariance = float(np.clip(covariance, -limit, limit))
            output[row, column] = covariance
            output[column, row] = covariance
    return output


def _in_gap(scenario: FlightScenario, time_s: float) -> bool:
    return any(start <= time_s < stop for start, stop in scenario.gps_gaps)


def _truth_trajectory(scenario: FlightScenario) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    time_s = np.arange(0.0, scenario.duration_s, IMU_DT_S)
    acceleration = np.asarray([scenario.acceleration(float(time)) for time in time_s])
    velocity = np.empty((len(time_s), 3), dtype=float)
    velocity[0] = scenario.initial_velocity_mps
    velocity[1:] = velocity[0] + np.cumsum(acceleration[:-1] * IMU_DT_S, axis=0)
    return time_s, acceleration, velocity


def run_doppler_eskf_scenario(
    scenario: FlightScenario,
    seed: int,
    paths: tuple[DopplerPath, ...] = PATH_ORDER,
) -> dict[DopplerPath, DopplerEskfTrace]:
    rng = np.random.default_rng(seed)
    time_s, truth_acceleration, truth_velocity = _truth_trajectory(scenario)
    satellite_count = len(local_los_matrix())
    baseline_cn0 = np.asarray([47.0, 45.0, 43.0, 38.0, 46.0, 42.0, 40.0, 36.0])
    filtered_cn0 = baseline_cn0.copy()
    fixed_measurement: np.ndarray | None = None
    adaptive_measurement: np.ndarray | None = None
    latest_raw: np.ndarray | None = None
    alpha_selector = AdaptiveAlpha()
    latest_alpha = FIXED_DOPPLER_ALPHA

    initial_error = rng.normal(0.0, 0.8, 3)
    filters = {
        path: VelocityBiasEskf.initialize(
            scenario.initial_velocity_mps,
            initial_error,
        )
        for path in paths
    }
    fused_history = {
        path: np.empty((len(time_s), 3), dtype=float)
        for path in paths
    }
    fused_covariance_history = {
        path: np.empty((len(time_s), 3, 3), dtype=float)
        for path in paths
    }
    wls_times = {path: [] for path in paths}
    wls_velocity = {path: [] for path in paths}
    wls_clock_drift = {path: [] for path in paths}
    wls_covariance = {path: [] for path in paths}
    wls_covariance_scale = {path: [] for path in paths}
    wls_nis = {path: [] for path in paths}
    wls_accepted = {path: [] for path in paths}
    wls_alpha = {path: [] for path in paths}
    excluded = {path: [] for path in paths}

    true_accel_bias = rng.normal(0.0, 0.04, 3)
    latest_acceleration_measurement = truth_acceleration[0] + true_accel_bias
    los = local_los_matrix()

    for sample in range(len(time_s)):
        if sample > 0:
            true_accel_bias += rng.normal(
                0.0,
                ACCEL_BIAS_RANDOM_WALK * np.sqrt(IMU_DT_S),
                3,
            )
            imu_noise_std = PROCESS_ACCEL_NOISE_DENSITY / np.sqrt(IMU_DT_S)
            latest_acceleration_measurement = (
                truth_acceleration[sample - 1]
                + true_accel_bias
                + rng.normal(0.0, imu_noise_std, 3)
            )
            for path in paths:
                filters[path].propagate(latest_acceleration_measurement, IMU_DT_S)

        if sample % TRACKING_STEPS == 0:
            current_time = float(time_s[sample])
            true_cn0 = baseline_cn0 + scenario.cn0_offset(current_time)
            raw_cn0 = np.clip(true_cn0 + rng.normal(0.0, 1.2, satellite_count), 10.0, 60.0)
            filtered_cn0 = CN0_ALPHA * raw_cn0 + (1.0 - CN0_ALPHA) * filtered_cn0
            nominal_sigma = velocity_sigma_mps(true_cn0)
            actual_sigma = nominal_sigma * np.asarray(
                [
                    scenario.noise_multiplier(current_time, satellite)
                    for satellite in range(satellite_count)
                ]
            )
            raw = (
                -los @ truth_velocity[sample]
                + scenario.receiver_clock_drift_mps(current_time)
                + rng.normal(0.0, actual_sigma)
                + np.asarray(
                    [
                        scenario.outlier_mps(current_time, satellite)
                        for satellite in range(satellite_count)
                    ]
                )
            )
            if fixed_measurement is None:
                fixed_measurement = raw.copy()
                adaptive_measurement = raw.copy()
            else:
                fixed_measurement += FIXED_DOPPLER_ALPHA * (raw - fixed_measurement)
                latest_alpha = alpha_selector.update(
                    latest_acceleration_measurement,
                    float(np.median(filtered_cn0)),
                )
                adaptive_measurement += latest_alpha * (raw - adaptive_measurement)
            latest_raw = raw

        if (
            sample > 0
            and sample % WLS_STEPS == 0
            and latest_raw is not None
            and not _in_gap(scenario, float(time_s[sample]))
        ):
            for path in paths:
                if path == DopplerPath.FIXED_IIR:
                    observations = fixed_measurement
                elif path == DopplerPath.ADAPTIVE_IIR:
                    observations = adaptive_measurement
                else:
                    observations = latest_raw
                result = velocity_wls(
                    np.asarray(observations),
                    filtered_cn0,
                    robust=(path == DopplerPath.RAW_ROBUST),
                )
                accepted = filters[path].update_velocity(
                    result.velocity_mps,
                    fusion_velocity_covariance(result.covariance_mps2),
                )
                wls_times[path].append(float(time_s[sample]))
                wls_velocity[path].append(result.velocity_mps)
                wls_clock_drift[path].append(result.clock_drift_mps)
                wls_covariance[path].append(result.covariance_mps2)
                wls_covariance_scale[path].append(result.covariance_scale)
                wls_nis[path].append(filters[path].last_nis)
                wls_accepted[path].append(accepted)
                wls_alpha[path].append(
                    latest_alpha
                    if path == DopplerPath.ADAPTIVE_IIR
                    else (FIXED_DOPPLER_ALPHA if path == DopplerPath.FIXED_IIR else 1.0)
                )
                excluded[path].append(
                    -1 if result.excluded_satellite is None else result.excluded_satellite
                )

        for path in paths:
            fused_history[path][sample] = filters[path].velocity_mps
            fused_covariance_history[path][sample] = filters[path].covariance[:3, :3]

    return {
        path: DopplerEskfTrace(
            scenario=scenario.name,
            path=path,
            time_s=time_s,
            truth_velocity_mps=truth_velocity,
            fused_velocity_mps=fused_history[path],
            fused_velocity_covariance_mps2=fused_covariance_history[path],
            acceleration_mps2=truth_acceleration,
            wls_time_s=np.asarray(wls_times[path]),
            wls_velocity_mps=np.asarray(wls_velocity[path]),
            wls_clock_drift_mps=np.asarray(wls_clock_drift[path]),
            wls_covariance_mps2=np.asarray(wls_covariance[path]),
            wls_covariance_scale=np.asarray(wls_covariance_scale[path]),
            wls_nis=np.asarray(wls_nis[path]),
            wls_accepted=np.asarray(wls_accepted[path], dtype=bool),
            adaptive_alpha=np.asarray(wls_alpha[path]),
            excluded_satellites=np.asarray(excluded[path], dtype=int),
        )
        for path in paths
    }