"""Two bounded negative checks using an already captured causal snapshot."""

import argparse
import json
from pathlib import Path

import numpy as np

from k1link.device_plugins.xgrids_k1.localization_source import extract_submap
from k1link.missions.causal_replay import digest
from k1link.missions.registration import path_hint
from k1link.missions.registration_worker import run_registration


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--replay", type=Path, required=True)
    p.add_argument("--reference-raw", type=Path, required=True)
    p.add_argument("--reference-planning", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    args = p.parse_args()
    base = json.loads((args.replay / "report.json").read_text())
    planning = json.loads(args.reference_planning.read_text())
    # Fixed third causal snapshot, never the final offline B fit.
    snapshot = args.replay / "step-003/registration-input.npz"
    query_path_file = snapshot.with_name("query-path.npy")
    files = {
        snapshot: base["artifacts"][str(snapshot.relative_to(args.replay))],
        query_path_file: base["artifacts"][str(query_path_file.relative_to(args.replay))],
        args.reference_raw: planning["source_digests"]["raw-transport-primary"],
        args.reference_raw.with_name("mqtt.metadata.jsonl"): planning["source_digests"][
            "raw-transport-index"
        ],
    }
    for file, expected in files.items():
        if digest(file) != expected:
            raise ValueError("Control input digest mismatch.")
    args.output.mkdir(parents=True, exist_ok=False)
    with np.load(snapshot, allow_pickle=False) as data:
        reference, query, hint = data["reference"], data["query"], data["initial"]
    path = np.load(query_path_file, allow_pickle=False)
    distant = hint.copy()
    distant[:3, 3] += 1000
    far_dir = args.output / "far-seed"
    far_dir.mkdir()
    far = run_registration(far_dir, reference, query, distant)
    # This disjoint A interval was specified before executing the controls.
    poses = planning["poses"]
    start = next(i for i, x in enumerate(poses) if x["distance_m"] >= 130)
    end = next(i for i, x in enumerate(poses) if x["distance_m"] >= 155)
    wrong, meta = extract_submap(args.reference_raw, planning, start, end)
    wrong_path = np.array([x["position"] for x in poses[start : end + 1]])
    wrong_dir = args.output / "wrong-region"
    wrong_dir.mkdir()
    other = run_registration(wrong_dir, wrong, query, path_hint(wrong_path, path))

    def clean(value):
        return {k: v for k, v in value.items() if k != "matched_query_indices"}

    report = dict(
        schema_version="missioncore.causal-replay-controls/v1",
        source_step="step-003",
        reference_interval_m=[130, 155],
        reference_interval_indices=[start, end],
        reference_extraction=meta,
        results={"far-seed": clean(far), "wrong-region": clean(other)},
        input_digests={str(k): v for k, v in files.items()},
        source_integrity_verified=all(digest(k) == v for k, v in files.items()),
        vehicle_control=False,
        localization_confirmed=False,
    )
    report["artifacts"] = {
        str(x.relative_to(args.output)): digest(x) for x in args.output.rglob("*") if x.is_file()
    }
    (args.output / "report.json").write_text(json.dumps(report, allow_nan=False))
    print(
        json.dumps(
            {
                k: {
                    j: v.get(j)
                    for j in [
                        "status",
                        "reasons",
                        "overlap",
                        "inlier_rmse_m",
                        "registration_seconds",
                    ]
                }
                for k, v in report["results"].items()
            }
        ),
        flush=True,
    )


if __name__ == "__main__":
    main()