"""Additional checks for an existing correction, never a re-fit or threshold change."""

import argparse
import json
from pathlib import Path

import numpy as np
from reconstruct_recorded_ring import PROFILE, digest, register, voxel, write_json
from scipy.spatial import cKDTree

from k1link.missions.registration import transform
from k1link.reconstruction.closure import review_acceptance
from k1link.reconstruction.smooth_correction import CorrectionField


def stats(values):
    return dict(
        zip(
            ["min", "median", "p95", "max"],
            np.quantile(values, [0, 0.5, 0.95, 1]).tolist(),
            strict=True,
        )
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("result", type=Path)
    parser.add_argument("--output-name", default="review.json")
    parser.add_argument("--include-group", type=int, action="append", default=[])
    args = parser.parse_args()
    root = args.result
    summary = json.loads((root / "summary.json").read_text())
    fit = json.loads((root / "correction.json").read_text())
    with np.load(Path(summary["source_cache"]) / "index.npz") as d:
        data = dict(d)
    p, f, s = data["poses"], data["frames"], data["frame_distance"]
    pts, ids, held = data["sample_points"], data["sample_frame"], data["heldout"]
    field = CorrectionField(fit["knots_m"], fit["parameters"], fit.get("origin_m"))
    length = float(data["distance"][-1])
    baseline = CorrectionField([0, length], np.zeros((2, 6)))
    runs = json.loads((root / "validation.json").read_text())
    failures = sorted(
        set(args.include_group)
        | {r["group"] for run in runs for r in run["windows"] if r["fit"]["status"] != "candidate"}
    )
    offsets = np.searchsorted(ids, np.arange(len(f) + 1))
    results = []
    for label, correction, run in zip(
        ["original", "corrected"], [baseline, field], runs, strict=True
    ):
        transformed = np.empty_like(pts)
        for i, distance in enumerate(s):
            start, end = offsets[i : i + 2]
            transformed[start:end] = correction.points(pts[start:end], distance)
        target = voxel(transformed[~held[ids]])
        tree = cKDTree(target)
        # Cross-visit check: last held-out frames against ONLY first training window.
        first = (f[:, 0] <= f[0, 0] + PROFILE["seam_reference_s"]) & ~held
        last = (f[:, 0] >= f[-1, 0] - 15) & held
        # Identical source point membership before/after, selected before correction.
        near = np.linalg.norm(pts - p[0, 1:4], axis=1) <= 25
        first_pts = transformed[first[ids] & near]
        last_pts = transformed[last[ids] & near]
        distances, _ = cKDTree(voxel(first_pts)).query(last_pts, workers=1)
        seam = dict(
            query_frames=int(last.sum()),
            points=len(last_pts),
            overlap_05m=float(np.mean(distances <= 0.5)),
            all_point_distances_m=stats(distances),
            inlier_rmse_m=float(np.sqrt(np.mean(distances[distances <= 0.5] ** 2))),
            refit=False,
        )
        focused = []
        for group in failures:
            # Same prior as the original two-second query; then only fresh held-out halves.
            full = next(r for r in run["windows"] if r["group"] == group)
            prior = np.asarray(full["fit"]["initial_T_reference_query"])
            t0 = f[0, 0] + group * PROFILE["holdout_period_s"] + PROFILE["holdout_start_s"]
            for half in [0, 1]:
                mask = held & (f[:, 0] >= t0 + half) & (f[:, 0] < t0 + half + 1)
                position = np.array(
                    [np.interp(float(np.mean(f[mask, 0])), p[:, 0], p[:, j]) for j in range(1, 4)]
                )
                query = pts[mask[ids]]
                query = query[np.linalg.norm(query - position, axis=1) <= 40]
                estimated = transform(position[None], prior)[0]
                reference = target[tree.query_ball_point(estimated, 45)]
                match = register(reference, query, prior)
                focused.append(dict(group=group, half=half, frames=int(mask.sum()), fit=match))
                if match["status"] == "candidate":
                    prior = np.asarray(match["T_reference_query"])
        results.append(
            dict(
                label=label,
                seam_holdout=seam,
                focused=focused,
                overlap=stats([r["fit"]["overlap"] for r in run["windows"]]),
                inlier_rmse_m=stats([r["fit"]["inlier_rmse_m"] for r in run["windows"]]),
                model_consistency_m=stats([r["model_consistency_m"] for r in run["windows"]]),
            )
        )
    acceptance = review_acceptance(results, runs)
    write_json(
        root / args.output_name,
        dict(
            frozen_correction=True,
            unchanged_registration_policy=True,
            selected_failure_groups=failures,
            results=results,
            acceptance=acceptance,
        ),
    )
    sources = [
        Path(__file__),
        Path(__file__).with_name("reconstruct_recorded_ring.py"),
        Path(__file__).parents[1] / "src/k1link/reconstruction/smooth_correction.py",
        Path(__file__).parents[1] / "src/k1link/reconstruction/closure.py",
        Path(__file__).parents[1] / "src/k1link/missions/registration.py",
    ]
    evidence = [
        root / name
        for name in [
            "summary.json",
            "correction.json",
            "validation.json",
            "registrations.json",
            "surface-links.json",
            "corrected-points.f32",
            "corrected-trajectory.npz",
            args.output_name,
        ]
    ]
    if summary["schema_version"] == "missioncore.recorded-ring-experiment/v2":
        evidence.append(root / "closure-search.json")
    write_json(
        root / (args.output_name + ".seal.json"),
        {str(path.resolve()): digest(path) for path in sources + evidence},
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()