#!/usr/bin/env python3
"""Audit BESTNAVA/Doppler factor yield for every unified RTK--IMU session."""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import replace
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from rtk_imu.rtk_imu_engineering import (
    _all_hpr,
    _height_reference,
    _nodes,
    _position_valid,
)
from rtk_imu.rtk_imu_multisource import _f, _nearest_index, _truth, load_unified_sessions


def _finite_doppler(row: dict[str, str]) -> bool:
    value = np.asarray([
        _f(row, "velocity_east_m_s"),
        _f(row, "velocity_north_m_s"),
        _f(row, "vertical_speed_m_s"),
    ])
    return bool(_truth(row, "doppler_velocity_valid") and np.all(np.isfinite(value)))


def _source_only_nodes(session, source: str, period_s: float):
    rows = {source: session.rtk_by_type.get(source, [])}
    if "GNHPR" in session.rtk_by_type:
        rows["GNHPR"] = session.rtk_by_type["GNHPR"]
    return _nodes(replace(session, rtk_by_type=rows), _height_reference([session]), period_s)


def _audit_session(session, period_s: float) -> dict[str, object]:
    best = session.rtk_by_type.get("BESTNAVA", [])
    hpr = _all_hpr(session)
    counts = {
        "raw_bestnava": len(best),
        "checksum_valid": 0,
        "fixed": 0,
        "finite_position": 0,
        "finite_doppler": 0,
        "hpr_near": 0,
        "hpr_q4": 0,
        "imu_near": 0,
        "raw_bestnava_candidate": 0,
    }
    for row in best:
        if not _truth(row, "checksum_valid"):
            continue
        counts["checksum_valid"] += 1
        if not _truth(row, "position_fixed"):
            continue
        counts["fixed"] += 1
        if not _position_valid(row, "BESTNAVA"):
            continue
        counts["finite_position"] += 1
        if not _finite_doppler(row):
            continue
        counts["finite_doppler"] += 1
        t_s = _f(row, "t_device_s")
        hpr_index = _nearest_index(hpr.t_s, t_s, 0.12)
        if hpr_index is None:
            continue
        counts["hpr_near"] += 1
        if not hpr.valid[hpr_index]:
            continue
        counts["hpr_q4"] += 1
        if _nearest_index(session.imu.t_s, t_s, 0.03) is None:
            continue
        counts["imu_near"] += 1
        counts["raw_bestnava_candidate"] += 1

    reference = _height_reference([session])
    combined_nodes = _nodes(session, reference, period_s) if reference is not None else []
    best_nodes = _source_only_nodes(session, "BESTNAVA", period_s) if reference is not None else []
    combined_best = [node for node in combined_nodes if node.source == "BESTNAVA"]
    best_velocity_nodes = [node for node in combined_best if node.velocity_enu_m_s is not None]
    source_best_velocity = [node for node in best_nodes if node.velocity_enu_m_s is not None]
    run_nodes: dict[int, list] = {}
    for node in combined_nodes:
        run_nodes.setdefault(node.continuity_id, []).append(node)
    qualifying_runs = [
        run for run in run_nodes.values()
        if len(run) >= 6 and run[-1].t_s - run[0].t_s >= 5.0
    ]
    qualifying_best = [
        node for run in qualifying_runs for node in run if node.source == "BESTNAVA"
    ]
    qualifying_doppler = [node for node in qualifying_best if node.velocity_enu_m_s is not None]
    return {
        "session_id": session.session_id,
        "batch_id": session.batch_id,
        "counts": counts,
        "combined_selected_nodes": len(combined_nodes),
        "combined_selected_bestnava": len(combined_best),
        "combined_selected_doppler": len(best_velocity_nodes),
        "best_only_selected_nodes": len(best_nodes),
        "best_only_selected_doppler": len(source_best_velocity),
        "best_suppressed_by_mixed_selection": max(0, len(best_nodes) - len(combined_best)),
        "doppler_suppressed_by_mixed_selection": max(0, len(source_best_velocity) - len(best_velocity_nodes)),
        "continuous_run_count": len(run_nodes),
        "qualified_run_count": len(qualifying_runs),
        "qualified_bestnava_factors": len(qualifying_best),
        "qualified_doppler_factors": len(qualifying_doppler),
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--sample-period-s", type=float, default=1.0)
    args = parser.parse_args(argv)
    sessions = load_unified_sessions(args.manifest)
    audit = [_audit_session(session, args.sample_period_s) for session in sessions]
    totals: dict[str, int] = {}
    for item in audit:
        for key, value in item["counts"].items():
            totals[key] = totals.get(key, 0) + int(value)
        for key in (
            "combined_selected_nodes", "combined_selected_bestnava", "combined_selected_doppler",
            "best_only_selected_nodes", "best_only_selected_doppler",
            "best_suppressed_by_mixed_selection", "doppler_suppressed_by_mixed_selection",
            "continuous_run_count", "qualified_run_count", "qualified_bestnava_factors",
            "qualified_doppler_factors",
        ):
            totals[key] = totals.get(key, 0) + int(item[key])
    payload = {
        "sample_period_s": args.sample_period_s,
        "session_count": len(audit),
        "totals": totals,
        "sessions": audit,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    print(json.dumps({"session_count": len(audit), "totals": totals}, ensure_ascii=False, indent=2))
    return 0


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