"""Python mirror of the firmware's code-measurement WLS policy.

The model copies the receiver-specific parts of NavEquations.c:

* four states: ECEF x/y/z and common receiver clock bias in metres;
* URA + elevation + C/N0 variance model;
* normalized weights for the numerical solve and absolute weights for
  covariance/integrity;
* the 99-percent chi-square table;
* single-satellite leave-one-out FDE only when at least seven measurements
  are available.

Atmospheric and satellite-clock corrections are represented as already
corrected measurement errors.  That keeps this study focused on tracking
quality rather than duplicating the existing PPS orbit/atmosphere model.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum

import numpy as np


URA_SIGMA_M = np.asarray(
    [
        2.4,
        3.4,
        4.85,
        6.85,
        9.65,
        13.65,
        24.0,
        48.0,
        96.0,
        192.0,
        384.0,
        768.0,
        1536.0,
        3072.0,
        6144.0,
    ],
    dtype=float,
)
CHI_SQUARE_99 = np.asarray(
    [0.0, 6.634896601, 9.210340372, 11.34486673, 13.27670414, 15.08627247],
    dtype=float,
)


class SigmaMethod(str, Enum):
    CURRENT = "current_cn0"
    NAIVE_ADD_DLL = "naive_add_dll"
    DLL_EXCESS = "dll_excess"
    DLL_EXCESS_GUARD = "dll_excess_guard"
    ORACLE = "oracle_variance"


@dataclass(frozen=True)
class SigmaInputs:
    ura_index: np.ndarray
    elevation_deg: np.ndarray
    cn0_dbhz: np.ndarray
    dll_variance_ratio: np.ndarray
    dll_mean_z: np.ndarray
    actual_jitter_multiplier: np.ndarray


@dataclass(frozen=True)
class WlsSolution:
    position_ecef_m: np.ndarray
    clock_bias_m: float
    residual_m: np.ndarray
    sigma_m: np.ndarray
    covariance: np.ndarray
    chi_square: float
    chi_square_ratio: float
    chi_square_pass: bool
    valid: bool
    used_indices: np.ndarray
    excluded_index: int | None


def cn0_code_sigma_m(cn0_dbhz: np.ndarray | float) -> np.ndarray:
    sigma = 3.0 * np.power(10.0, (45.0 - np.asarray(cn0_dbhz)) / 20.0)
    return np.clip(sigma, 0.5, 30.0)


def elevation_sigma_m(elevation_deg: np.ndarray | float) -> np.ndarray:
    bounded = np.clip(np.asarray(elevation_deg), 10.0, 90.0)
    return 3.0 / np.sin(np.deg2rad(bounded))


def firmware_sigma_m(
    inputs: SigmaInputs,
    method: SigmaMethod,
) -> np.ndarray:
    ura = URA_SIGMA_M[np.asarray(inputs.ura_index, dtype=int)]
    elevation = elevation_sigma_m(inputs.elevation_deg)
    code = cn0_code_sigma_m(inputs.cn0_dbhz)
    ratio = np.clip(np.asarray(inputs.dll_variance_ratio), 0.05, 25.0)

    if method == SigmaMethod.CURRENT:
        code_variance_factor = np.ones_like(code)
    elif method == SigmaMethod.NAIVE_ADD_DLL:
        # Deliberately included as a negative control: at ratio=1 this counts
        # the C/N0-predicted thermal jitter twice.
        code_variance_factor = 1.0 + ratio
    elif method == SigmaMethod.DLL_EXCESS:
        # Only the measured excess over the C/N0 expectation is added.
        code_variance_factor = np.maximum(1.0, ratio)
    elif method == SigmaMethod.DLL_EXCESS_GUARD:
        mean_factor = np.square(
            np.maximum(1.0, np.asarray(inputs.dll_mean_z) / 3.0)
        )
        code_variance_factor = np.maximum.reduce(
            [np.ones_like(code), ratio, mean_factor]
        )
    elif method == SigmaMethod.ORACLE:
        code_variance_factor = np.square(
            np.maximum(0.05, np.asarray(inputs.actual_jitter_multiplier))
        )
    else:
        raise ValueError(f"unsupported sigma method: {method}")

    return np.sqrt(
        ura * ura
        + elevation * elevation
        + code * code * code_variance_factor
    )


def ecef_to_lla(ecef_m: np.ndarray) -> tuple[float, float, float]:
    x, y, z = np.asarray(ecef_m, dtype=float)
    a = 6_378_137.0
    inverse_flattening = 298.257223563
    flattening = 1.0 / inverse_flattening
    eccentricity_squared = flattening * (2.0 - flattening)
    longitude = np.arctan2(y, x)
    horizontal = np.hypot(x, y)
    latitude = np.arctan2(
        z,
        horizontal * (1.0 - eccentricity_squared),
    )
    height = 0.0
    for _ in range(12):
        sine = np.sin(latitude)
        prime_vertical = a / np.sqrt(1.0 - eccentricity_squared * sine * sine)
        height = horizontal / max(np.cos(latitude), 1.0e-15) - prime_vertical
        next_latitude = np.arctan2(
            z,
            horizontal
            * (
                1.0
                - eccentricity_squared
                * prime_vertical
                / max(prime_vertical + height, 1.0)
            ),
        )
        if abs(next_latitude - latitude) < 1.0e-14:
            latitude = next_latitude
            break
        latitude = next_latitude
    return (
        float(np.rad2deg(latitude)),
        float(np.rad2deg(longitude)),
        float(height),
    )


def lla_to_ecef(latitude_deg: float, longitude_deg: float, height_m: float) -> np.ndarray:
    a = 6_378_137.0
    inverse_flattening = 298.257223563
    flattening = 1.0 / inverse_flattening
    eccentricity_squared = flattening * (2.0 - flattening)
    latitude = np.deg2rad(latitude_deg)
    longitude = np.deg2rad(longitude_deg)
    sine = np.sin(latitude)
    cosine = np.cos(latitude)
    prime_vertical = a / np.sqrt(1.0 - eccentricity_squared * sine * sine)
    return np.asarray(
        [
            (prime_vertical + height_m) * cosine * np.cos(longitude),
            (prime_vertical + height_m) * cosine * np.sin(longitude),
            (
                prime_vertical * (1.0 - eccentricity_squared)
                + height_m
            )
            * sine,
        ],
        dtype=float,
    )


def enu_axes(receiver_ecef_m: np.ndarray) -> np.ndarray:
    latitude_deg, longitude_deg, _ = ecef_to_lla(receiver_ecef_m)
    latitude = np.deg2rad(latitude_deg)
    longitude = np.deg2rad(longitude_deg)
    east = np.asarray([-np.sin(longitude), np.cos(longitude), 0.0])
    north = np.asarray(
        [
            -np.sin(latitude) * np.cos(longitude),
            -np.sin(latitude) * np.sin(longitude),
            np.cos(latitude),
        ]
    )
    up = np.asarray(
        [
            np.cos(latitude) * np.cos(longitude),
            np.cos(latitude) * np.sin(longitude),
            np.sin(latitude),
        ]
    )
    return np.vstack([east, north, up])


def satellites_from_azel(
    receiver_ecef_m: np.ndarray,
    azimuth_deg: np.ndarray,
    elevation_deg: np.ndarray,
    range_m: float = 21_000_000.0,
) -> np.ndarray:
    axes = enu_axes(receiver_ecef_m)
    satellites: list[np.ndarray] = []
    for azimuth, elevation in zip(azimuth_deg, elevation_deg):
        azimuth_rad = np.deg2rad(azimuth)
        elevation_rad = np.deg2rad(elevation)
        los_enu = np.asarray(
            [
                np.cos(elevation_rad) * np.sin(azimuth_rad),
                np.cos(elevation_rad) * np.cos(azimuth_rad),
                np.sin(elevation_rad),
            ]
        )
        los_ecef = axes.T @ los_enu
        satellites.append(receiver_ecef_m + range_m * los_ecef)
    return np.asarray(satellites)


def position_error_enu_m(
    estimated_ecef_m: np.ndarray,
    true_ecef_m: np.ndarray,
) -> np.ndarray:
    return enu_axes(true_ecef_m) @ (estimated_ecef_m - true_ecef_m)


def _candidate(
    satellite_positions_m: np.ndarray,
    observations_m: np.ndarray,
    sigma_m: np.ndarray,
    initial_state: np.ndarray,
    used_indices: np.ndarray,
) -> WlsSolution:
    state = np.asarray(initial_state, dtype=float).copy()
    weights = 1.0 / np.square(sigma_m)
    design = np.empty((len(observations_m), 4), dtype=float)

    for _ in range(10):
        difference = satellite_positions_m - state[:3]
        geometric_range = np.linalg.norm(difference, axis=1)
        design[:, :3] = -difference / geometric_range[:, None]
        design[:, 3] = 1.0
        residual = observations_m - geometric_range - state[3]
        normal = design.T @ (weights[:, None] * design)
        rhs = design.T @ (weights * residual)
        update = np.linalg.solve(normal, rhs)
        state += update
        if np.linalg.norm(update) < 1.0e-4:
            break

    difference = satellite_positions_m - state[:3]
    geometric_range = np.linalg.norm(difference, axis=1)
    design[:, :3] = -difference / geometric_range[:, None]
    design[:, 3] = 1.0
    residual = observations_m - geometric_range - state[3]
    normal = design.T @ (weights[:, None] * design)
    covariance = np.linalg.inv(normal)
    chi_square = float(np.sum(weights * np.square(residual)))
    degrees_of_freedom = len(observations_m) - 4
    threshold = (
        float(CHI_SQUARE_99[degrees_of_freedom])
        if 0 < degrees_of_freedom < len(CHI_SQUARE_99)
        else 0.0
    )
    chi_square_pass = degrees_of_freedom > 0 and chi_square <= threshold
    valid = (
        (len(observations_m) == 4 and degrees_of_freedom == 0)
        or chi_square_pass
    )
    ratio = chi_square / threshold if threshold > 0.0 else float("inf")
    return WlsSolution(
        position_ecef_m=state[:3].copy(),
        clock_bias_m=float(state[3]),
        residual_m=residual.copy(),
        sigma_m=sigma_m.copy(),
        covariance=covariance,
        chi_square=chi_square,
        chi_square_ratio=ratio,
        chi_square_pass=chi_square_pass,
        valid=valid,
        used_indices=used_indices.copy(),
        excluded_index=None,
    )


def firmware_wls(
    satellite_positions_m: np.ndarray,
    observations_m: np.ndarray,
    sigma_m: np.ndarray,
    initial_state: np.ndarray,
    *,
    enable_fde: bool = True,
) -> WlsSolution:
    count = len(observations_m)
    if count < 4:
        raise ValueError("position WLS needs at least four satellites")
    full_indices = np.arange(count, dtype=int)
    full = _candidate(
        satellite_positions_m,
        observations_m,
        sigma_m,
        initial_state,
        full_indices,
    )
    if full.chi_square_pass or not enable_fde or count < 7:
        return full

    selected = full
    passing_found = False
    full_state = np.concatenate([full.position_ecef_m, [full.clock_bias_m]])
    for excluded in range(count):
        used = full_indices[full_indices != excluded]
        candidate = _candidate(
            satellite_positions_m[used],
            observations_m[used],
            sigma_m[used],
            full_state,
            used,
        )
        candidate = WlsSolution(
            **{
                **candidate.__dict__,
                "excluded_index": excluded,
            }
        )
        if candidate.chi_square_pass:
            if not passing_found or candidate.chi_square_ratio < selected.chi_square_ratio:
                selected = candidate
            passing_found = True
        elif not passing_found and candidate.chi_square_ratio < selected.chi_square_ratio:
            selected = candidate
    return selected