#!/usr/bin/env python3
"""Run the C-mirroring WLS + timestamped PPS PI simulation."""

from __future__ import annotations

import argparse
import os
from pathlib import Path

os.environ.setdefault("MPLCONFIGDIR", "/tmp/pps_python_matplotlib")

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

try:
    from .pps_loop import PIConfig, ControllerTrace, simulate_controller
    from .pps_model import (
        C,
        DEFAULT_IONO,
        POSTPONE_TOW_S,
        RF_HZ,
        START_OFFSET_S,
        TRACKING_HZ,
        ReceiverEpochSimulator,
        build_clock_trajectory,
        c_tropo_delay_m,
        corrected_geometric_range,
        ecef_to_lla,
        azel_from_ecef,
        klobuchar_delay_m,
        lla_to_ecef,
        make_satellite_constellation,
    )
except ImportError:  # allow: python3 pps_python/simulate_pps.py
    from pps_loop import PIConfig, ControllerTrace, simulate_controller
    from pps_model import (
        C,
        DEFAULT_IONO,
        POSTPONE_TOW_S,
        RF_HZ,
        START_OFFSET_S,
        TRACKING_HZ,
        ReceiverEpochSimulator,
        build_clock_trajectory,
        c_tropo_delay_m,
        corrected_geometric_range,
        ecef_to_lla,
        azel_from_ecef,
        klobuchar_delay_m,
        lla_to_ecef,
        make_satellite_constellation,
    )


def finite_rms(values: np.ndarray) -> float:
    finite = values[np.isfinite(values)]
    return float(np.sqrt(np.mean(finite * finite))) if finite.size else np.nan


def finite_percentile_abs(values: np.ndarray, percentile: float) -> float:
    finite = values[np.isfinite(values)]
    return float(np.percentile(np.abs(finite), percentile)) if finite.size else np.nan


def tail(values: np.ndarray, fraction: float = 0.5) -> np.ndarray:
    start = int(values.size * (1.0 - fraction))
    return values[start:]


def save_epoch_csv(output_dir: Path, epochs, receiver_ecef: np.ndarray) -> None:
    rows = []
    for epoch in epochs:
        rows.append(
            [
                epoch.sequence,
                epoch.transmit_true_s,
                epoch.reference_arrival_true_s,
                epoch.snapshot_ready_true_s,
                epoch.apply_true_s,
                epoch.satellite_count,
                epoch.reference_svid,
                epoch.wls.clock_bias_m,
                epoch.wls.clock_bias_m / C,
                epoch.timing_measurement_error_s,
                np.linalg.norm(epoch.wls.position_ecef_m - receiver_ecef),
                epoch.compute_delay_s,
                epoch.apb_delay_s,
                epoch.wls.iterations,
                epoch.wls.final_clock_update_m,
            ]
        )
    header = (
        "sequence,tx_true_s,reference_true_s,snapshot_ready_true_s,apply_true_s,"
        "satellite_count,reference_svid,clock_bias_m,clock_bias_s,"
        "reference_error_s,position_error_m,compute_delay_s,apb_delay_s,"
        "wls_iterations,final_clock_update_m"
    )
    np.savetxt(
        output_dir / "epoch_results.csv",
        np.asarray(rows),
        delimiter=",",
        header=header,
        comments="",
        fmt="%.12g",
    )


def plot_wls(output_dir: Path, epochs, receiver_ecef: np.ndarray) -> None:
    time_s = np.asarray([epoch.transmit_true_s for epoch in epochs])
    bias_ms = np.asarray([epoch.raw_bias_phase_s for epoch in epochs]) * 1.0e3
    reference_error_ns = (
        np.asarray([epoch.timing_measurement_error_s for epoch in epochs]) * 1.0e9
    )
    position_error_m = np.asarray(
        [np.linalg.norm(epoch.wls.position_ecef_m - receiver_ecef) for epoch in epochs]
    )
    satellite_count = np.asarray([epoch.satellite_count for epoch in epochs])
    reference_svid = np.asarray([epoch.reference_svid for epoch in epochs])

    figure, axes = plt.subplots(4, 1, figsize=(12, 11), sharex=True)
    axes[0].plot(time_s, bias_ms, linewidth=1.0)
    axes[0].axhline(0.0, color="black", linewidth=0.6)
    axes[0].set_ylabel("WLS bias / c [ms]")
    axes[0].set_title(
        "Current C observation datum: raw WLS bias is not the PPS phase error"
    )
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(time_s, reference_error_ns, linewidth=0.8, color="tab:green")
    axes[1].axhline(0.0, color="black", linewidth=0.6)
    axes[1].set_ylabel("TOW+bias reference error [ns]")
    axes[1].grid(True, alpha=0.3)

    axes[2].semilogy(time_s, np.maximum(position_error_m, 1.0e-3), linewidth=0.8)
    axes[2].set_ylabel("Position error [m]")
    axes[2].grid(True, which="both", alpha=0.3)

    axes[3].step(time_s, satellite_count, where="post", label="satellite count")
    axes[3].step(time_s, reference_svid, where="post", label="minimum/reference SVID")
    axes[3].set_ylabel("count / SVID")
    axes[3].set_xlabel("True elapsed time [s]")
    axes[3].legend(loc="best")
    axes[3].grid(True, alpha=0.3)
    figure.tight_layout()
    figure.savefig(output_dir / "01_wls_time_model.png", dpi=160)
    plt.close(figure)


