"""Compare indexed and full-scan windows on frozen physical-pass fit inputs."""

import argparse
import json
import time
from pathlib import Path

import numpy as np

from k1link.artifacts import utc_now_iso
from k1link.missions.causal_replay import digest
from k1link.missions.reference_window import ReferenceWindowIndex, reference_window


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    report = json.loads((args.run / "report.json").read_text())
    artifacts = {args.run / p: sha for p, sha in report["artifacts"].items()}
    assert all(digest(p) == sha for p, sha in artifacts.items())
    code = {p: digest(p) for p in [Path(__file__), Path("src/k1link/missions/reference_window.py")]}
    reference = np.load(args.run / "reference.npy", allow_pickle=False)
    started = time.perf_counter()
    index = ReferenceWindowIndex(reference)
    preparation = time.perf_counter() - started
    measurements = []
    for step in sorted(args.run.glob("step-*")):
        source = json.loads((step / "source.json").read_text())
        if source["role"] != "fresh-validation":
            continue
        input_path = step / "registration-input.npz"
        if not input_path.exists():
            continue
        with np.load(input_path, allow_pickle=False) as frozen:
            sample = dict(
                points=frozen["query"], path=np.asarray(source["query_path"])
            )
            t0 = time.perf_counter()
            expected, old = reference_window(reference, sample, frozen["initial"])
            t1 = time.perf_counter()
            actual, new = reference_window(reference, sample, frozen["initial"], index=index)
            t2 = time.perf_counter()
            assert np.array_equal(expected, actual)
            assert np.array_equal(actual, frozen["reference"])
            assert new["target_sha256"] == old["target_sha256"]
            measurements.append(
                dict(
                    step=step.name,
                    full_scan_s=t1 - t0,
                    indexed_s=t2 - t1,
                    map_points=len(reference),
                    target_points=len(actual),
                    examined_points=new["examined_points"],
                    target_sha256=new["target_sha256"],
                )
            )
    assert measurements
    assert all(digest(p) == sha for p, sha in artifacts.items())
    assert all(digest(p) == sha for p, sha in code.items())
    result = dict(
        created_at_utc=utc_now_iso(),
        run=str(args.run),
        original_report_sha256=digest(args.run / "report.json"),
        implementation_sha256={str(p): sha for p, sha in code.items()},
        index_preparation_s=preparation,
        index_bytes=index.order.nbytes,
        spatial_cells=len(index.slices),
        measurements=measurements,
        all_exact=True,
        sources_unchanged=True,
        limitation="One bounded pass on real saved inputs, not kilometre qualification.",
    )
    with args.output.open("x") as stream:
        json.dump(result, stream, indent=2, allow_nan=False)
    print(
        json.dumps(
            dict(
                exact_windows=len(measurements),
                map_points=len(reference),
                index_preparation_s=preparation,
                median_full_s=float(np.median([m["full_scan_s"] for m in measurements])),
                median_indexed_s=float(np.median([m["indexed_s"] for m in measurements])),
                median_examined=float(np.median([m["examined_points"] for m in measurements])),
            )
        )
    )


if __name__ == "__main__":
    main()