"""Offline first-loop experiment on an immutable K1 recording, never an API action.

Run with the project's map-correction extra. All outputs are private derivatives;
the caller supplies a fresh output directory, source session identity and digest.
No raw overwrites, scene publication, hardware commands or planner mutations.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import time
from dataclasses import asdict
from datetime import UTC, datetime
from importlib.metadata import version
from pathlib import Path

import numpy as np
from scipy.spatial import cKDTree
from scipy.spatial.transform import Rotation

from k1link.device_plugins.xgrids_k1.protocol.streams import decode_lio_pcl, decode_lio_pose
from k1link.device_plugins.xgrids_k1.viewer.replay import iter_replay_messages
from k1link.missions.registration import POLICY, PreparedReference, transform
from k1link.reconstruction.closure import ClosurePolicy, ClosureUnavailable, acquire_closure
from k1link.reconstruction.smooth_correction import (
    CorrectionField,
    CorrectionPolicy,
    SurfaceLink,
    fit_correction,
)

PROFILE = dict(
    version="recorded-ring-experiment/v2",
    sample_stride=8,
    holdout_period_s=10.0,
    holdout_start_s=4.0,
    holdout_duration_s=2.0,
    seam_reference_s=20.0,
    seam_query_s=5.0,
    seam_radius_m=25.0,
    local_validation_radius_m=40.0,
    neighbor_radius_m=30.0,
    voxel_m=0.25,
    seam_translation_weight_m=0.03,
    seam_rotation_weight_deg=0.05,
    neighbor_translation_weight_m=0.2,
    neighbor_rotation_weight_deg=0.5,
)

# A separate first-fit policy, never a mutation of the live tracking policy.
# Quality/shape/information/correspondence gates remain identical.
CLOSURE_REGISTRATION_POLICY = {
    **POLICY,
    "version": "offline-closure-gicp/v1",
    "maximum_correction_m": 25.0,
    "maximum_correction_deg": 180.0,
}
DECODE_KEYS = ("sample_stride", "holdout_period_s", "holdout_start_s", "holdout_duration_s")


def compatible_cache_profile(profile):
    # Earlier sealed caches include fitting settings, though decoding never uses
    # them. Reuse is safe only when every actual decode/split setting agrees.
    return all(profile.get(key) == PROFILE[key] for key in DECODE_KEYS)


def digest(path):
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def write_json(path, value):
    with path.open("x") as stream:
        json.dump(value, stream, indent=2, allow_nan=False)


def extract(raw, output, expected):
    if digest(raw) != expected:
        raise ValueError("Source digest mismatch before decoding.")
    output.mkdir()  # Exclusive new directory, never replace an earlier experiment.
    started = time.monotonic()
    poses, frames, samples, sample_ids = [], [], [], []
    first = None
    total = 0
    sequences = {"lio_pose": [], "lio_pcl": []}
    with (
        (output / "source-points.f32").open("xb") as points_file,
        (output / "source-intensity.u8").open("xb") as intensity_file,
    ):
        for message in iter_replay_messages(raw):
            if message.received_monotonic_ns is None:
                raise ValueError("Source has no monotonic receipt timestamp.")
            clock = message.received_monotonic_ns / 1e9
            if first is None:
                first = clock
            t = clock - first
            if message.topic.endswith("/lio_pose"):
                pose = decode_lio_pose(message.payload)
                sequences["lio_pose"].append(pose.header.seq)
                poses.append(
                    [
                        t,
                        *pose.position_xyz,
                        *pose.orientation_xyzw,
                        pose.pose_stamp,
                        pose.header.seq,
                    ]
                )
            elif message.topic.endswith("/lio_pcl"):
                cloud = decode_lio_pcl(message.payload)
                sequences["lio_pcl"].append(cloud.header.seq)
                data = np.asarray(cloud.points, dtype=np.int64).reshape(-1, 4)
                xyz = (data[:, :3] / cloud.header.scaler).astype("<f4")
                if not np.isfinite(xyz).all():
                    raise ValueError("Non-finite source geometry.")
                xyz.tofile(points_file)
                (data[:, 3] & 255).astype("u1").tofile(intensity_file)
                sampled = xyz[:: PROFILE["sample_stride"]]
                samples.append(sampled)
                sample_ids.append(np.full(len(sampled), len(frames), dtype=np.int32))
                frames.append([t, cloud.header.seq, total, len(xyz)])
                total += len(xyz)
                if len(frames) % 500 == 0:
                    print(f"decode {len(frames)} frames, {total} points", flush=True)
    p, f = np.asarray(poses), np.asarray(frames)
    if min(len(p), len(f)) < 2 or (np.diff(p[:, 0]) <= 0).any():
        raise ValueError("Missing or unordered source trajectory.")
    if (np.diff(f[:, 0]) < 0).any():
        raise ValueError("Cloud receipt clock moved backwards.")
    if any((np.diff(seq) != 1).any() for seq in sequences.values()):
        raise ValueError("Source sequence gaps or resets require a separate review.")
    distance = np.r_[0.0, np.cumsum(np.linalg.norm(np.diff(p[:, 1:4], axis=0), axis=1))]
    frame_distance = np.interp(f[:, 0], p[:, 0], distance)
    elapsed = f[:, 0] - f[0, 0]
    phase = elapsed % PROFILE["holdout_period_s"]
    held = (phase >= PROFILE["holdout_start_s"]) & (
        phase < PROFILE["holdout_start_s"] + PROFILE["holdout_duration_s"]
    )
    nearest = np.clip(np.searchsorted(p[:, 0], f[:, 0]), 1, len(p) - 1)
    gap = np.minimum(abs(p[nearest, 0] - f[:, 0]), abs(p[nearest - 1, 0] - f[:, 0]))
    if digest(raw) != expected:
        raise ValueError("Source changed while decoding; derivative is not admissible.")
    np.savez(
        output / "index.npz",
        poses=p,
        frames=f,
        distance=distance,
        frame_distance=frame_distance,
        heldout=held,
        sample_points=np.concatenate(samples),
        sample_frame=np.concatenate(sample_ids),
    )
    meta = dict(
        source_sha256=expected,
        profile=PROFILE,
        frames=len(f),
        poses=len(p),
        points=total,
        seconds=time.monotonic() - started,
        path_m=float(distance[-1]),
        receipt_pose_nearest_gap_p95_s=float(np.quantile(gap, 0.95)),
        receipt_pose_nearest_gap_max_s=float(gap.max()),
        clock_binding="host-monotonic interpolation; NOT hardware synchronization",
        mapped_increment_not_native_sweep=True,
        heldout_frames=int(held.sum()),
        training_frames=int((~held).sum()),
    )
    write_json(output / "source.json", meta)
    return meta


def voxel(points):
    if not len(points):
        return points
    _, idx = np.unique(
        np.floor(points / PROFILE["voxel_m"]).astype(np.int64), axis=0, return_index=True
    )
    return points[np.sort(idx)]


def register(reference, query, seed, *, acquisition=False):
    try:
        result = PreparedReference(reference).register(
            query, seed, policy=CLOSURE_REGISTRATION_POLICY if acquisition else POLICY
        )
        result.pop("matched_query_indices", None)
        return result
    except ValueError as exc:
        return dict(status="unavailable", reasons=[str(exc)])


def links_for(data):
    p, s = data["poses"], data["frame_distance"]
    pts, ids, held = data["sample_points"], data["sample_frame"], data["heldout"]
    training = ~held[ids]

    def take(frame_mask, center, radius):
        mask = training & frame_mask[ids]
        chunk = pts[mask]
        return chunk[np.linalg.norm(chunk - center, axis=1) <= radius]

    link, search = acquire_closure(data, register)
    selected = search["attempts"][search["selected_attempt"]]
    links = [link]
    audits = [dict(kind="seam", **selected, search=search)]
    # Disjoint time/distance windows: no shared frame can self-match across an edge.
    centers = np.linspace(0, data["distance"][-1], int(np.ceil(data["distance"][-1] / 20)) + 1)
    for i, (sa, sb) in enumerate(zip(centers[:-1], centers[1:], strict=True)):
        width = (sb - sa) / 3
        amask, bmask = abs(s - sa) <= width, abs(s - sb) <= width
        assert not np.any(amask & bmask)
        pivot = np.array([np.interp((sa + sb) / 2, data["distance"], p[:, j]) for j in range(1, 4)])
        a = take(amask, pivot, PROFILE["neighbor_radius_m"])
        b = take(bmask, pivot, PROFILE["neighbor_radius_m"])
        fit = register(a, b, np.eye(4))
        audits.append(dict(kind="neighbor", distances_m=[float(sa), float(sb)], fit=fit))
        if fit["status"] == "candidate":
            links.append(
                SurfaceLink(
                    float(np.mean(s[amask & ~held])),
                    float(np.mean(s[bmask & ~held])),
                    np.asarray(fit["T_reference_query"]),
                    np.median(b, axis=0),
                    PROFILE["neighbor_translation_weight_m"],
                    PROFILE["neighbor_rotation_weight_deg"],
                    f"neighbor-{i}",
                )
            )
        print(f"neighbor {i + 1}/{len(centers) - 1}: {fit['status']}", flush=True)
    return links, audits


def evaluate(data, field, label):
    pts, ids = data["sample_points"], data["sample_frame"]
    s, f, p = data["frame_distance"], data["frames"], data["poses"]
    train = ~data["heldout"][ids]
    corrected = np.empty_like(pts)
    offsets = np.searchsorted(ids, np.arange(len(f) + 1))
    for i, distance in enumerate(s):
        start, end = offsets[i : i + 2]
        corrected[start:end] = field.points(pts[start:end], distance)
    target = voxel(corrected[train])
    tree = cKDTree(target)
    groups = np.floor((f[:, 0] - f[0, 0]) / PROFILE["holdout_period_s"]).astype(int)
    seed, rows = np.eye(4), []
    for group in np.unique(groups[data["heldout"]]):
        frames = data["heldout"] & (groups == group)
        seconds, distance = float(np.mean(f[frames, 0])), float(np.mean(s[frames]))
        position = np.array([np.interp(seconds, p[:, 0], p[:, j]) for j in range(1, 4)])
        query = pts[frames[ids]]  # Uncorrected, held-out scanner output.
        radius = PROFILE["local_validation_radius_m"]
        query = query[np.linalg.norm(query - position, axis=1) <= radius]
        estimated = transform(position[None], seed)[0]
        reference = target[tree.query_ball_point(estimated, radius + 5)]
        fit = register(reference, query, seed)
        row = dict(
            group=int(group),
            distance_m=distance,
            source_frames=int(frames.sum()),
            source_query_points=len(query),
            fit=fit,
        )
        if "T_reference_query" in fit:
            fitted = np.asarray(fit["T_reference_query"])
            implied = field.matrices(distance)[0]
            row["model_consistency_m"] = float(
                np.linalg.norm(
                    transform(position[None], fitted) - transform(position[None], implied)
                )
            )
            row["model_consistency_deg"] = float(
                np.rad2deg(
                    np.linalg.norm(
                        Rotation.from_matrix(fitted[:3, :3].T @ implied[:3, :3]).as_rotvec()
                    )
                )
            )
            # All-point tails remain visible, not just accepted correspondences.
            dist, _ = tree.query(transform(query, fitted), workers=1)
            row["all_point_distance_p95_m"] = float(np.quantile(dist, 0.95))
            if fit["status"] == "candidate":
                seed = fitted  # causal last accepted transform, never field oracle.
        rows.append(row)
        if len(rows) % 10 == 0:
            print(f"validation {label}: {len(rows)} windows", flush=True)
    return dict(
        label=label,
        training_map_points=len(target),
        windows=rows,
        candidate_count=sum(r["fit"]["status"] == "candidate" for r in rows),
        total=len(rows),
        interpretation="same-source held-out-frame local matching, not independent truth",
        seed="identity then previous accepted transform; no current-field seed",
    )


def materialize(cache, data, field, output):
    frames = data["frames"]
    count = int(frames[-1, 2] + frames[-1, 3])
    source = np.memmap(cache / "source-points.f32", dtype="<f4", mode="r", shape=(count, 3))
    with (output / "corrected-points.f32").open("xb") as stream:
        for frame, distance in zip(frames, data["frame_distance"], strict=True):
            offset, size = int(frame[2]), int(frame[3])
            field.points(source[offset : offset + size], distance).astype("<f4").tofile(stream)
    pos, q = field.poses(data["poses"][:, 1:4], data["poses"][:, 4:8], data["distance"])
    np.savez(
        output / "corrected-trajectory.npz",
        positions=pos,
        orientations_xyzw=q,
        receipt_time_s=data["poses"][:, 0],
        distance_m=data["distance"],
        frame_distance_m=data["frame_distance"],
        frames=frames,
    )
    return dict(
        points=count,
        path_before_m=float(data["distance"][-1]),
        path_after_m=float(np.linalg.norm(np.diff(pos, axis=0), axis=1).sum()),
        endpoint_delta_before_m=(data["poses"][-1, 1:4] - data["poses"][0, 1:4]).tolist(),
        endpoint_delta_after_m=(pos[-1] - pos[0]).tolist(),
        corrected_points_sha256=digest(output / "corrected-points.f32"),
        corrected_trajectory_sha256=digest(output / "corrected-trajectory.npz"),
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--raw", type=Path, required=True)
    parser.add_argument("--sha256", required=True)
    parser.add_argument("--session-id", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--cache", type=Path)
    args = parser.parse_args()
    if args.raw.resolve().is_relative_to(args.output.resolve()):
        raise ValueError("Output must not contain the source recording.")
    args.output.mkdir(parents=True, exist_ok=False)
    started = time.monotonic()
    cache = args.cache or args.output / "decoded"
    if args.cache:
        meta = json.loads((cache / "source.json").read_text())
        if (
            meta["source_sha256"] != args.sha256
            or digest(args.raw) != args.sha256
            or not compatible_cache_profile(meta["profile"])
        ):
            raise ValueError("Cached source identity mismatch.")
        seal = json.loads((cache / "seal.json").read_text())
        if any(digest(cache / name) != value for name, value in seal.items()):
            raise ValueError("Decoded cache integrity mismatch.")
    else:
        meta = extract(args.raw, cache, args.sha256)
        write_json(
            cache / "seal.json",
            {
                name: digest(cache / name)
                for name in ["source-points.f32", "source-intensity.u8", "index.npz", "source.json"]
            },
        )
    with np.load(cache / "index.npz") as stored:
        data = dict(stored)
    try:
        links, audits = links_for(data)
    except ClosureUnavailable as exc:
        write_json(args.output / "closure-search.json", exc.report)
        raise
    write_json(args.output / "closure-search.json", audits[0]["search"])
    write_json(args.output / "registrations.json", audits)
    # Keep each measured link and its explicit weights reproducible.
    link_doc = []
    for e in links:
        d = asdict(e)
        d["T_reference_query"] = e.T_reference_query.tolist()
        d["query_center"] = e.query_center.tolist()
        link_doc.append(d)
    write_json(args.output / "surface-links.json", link_doc)
    length = float(data["distance"][-1])
    field, fit = fit_correction(length, links)
    write_json(args.output / "correction.json", fit)
    if not fit["converged"]:
        raise ValueError("Correction solver did not converge; do not materialize.")
    original = CorrectionField([0, length], np.zeros((2, 6)))
    validation = [evaluate(data, original, "original"), evaluate(data, field, "corrected")]
    write_json(args.output / "validation.json", validation)
    sensitivity = []
    grid = np.linspace(0, length, 1001)
    route = np.stack(
        [np.interp(grid, data["distance"], data["poses"][:, j]) for j in range(1, 4)], axis=1
    )
    for strength in [0.5, 2.0]:
        other, report = fit_correction(length, links, CorrectionPolicy(strain_weight=strength))
        delta = np.linalg.norm(other.points(route, grid) - field.points(route, grid), axis=1)
        sensitivity.append(
            dict(
                strain_weight=strength,
                converged=report["converged"],
                maximum_route_difference_m=float(delta.max()),
                p95_route_difference_m=float(np.quantile(delta, 0.95)),
            )
        )
    corrected_route = field.points(route, grid)
    displacement = corrected_route - route
    gradient = np.linalg.norm(np.diff(displacement, axis=0), axis=1) / np.diff(grid)
    rotation_gradient = np.rad2deg(np.linalg.norm(field.spline(grid, 1)[:, 3:], axis=1))
    product = materialize(cache, data, field, args.output)
    if digest(args.raw) != args.sha256:
        raise ValueError("Raw source changed during experiment.")
    summary = dict(
        schema_version="missioncore.recorded-ring-experiment/v2",
        created_at=datetime.now(UTC).isoformat(),
        source_session=args.session_id,
        source=meta,
        profile=PROFILE,
        registration_policy=POLICY,
        closure_registration_policy=CLOSURE_REGISTRATION_POLICY,
        closure_policy=asdict(ClosurePolicy()),
        closure_acquisition=audits[0]["search"]["status"],
        versions={m: version(m) for m in ["numpy", "scipy", "small-gicp"]},
        elapsed_s=time.monotonic() - started,
        source_cache=str(cache.resolve()),
        product=product,
        surface_links=len(links),
        sensitivity=sensitivity,
        deformation=dict(
            maximum_route_displacement_m=float(np.linalg.norm(displacement, axis=1).max()),
            max_route_displacement_gradient_m_per_m=float(gradient.max()),
            p95_route_displacement_gradient_m_per_m=float(np.quantile(gradient, 0.95)),
            max_rotation_parameter_gradient_deg_per_m=float(rotation_gradient.max()),
            individual_frame_transform="rigid; no internal scale/shear",
        ),
        validation=[
            dict(label=v["label"], candidates=v["candidate_count"], windows=v["total"])
            for v in validation
        ],
        status="experimental-candidate-not-promoted",
        production_promotion=False,
        vehicle_control=False,
        original_modified=False,
        limitations=[
            "Single source: frame holdout is not an independent pass or ground truth.",
            "K1 mapped increments, no per-point motion reconstruction.",
            "Known start-area revisit; no automatic arbitrary-loop discovery.",
            "Weights are engineering priors, not calibrated sensor covariance.",
        ],
    )
    write_json(args.output / "summary.json", summary)
    print(json.dumps(summary, indent=2), flush=True)


if __name__ == "__main__":
    main()