def plot_latency(output_dir: Path, epochs) -> None:
    time_s = np.asarray([epoch.transmit_true_s for epoch in epochs])
    propagation_span_ms = np.asarray(
        [epoch.snapshot_ready_true_s - epoch.reference_arrival_true_s for epoch in epochs]
    ) * 1.0e3
    compute_ms = np.asarray([epoch.compute_delay_s for epoch in epochs]) * 1.0e3
    apb_ms = np.asarray([epoch.apb_delay_s for epoch in epochs]) * 1.0e3
    total_ms = np.asarray(
        [epoch.apply_true_s - epoch.reference_arrival_true_s for epoch in epochs]
    ) * 1.0e3

    figure, axes = plt.subplots(2, 1, figsize=(12, 7), sharex=True)
    axes[0].plot(time_s, propagation_span_ms, label="satellite arrival span")
    axes[0].plot(time_s, compute_ms, label="WLS compute")
    axes[0].plot(time_s, apb_ms, label="APB")
    axes[0].set_ylabel("Delay component [ms]")
    axes[0].legend(loc="best")
    axes[0].grid(True, alpha=0.3)
    axes[1].plot(time_s, total_ms, color="tab:red")
    axes[1].set_ylabel("Capture to commit [ms]")
    axes[1].set_xlabel("True elapsed time [s]")
    axes[1].grid(True, alpha=0.3)
    figure.suptitle("Compute/APB latency belongs to an old, timestamped epoch")
    figure.tight_layout()
    figure.savefig(output_dir / "02_latency_timeline.png", dpi=160)
    plt.close(figure)


def plot_pi(
    output_dir: Path,
    traces: list[ControllerTrace],
    tracking_clock,
) -> None:
    timestamped, free_running, untagged = traces
    figure, axes = plt.subplots(4, 1, figsize=(12, 12), sharex=True)
    axes[0].plot(
        timestamped.time_s,
        timestamped.phase_error_s * 1.0e9,
        label="timestamped PI",
    )
    axes[0].set_ylabel("PI phase [ns]")
    axes[0].legend(loc="best")
    axes[0].grid(True, alpha=0.3)

    axes[1].plot(
        free_running.time_s,
        free_running.phase_error_s * 1.0e6,
        color="tab:orange",
        label="timestamped, PI disabled",
    )
    axes[1].set_ylabel("Free-run phase [us]")
    axes[1].legend(loc="best")
    axes[1].grid(True, alpha=0.3)

    axes[2].plot(
        untagged.time_s,
        untagged.phase_error_s * 1.0e3,
        color="tab:red",
        label="old reference treated as APB-now",
    )
    axes[2].set_ylabel("Untagged phase [ms]")
    axes[2].legend(loc="best")
    axes[2].grid(True, alpha=0.3)

    axes[3].plot(
        timestamped.time_s,
        timestamped.frequency_command * 1.0e6,
        label="PI frequency command",
    )
    axes[3].plot(
        timestamped.time_s,
        tracking_clock.ppm_at(timestamped.time_s),
        label="raw 125 MHz oscillator",
        alpha=0.8,
    )
    axes[3].set_ylabel("Frequency [ppm]")
    axes[3].set_xlabel("True elapsed time [s]")
    axes[3].legend(loc="best")
    axes[3].grid(True, alpha=0.3)
    figure.suptitle("125 MHz disciplined time: timestamping versus delayed APB-now")
    figure.tight_layout()
    figure.savefig(output_dir / "03_pi_loop_response.png", dpi=160)
    plt.close(figure)


