"""One-dimensional code smoothing study with carrier/Doppler aiding."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class SmoothingResult:
    raw_rmse_m: float
    scalar_iir_rmse_m: float
    doppler_smoothed_rmse_m: float
    firmware_iir_doppler_smoothed_rmse_m: float
    scalar_iir_peak_dynamic_error_m: float
    doppler_peak_dynamic_error_m: float
    firmware_iir_doppler_peak_dynamic_error_m: float


def run_smoothing_trial(
    seed: int,
    *,
    duration_s: float = 80.0,
    tracking_dt_s: float = 0.020,
    wls_dt_s: float = 0.200,
    code_sigma_m: float = 4.0,
    doppler_sigma_mps: float = 0.25,
    doppler_cfo_mps: float = 0.0,
    alpha: float = 0.10,
    firmware_doppler_alpha: float = 0.05,
) -> SmoothingResult:
    """Compare raw code, blind IIR and Doppler-propagated code smoothing.

    The truth includes a range-rate step so a visually smooth but lagging
    scalar IIR is penalized.  ``doppler_cfo_mps`` represents residual receiver
    carrier frequency offset after the nominal IF subtraction.
    """

    if abs(round(wls_dt_s / tracking_dt_s) - wls_dt_s / tracking_dt_s) > 1.0e-12:
        raise ValueError("WLS period must be an integer number of tracking periods")
    rng = np.random.default_rng(seed)
    tracking_time_s = np.arange(0.0, duration_s, tracking_dt_s)
    true_rate_tracking = np.where(
        tracking_time_s < 20.0,
        -120.0,
        np.where(tracking_time_s < 45.0, 230.0, -40.0),
    )
    true_range_tracking = np.empty_like(tracking_time_s)
    true_range_tracking[0] = 21_000_000.0
    true_range_tracking[1:] = (
        true_range_tracking[0]
        + np.cumsum(true_rate_tracking[:-1] * tracking_dt_s)
    )
    measured_rate_tracking = (
        true_rate_tracking
        + doppler_cfo_mps
        + rng.normal(0.0, doppler_sigma_mps, len(tracking_time_s))
    )
    filtered_rate_tracking = np.empty_like(measured_rate_tracking)
    filtered_rate_tracking[0] = measured_rate_tracking[0]
    for index in range(1, len(filtered_rate_tracking)):
        filtered_rate_tracking[index] = (
            firmware_doppler_alpha * measured_rate_tracking[index]
            + (1.0 - firmware_doppler_alpha) * filtered_rate_tracking[index - 1]
        )

    tracking_per_wls = int(round(wls_dt_s / tracking_dt_s))
    wls_indices = np.arange(0, len(tracking_time_s), tracking_per_wls)
    time_s = tracking_time_s[wls_indices]
    true_range = true_range_tracking[wls_indices]
    raw_code = true_range + rng.normal(0.0, code_sigma_m, len(time_s))

    scalar = np.empty_like(raw_code)
    doppler_smoothed = np.empty_like(raw_code)
    firmware_iir_doppler_smoothed = np.empty_like(raw_code)
    scalar[0] = raw_code[0]
    doppler_smoothed[0] = raw_code[0]
    firmware_iir_doppler_smoothed[0] = raw_code[0]
    for index in range(1, len(time_s)):
        scalar[index] = alpha * raw_code[index] + (1.0 - alpha) * scalar[index - 1]
        start = wls_indices[index - 1]
        stop = wls_indices[index]
        raw_range_delta = (
            float(np.sum(measured_rate_tracking[start:stop]))
            * tracking_dt_s
        )
        filtered_range_delta = (
            float(np.sum(filtered_rate_tracking[start:stop]))
            * tracking_dt_s
        )
        prediction = doppler_smoothed[index - 1] + raw_range_delta
        doppler_smoothed[index] = (
            prediction + alpha * (raw_code[index] - prediction)
        )
        filtered_prediction = (
            firmware_iir_doppler_smoothed[index - 1]
            + filtered_range_delta
        )
        firmware_iir_doppler_smoothed[index] = (
            filtered_prediction
            + alpha * (raw_code[index] - filtered_prediction)
        )

    warm = time_s >= 5.0
    dynamic = ((time_s >= 20.0) & (time_s < 25.0)) | (
        (time_s >= 45.0) & (time_s < 50.0)
    )
    rmse = lambda error, mask: float(np.sqrt(np.mean(np.square(error[mask]))))
    return SmoothingResult(
        raw_rmse_m=rmse(raw_code - true_range, warm),
        scalar_iir_rmse_m=rmse(scalar - true_range, warm),
        doppler_smoothed_rmse_m=rmse(doppler_smoothed - true_range, warm),
        firmware_iir_doppler_smoothed_rmse_m=rmse(
            firmware_iir_doppler_smoothed - true_range,
            warm,
        ),
        scalar_iir_peak_dynamic_error_m=float(
            np.max(np.abs(scalar[dynamic] - true_range[dynamic]))
        ),
        doppler_peak_dynamic_error_m=float(
            np.max(np.abs(doppler_smoothed[dynamic] - true_range[dynamic]))
        ),
        firmware_iir_doppler_peak_dynamic_error_m=float(
            np.max(
                np.abs(
                    firmware_iir_doppler_smoothed[dynamic]
                    - true_range[dynamic]
                )
            )
        ),
    )