#!/usr/bin/env python3
"""Re-run SignalSim E1B acquisition/tracking and diagnose lock transitions."""

from __future__ import annotations

import argparse
from dataclasses import asdict
import json
import math
from pathlib import Path
import sys
from typing import Any

import numpy as np
import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from galileo_rx.acquisition.acquisition_manager import AcquisitionManager
from galileo_rx.config import AcquisitionConfig, TrackingConfig
from galileo_rx.diagnostics.lock_analysis import (
    LockThresholds,
    analyze_tracking_lock,
    plot_tracking_lock_diagnostics,
    save_tracking_measurements,
    write_lock_report,
    write_transition_csv,
)
from galileo_rx.io.signalsim_reader import SignalSimReader, SignalSimSampleSpec
from galileo_rx.signal_profiles import galileo_e1_profile
from galileo_rx.signals.signalsim_codes import SignalSimCodeProvider
from galileo_rx.tracking.tracking_channel import TrackingChannel


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="configs/navigation_collection.yaml")
    parser.add_argument("--output-dir", default="reports/tracking_lock_baseline")
    parser.add_argument("--prn", type=int, action="append")
    args = parser.parse_args()
    config_path = Path(args.config).resolve()
    root = config_path.parent.parent
    cfg = yaml.safe_load(config_path.read_text(encoding="utf-8"))
    fs = float(cfg["rf"]["sampling_frequency_hz"])
    if_hz = float(cfg["rf"]["intermediate_frequency_hz"])
    source = (root / cfg["signalsim"]["path"]).resolve()
    samples_path = (root / cfg["signalsim"]["samples"]).resolve()
    samples = SignalSimReader(
        samples_path, SignalSimSampleSpec(cfg["signalsim"]["sample_format"]), memory_map=True
    ).read()
    profile = galileo_e1_profile()
    acquisition_cfg = AcquisitionConfig(**cfg["acquisition"], parallelism="doppler_vectorized")
    provider = SignalSimCodeProvider(source)
    manager = AcquisitionManager(profile, provider, fs, acquisition_cfg)
    acquisition_count = manager.engine.coherent_samples * acquisition_cfg.noncoherent_accumulations
    primary_samples = int(round(fs * profile.component("E1B").primary_code_period_s))
    output = Path(args.output_dir).resolve()
    output.mkdir(parents=True, exist_ok=True)
    report: dict[str, Any] = {
        "analysis": "unmodified raw lock detector baseline",
        "sample_file": str(samples_path),
        "sample_count": int(samples.size),
        "thresholds": asdict(LockThresholds()),
        "satellites": {},
    }
    all_rows: list[dict[str, Any]] = []
    for prn in (args.prn or [int(value) for value in cfg["satellites"]["prns"]]):
        acquisition, _ = manager.acquire_one(samples[:acquisition_count], "E1B", prn, if_hz)
        start_sample = int(round(
            acquisition.code_phase_samples
            + math.ceil((acquisition_count - acquisition.code_phase_samples) / primary_samples) * primary_samples
        ))
        tracking_values = cfg["tracking"]
        tracking_cfg = TrackingConfig(
            integration_ms=float(tracking_values["integration_ms"]),
            pull_in_duration_ms=float(tracking_values["pull_in_duration_ms"]),
            pll_noise_bandwidth_hz=float(tracking_values["pll_noise_bandwidth_hz"]),
            fll_noise_bandwidth_hz=float(tracking_values["fll_noise_bandwidth_hz"]),
            dll_noise_bandwidth_hz=float(tracking_values["dll_noise_bandwidth_hz"]),
            secondary_code_mode="none",
            loss_of_lock_epochs=40,
        )
        channel = TrackingChannel(profile, provider, fs, tracking_cfg, acquisition, start_sample)
        measurements = channel.process(samples[start_sample:])
        summary, rows, nav_transitions = analyze_tracking_lock(
            prn,
            measurements,
            carrier_center_hz=channel.carrier_center_hz,
            integration_ms=tracking_cfg.integration_ms,
        )
        summary["acquisition"] = asdict(acquisition)
        report["satellites"][prn] = summary
        all_rows.extend(rows)
        plot_tracking_lock_diagnostics(
            output / f"prn_{prn:02d}_lock_diagnostics.png",
            measurements,
            nav_transitions,
            carrier_center_hz=channel.carrier_center_hz,
        )
        save_tracking_measurements(output / f"prn_{prn:02d}_measurements.npz", measurements)
        print(json.dumps({"prn": prn, "false_transitions": len(rows), "post_pull_in": summary["post_pull_in"]}))
    collection_path = root / "reports" / "inav_collection_report.json"
    if collection_path.is_file():
        collection = json.loads(collection_path.read_text(encoding="utf-8"))
        report["navigation_result"] = {
            "pages": sum(item["pages"] for item in collection["satellites"].values()),
            "crc_passes": sum(item["crc_passes"] for item in collection["satellites"].values()),
            "crc_failures": sum(item["crc_failures"] for item in collection["satellites"].values()),
            "complete_ephemeris_prns": [
                int(prn) for prn, item in collection["satellites"].items() if item["complete_ephemeris"]
            ],
        }
    write_transition_csv(output / "lock_false_transitions.csv", all_rows)
    write_lock_report(output / "tracking_lock_baseline.json", report)
    return 0


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