def plot_pps(
    output_dir: Path,
    traces: list[ControllerTrace],
    gps_tow_start_s: float,
) -> None:
    timestamped, free_running, untagged = traces
    ts_time, ts_error = timestamped.pps_samples(gps_tow_start_s)
    fr_time, fr_error = free_running.pps_samples(gps_tow_start_s)
    un_time, un_error = untagged.pps_samples(gps_tow_start_s)

    figure, axes = plt.subplots(3, 1, figsize=(12, 9), sharex=True)
    axes[0].plot(ts_time, ts_error * 1.0e9, marker=".", label="timestamped PI")
    axes[0].set_ylabel("Predicted PPS error [ns]")
    axes[0].legend(loc="best")
    axes[0].grid(True, alpha=0.3)
    axes[1].plot(fr_time, fr_error * 1.0e6, marker=".", color="tab:orange")
    axes[1].set_ylabel("Free-run PPS error [us]")
    axes[1].grid(True, alpha=0.3)
    axes[2].plot(un_time, un_error * 1.0e3, marker=".", color="tab:red")
    axes[2].set_ylabel("Untagged PPS error [ms]")
    axes[2].set_xlabel("True elapsed time [s]")
    axes[2].grid(True, alpha=0.3)
    figure.suptitle("PPS edge error, including 8 ns output quantization")
    figure.tight_layout()
    figure.savefig(output_dir / "04_pps_edge_error.png", dpi=160)
    plt.close(figure)


def simulate_startup_tow(
    receiver_ecef: np.ndarray,
    satellites,
    epochs,
    rng: np.random.Generator,
    *,
    trial_count: int = 4000,
) -> dict[str, np.ndarray]:
    latitude, longitude, height = ecef_to_lla(receiver_ecef)
    flight_times = []
    for satellite in satellites:
        geometric, rotated = corrected_geometric_range(receiver_ecef, satellite.ecef_m)
        azimuth, elevation = azel_from_ecef(receiver_ecef, rotated)
        ionosphere = klobuchar_delay_m(
            DEFAULT_IONO,
            345_600.0,
            latitude,
            longitude,
            azimuth,
            elevation,
        )
        troposphere = c_tropo_delay_m(
            np.sin(np.deg2rad(max(elevation, 0.0))), height / 1000.0
        )
        flight_times.append(
            (geometric + ionosphere + troposphere) / C
            - satellite.clock_correction(0.0)
        )
    flight_times = np.asarray(flight_times)

    # Acquisition/bit-lock availability is intentionally independent per
    # channel. A full subframe can only be decoded after the next complete
    # six-second interval, so the first HOW is not necessarily the nearest SV.
    acquisition_ready = rng.uniform(0.0, 18.0, size=(trial_count, len(satellites)))
    complete_subframe_tx = (np.floor(acquisition_ready / 6.0) + 2.0) * 6.0
    # The current hardware validates/delivers the completed subframe only
    # on the ninth following bit because validation reads the old shift
    # register. At 50 bit/s
    # this is a deterministic 180 ms readout delay (POSTPONE_TOW), distinct
    # from satellite signal flight time.
    how_interrupt_rx = (
        complete_subframe_tx + POSTPONE_TOW_S + flight_times[None, :]
    )
    first_index = np.argmin(how_interrupt_rx, axis=1)
    first_flight = flight_times[first_index]

    timing_noise = np.asarray(
        [epoch.timing_measurement_error_s for epoch in epochs]
    )
    wls_error = rng.choice(timing_noise, size=trial_count, replace=True)
    latency = np.asarray(
        [epoch.apply_true_s - epoch.reference_arrival_true_s for epoch in epochs]
    )
    sampled_latency = rng.choice(latency, size=trial_count, replace=True)
    return {
        "first_svid": first_index + 1,
        "first_flight_s": first_flight,
        # Parsed HOW labels transmit subframe time, not receive-now.
        "how_as_receive_error_s": -(POSTPONE_TOW_S + first_flight),
        "how_plus_pipeline_startoffset_error_s": START_OFFSET_S - first_flight,
        "wls_aligned_error_s": wls_error,
        # Old receive reference written as if it were valid at APB-now.
        "apb_now_error_s": wls_error - sampled_latency,
    }


def plot_startup_tow(output_dir: Path, startup: dict[str, np.ndarray]) -> None:
    first_svid = startup["first_svid"]
    unique, counts = np.unique(first_svid, return_counts=True)
    probability = counts / counts.sum()
    figure, axes = plt.subplots(2, 2, figsize=(12, 8))
    axes[0, 0].bar(unique, probability)
    axes[0, 0].set_xlabel("First valid HOW SVID")
    axes[0, 0].set_ylabel("Probability")
    axes[0, 0].set_title("Acquisition/subframe phase can dominate proximity")
    axes[0, 0].grid(True, axis="y", alpha=0.3)

    axes[0, 1].hist(
        startup["how_as_receive_error_s"] * 1.0e3,
        bins=35,
        alpha=0.8,
    )
    axes[0, 1].set_xlabel("Phase error [ms]")
    axes[0, 1].set_title("Wrong: parsed HOW used as receive-now")
    axes[0, 1].grid(True, alpha=0.3)

    axes[1, 0].hist(
        startup["how_plus_pipeline_startoffset_error_s"] * 1.0e3,
        bins=35,
        color="tab:orange",
        alpha=0.8,
    )
    axes[1, 0].set_xlabel("Phase error [ms]")
    axes[1, 0].set_title("Provisional: HOW + 180 ms + 68.802 ms")
    axes[1, 0].grid(True, alpha=0.3)

    axes[1, 1].hist(
        startup["wls_aligned_error_s"] * 1.0e9,
        bins=45,
        label="timestamped WLS alignment",
        alpha=0.8,
    )
    axes[1, 1].set_xlabel("Phase error [ns]")
    axes[1, 1].set_title("Correct: TOW + flight time paired with 125 MHz timestamp")
    axes[1, 1].grid(True, alpha=0.3)
    figure.suptitle("Cold-start TOW ownership and PPS coarse alignment")
    figure.tight_layout()
    figure.savefig(output_dir / "05_startup_tow_alignment.png", dpi=160)
    plt.close(figure)


def build_summary(
    epochs,
    receiver_ecef: np.ndarray,
    traces: list[ControllerTrace],
    gps_tow_start_s: float,
    startup: dict[str, np.ndarray],
) -> str:
    times = np.asarray([epoch.transmit_true_s for epoch in epochs])
    position_error = np.asarray(
        [np.linalg.norm(epoch.wls.position_ecef_m - receiver_ecef) for epoch in epochs]
    )
    timing_error = np.asarray([epoch.timing_measurement_error_s for epoch in epochs])
    raw_bias = np.asarray([epoch.raw_bias_phase_s for epoch in epochs])
    latency = np.asarray(
        [epoch.apply_true_s - epoch.reference_arrival_true_s for epoch in epochs]
    )

    lines = [
        "C-mirroring WLS + PPS timing simulation",
        "========================================",
        f"epochs: {len(epochs)}",
        f"raw WLS bias/c range: {np.min(raw_bias)*1e3:.6f} .. {np.max(raw_bias)*1e3:.6f} ms",
        f"derived GPS reference RMS error: {finite_rms(timing_error)*1e9:.3f} ns",
        f"capture-to-commit delay mean/p99: {np.mean(latency)*1e3:.3f} / {np.percentile(latency, 99)*1e3:.3f} ms",
        f"first HOW used as receive-now, median phase error: {np.median(startup['how_as_receive_error_s'])*1e3:.3f} ms",
        f"first HOW + 180 ms + 68.802 ms, median phase error: {np.median(startup['how_plus_pipeline_startoffset_error_s'])*1e3:.3f} ms",
    ]
    for start, stop, label in [
        (0.0, 20.0, "4 satellites"),
        (20.0, 50.0, "5 satellites"),
        (50.0, np.inf, "7 satellites / drop-reacquisition"),
    ]:
        selected = (times >= start) & (times < stop)
        if np.any(selected):
            lines.append(
                f"{label}: position RMS={finite_rms(position_error[selected]):.3f} m, "
                f"timing RMS={finite_rms(timing_error[selected])*1e9:.3f} ns"
            )

    for trace in traces:
        phase_tail = tail(trace.phase_error_s)
        pps_time, pps_error = trace.pps_samples(gps_tow_start_s)
        del pps_time
        lines.append(
            f"{trace.name}: tail phase RMS={finite_rms(phase_tail)*1e9:.3f} ns, "
            f"|phase| p95={finite_percentile_abs(phase_tail, 95)*1e9:.3f} ns, "
            f"PPS RMS={finite_rms(pps_error)*1e9:.3f} ns, "
            f"rejected stale commits={trace.rejected_updates}"
        )
    lines.extend(
        [
            "",
            "Interpretation:",
            "- bias/c is millisecond-scale because it contains the 68.802 ms synthetic datum and the minimum satellite path.",
            "- epoch-aligned TOW (first delivery: HOW + 180 ms) + STARTOFFSET - bias/c reconstructs GPS time at the paired timestamp.",
            "- treating that old reference as valid at APB arrival converts software latency into PPS phase error.",
            "- timestamping removes that phase misassociation; control application delay remains only as loop delay.",
            "- baremetal applies the fixed 180 ms readout correction once; the PPS peripheral must not add it again.",
        ]
    )
    return "\n".join(lines) + "\n"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--duration", type=float, default=180.0)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(__file__).resolve().parent / "output",
    )
    parser.add_argument("--no-plots", action="store_true")
    parser.add_argument("--loop-bandwidth-hz", type=float, default=0.035)
    parser.add_argument("--max-slew-ppm", type=float, default=20.0)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.duration < 5.0:
        raise SystemExit("duration must be at least 5 seconds")
    args.output_dir.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(args.seed)
    clock_duration = args.duration + 2.0

    shared = build_clock_trajectory(
        clock_duration,
        1.0,
        rng,
        base_ppm=1.45,
        temperature_ppm=0.45,
        temperature_period_s=140.0,
        random_walk_ppm_sqrt_s=0.002,
    )
    rf_clock = build_clock_trajectory(
        clock_duration,
        RF_HZ,
        rng,
        base_ppm=0.10,
        temperature_ppm=0.05,
        temperature_period_s=95.0,
        random_walk_ppm_sqrt_s=0.001,
        shared_fractional=shared.fractional_frequency,
        # Force an early 32-bit wrap to exercise the firmware unwrap logic.
        initial_ticks=float((1 << 32) - 100_000),
    )
    tracking_clock = build_clock_trajectory(
        clock_duration,
        TRACKING_HZ,
        rng,
        base_ppm=0.20,
        temperature_ppm=0.12,
        temperature_period_s=180.0,
        random_walk_ppm_sqrt_s=0.001,
        shared_fractional=shared.fractional_frequency,
    )

    receiver_ecef = lla_to_ecef(39.9208, 32.8541, 890.0)
    satellites = make_satellite_constellation(receiver_ecef)
    gps_tow_start_s = 345_600.0
    receiver = ReceiverEpochSimulator(
        receiver_ecef,
        satellites,
        rf_clock,
        tracking_clock,
        rng,
        gps_tow_start_s=gps_tow_start_s,
    )
    epochs = receiver.simulate(args.duration)
    if not epochs:
        raise SystemExit("no valid WLS epochs were generated")

    pi_config = PIConfig(
        natural_frequency_hz=args.loop_bandwidth_hz,
        damping=0.9,
        max_slew_ppm=args.max_slew_ppm,
    )
    disabled_config = PIConfig(
        natural_frequency_hz=0.0,
        damping=0.9,
        max_slew_ppm=args.max_slew_ppm,
    )
    timestamped = simulate_controller(
        "timestamped PI",
        "timestamped",
        epochs,
        tracking_clock,
        gps_tow_start_s,
        args.duration,
        pi_config,
    )
    free_running = simulate_controller(
        "timestamped, PI disabled",
        "timestamped",
        epochs,
        tracking_clock,
        gps_tow_start_s,
        args.duration,
        disabled_config,
    )
    untagged = simulate_controller(
        "untagged APB-now PI",
        "untagged",
        epochs,
        tracking_clock,
        gps_tow_start_s,
        args.duration,
        pi_config,
    )
    traces = [timestamped, free_running, untagged]
    startup = simulate_startup_tow(receiver_ecef, satellites, epochs, rng)

    save_epoch_csv(args.output_dir, epochs, receiver_ecef)
    summary = build_summary(
        epochs, receiver_ecef, traces, gps_tow_start_s, startup
    )
    (args.output_dir / "summary.txt").write_text(summary, encoding="utf-8")
    if not args.no_plots:
        plot_wls(args.output_dir, epochs, receiver_ecef)
        plot_latency(args.output_dir, epochs)
        plot_pi(args.output_dir, traces, tracking_clock)
        plot_pps(args.output_dir, traces, gps_tow_start_s)
        plot_startup_tow(args.output_dir, startup)
    print(summary, end="")
    print(f"output: {args.output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())