#!/usr/bin/env python3
"""Rigorous stationary LiDAR / dual-antenna RTK hand-eye calibration.

Convention: T_A_B maps points from frame B into frame A.
X = T_body_lidar, A_ij = T_W_Bi^-1 T_W_Bj, B_ij = T_Li_Lj,
therefore A_ij X = X B_ij. Raw sensor-frame points_raw are used.
"""
from __future__ import annotations

import argparse
import csv
import json
import math
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from scipy.optimize import least_squares
from scipy.spatial import cKDTree


def skew(v):
    x, y, z = v
    return np.array([[0.0, -z, y], [z, 0.0, -x], [-y, x, 0.0]])


def so3_exp(v):
    angle = float(np.linalg.norm(v))
    if angle < 1e-12:
        return np.eye(3) + skew(v)
    k = skew(np.asarray(v, float) / angle)
    return np.eye(3) + math.sin(angle) * k + (1.0 - math.cos(angle)) * k @ k


def so3_log(rotation):
    cosine = float(np.clip((np.trace(rotation) - 1.0) / 2.0, -1.0, 1.0))
    angle = math.acos(cosine)
    vee = np.array([
        rotation[2, 1] - rotation[1, 2],
        rotation[0, 2] - rotation[2, 0],
        rotation[1, 0] - rotation[0, 1],
    ])
    if angle < 1e-9:
        return vee / 2.0
    if abs(math.pi - angle) < 1e-5:
        values, vectors = np.linalg.eigh((rotation + np.eye(3)) / 2.0)
        return vectors[:, int(np.argmax(values))] * angle
    return vee * angle / (2.0 * math.sin(angle))


def quat_to_rotation(q):
    x, y, z, w = np.asarray(q, float) / np.linalg.norm(q)
    return np.array([
        [1-2*(y*y+z*z), 2*(x*y-z*w), 2*(x*z+y*w)],
        [2*(x*y+z*w), 1-2*(x*x+z*z), 2*(y*z-x*w)],
        [2*(x*z-y*w), 2*(y*z+x*w), 1-2*(x*x+y*y)],
    ])


def rotation_to_quat(rotation):
    from scipy.spatial.transform import Rotation
    return Rotation.from_matrix(rotation).as_quat()


def rpy_deg(rotation):
    from scipy.spatial.transform import Rotation
    return Rotation.from_matrix(rotation).as_euler("xyz", degrees=True).tolist()


def make_transform(translation, rotation):
    transform = np.eye(4)
    transform[:3, :3] = rotation
    transform[:3, 3] = translation
    return transform


def params_transform(params):
    return make_transform(params[:3], so3_exp(params[3:]))


def inverse_transform(transform):
    answer = np.eye(4)
    answer[:3, :3] = transform[:3, :3].T
    answer[:3, 3] = -answer[:3, :3] @ transform[:3, 3]
    return answer


def transform_points(points, transform):
    return points @ transform[:3, :3].T + transform[:3, 3]


def rotation_angle_deg(rotation):
    return math.degrees(np.linalg.norm(so3_log(rotation)))


@dataclass
class PoseSeries:
    time: np.ndarray
    transforms: np.ndarray


def read_poses(path):
    timestamps, transforms = [], []
    with Path(path).open(encoding="utf-8-sig", newline="") as stream:
        reader = csv.DictReader(stream)
        required = ("time", "x", "y", "z", "qx", "qy", "qz", "qw")
        missing = [key for key in required if key not in (reader.fieldnames or [])]
        if missing:
            raise ValueError(f"{path}: missing pose fields {missing}")
        for row in reader:
            timestamps.append(float(row["time"]))
            translation = np.array([float(row[k]) for k in ("x", "y", "z")])
            quaternion = np.array([float(row[k]) for k in ("qx", "qy", "qz", "qw")])
            transforms.append(make_transform(translation, quat_to_rotation(quaternion)))
    order = np.argsort(timestamps)
    return PoseSeries(np.asarray(timestamps)[order], np.asarray(transforms)[order])


def nearest_pose(series, timestamp):
    index = int(np.argmin(np.abs(series.time - timestamp)))
    return series.transforms[index], float(abs(series.time[index] - timestamp))


def npz_files(root):
    files = sorted(Path(root).rglob("*.npz"))
    if not files:
        raise FileNotFoundError(f"no NPZ files under {root}")
    return files


def load_npz_xyz(path, min_range=1.0, max_range=50.0):
    with np.load(path, allow_pickle=False) as data:
        if "points_raw" not in data:
            raise ValueError(f"{path}: points_raw is required; cart-frame points are forbidden")
        raw = np.asarray(data["points_raw"], dtype=np.float64)
        timestamp = float(np.ravel(data["unix_time_ns"])[0]) / 1e9
        counter = int(np.ravel(data["frame_counter"])[0])
    distance = raw[:, 0] * 0.001
    azimuth = np.deg2rad(raw[:, 1])
    altitude = np.deg2rad(raw[:, 2])
    valid = (
        np.isfinite(distance + azimuth + altitude)
        & (distance >= min_range)
        & (distance <= max_range)
    )
    distance, azimuth, altitude = distance[valid], azimuth[valid], altitude[valid]
    xyz = np.column_stack((
        distance * np.cos(altitude) * np.cos(azimuth),
        distance * np.cos(altitude) * np.sin(azimuth),
        distance * np.sin(altitude),
    ))
    return timestamp, counter, xyz


def load_stations(root, min_range, max_range, z_min=None, z_max=None):
    stations = []
    for path in npz_files(root):
        timestamp, counter, xyz = load_npz_xyz(path, min_range, max_range)
        if z_min is not None:
            xyz = xyz[(xyz[:, 2] >= z_min) & (xyz[:, 2] <= z_max)]
        stations.append((timestamp, counter, path, xyz))
    stations.sort(key=lambda item: item[0])
    return stations


def split_holdout(points, fraction, phase):
    stride = max(int(round(1.0 / fraction)), 2)
    index = np.arange(len(points))
    holdout = ((index + phase) % stride) == 0
    return points[~holdout], points[holdout]


def make_o3d_cloud(points, voxel):
    import open3d as o3d
    cloud = o3d.geometry.PointCloud()
    cloud.points = o3d.utility.Vector3dVector(np.asarray(points, float))
    return cloud.voxel_down_sample(voxel)


def align_open3d(target, source, initial, voxels, correspondences, iterations):
    import open3d as o3d
    registration = o3d.pipelines.registration
    estimate = registration.TransformationEstimationForGeneralizedICP()
    criteria = registration.ICPConvergenceCriteria(max_iteration=iterations)
    transform, stages = np.asarray(initial, float), []
    final_target = final_source = final_answer = None
    started = time.perf_counter()
    for voxel, correspondence in zip(voxels, correspondences):
        target_cloud = make_o3d_cloud(target, voxel)
        source_cloud = make_o3d_cloud(source, voxel)
        answer = registration.registration_generalized_icp(
            source_cloud, target_cloud, correspondence, transform, estimate, criteria
        )
        transform = np.asarray(answer.transformation, float)
        stages.append({
            "voxel_m": voxel,
            "max_correspondence_m": correspondence,
            "fitness": float(answer.fitness),
            "inlier_rmse_m": float(answer.inlier_rmse),
            "target_points": len(target_cloud.points),
            "source_points": len(source_cloud.points),
        })
        final_target, final_source, final_answer = target_cloud, source_cloud, answer
    information = registration.get_information_matrix_from_point_clouds(
        final_source, final_target, correspondences[-1], transform
    )
    inliers = int(round(float(final_answer.fitness) * len(final_source.points)))
    return {
        "transform": transform,
        "hessian": np.asarray(information, float),
        "converged": None,
        "iterations": None,
        "num_inliers": inliers,
        "objective": float(final_answer.inlier_rmse ** 2 * max(inliers, 1)),
        "elapsed_sec": time.perf_counter() - started,
        "stages": stages,
    }


def align_small_gicp(target, source, initial, voxels, correspondences, iterations, threads):
    import small_gicp
    transform, stages, result = np.asarray(initial, float), [], None
    started = time.perf_counter()
    for voxel, correspondence in zip(voxels, correspondences):
        result = small_gicp.align(
            np.ascontiguousarray(target),
            np.ascontiguousarray(source),
            transform,
            registration_type="GICP",
            downsampling_resolution=voxel,
            max_correspondence_distance=correspondence,
            num_threads=threads,
            max_iterations=iterations,
            rotation_epsilon=math.radians(0.005),
            translation_epsilon=0.0005,
            verbose=False,
        )
        transform = np.asarray(result.T_target_source, float)
        stages.append({
            "voxel_m": voxel,
            "max_correspondence_m": correspondence,
            "converged": bool(result.converged),
            "iterations": int(result.iterations),
            "num_inliers": int(result.num_inliers),
            "objective": float(result.error),
        })
    return {
        "transform": transform,
        "hessian": np.asarray(result.H, float),
        "converged": bool(result.converged),
        "iterations": int(result.iterations),
        "num_inliers": int(result.num_inliers),
        "objective": float(result.error),
        "elapsed_sec": time.perf_counter() - started,
        "stages": stages,
    }


def align_backend(backend, target, source, initial, args):
    if backend == "open3d":
        return align_open3d(
            target, source, initial, args.voxels, args.correspondences, args.iterations
        )
    return align_small_gicp(
        target, source, initial, args.voxels, args.correspondences,
        args.iterations, args.threads
    )


def symmetric_heldout_metrics(target_fit, target_holdout, source_fit, source_holdout,
                              transform, threshold):
    transformed_source_fit = transform_points(source_fit, transform)
    transformed_source_holdout = transform_points(source_holdout, transform)
    forward = cKDTree(target_fit).query(transformed_source_holdout, workers=-1)[0]
    reverse = cKDTree(transformed_source_fit).query(target_holdout, workers=-1)[0]
    distances = np.concatenate((forward, reverse))
    inliers = distances[distances <= threshold]
    return {
        "evaluated": int(len(distances)),
        "inliers": int(len(inliers)),
        "inlier_ratio": float(len(inliers) / max(len(distances), 1)),
        "inlier_rmse_m": float(np.sqrt(np.mean(inliers**2))) if len(inliers) else None,
        "median_m": float(np.median(distances)),
        "p90_m": float(np.quantile(distances, 0.90)),
        "p95_m": float(np.quantile(distances, 0.95)),
    }


def hessian_metrics(hessian, characteristic_length=10.0):
    hessian = 0.5 * (np.asarray(hessian, float) + np.asarray(hessian, float).T)
    scale = np.diag([1.0 / characteristic_length] * 3 + [1.0] * 3)
    scaled = scale.T @ hessian @ scale
    values, vectors = np.linalg.eigh(scaled)
    largest = max(float(np.max(np.abs(values))), np.finfo(float).eps)
    positive = values[values > largest * 1e-9]
    condition = float(positive[-1] / positive[0]) if len(positive) else float("inf")
    return {
        "native_order": ["rx_rad", "ry_rad", "rz_rad", "tx_m", "ty_m", "tz_m"],
        "scaled_eigenvalues": values.tolist(),
        "effective_rank": int(len(positive)),
        "scaled_condition_number": condition,
        "weakest_scaled_direction": vectors[:, int(np.argmin(values))].tolist(),
    }


def transform_difference(reference, candidate):
    delta = inverse_transform(reference) @ candidate
    return {
        "translation_m": float(np.linalg.norm(delta[:3, 3])),
        "rotation_deg": rotation_angle_deg(delta[:3, :3]),
    }


def loop_metrics(transforms):
    loops = []
    for (i, j), b_ij in transforms.items():
        for (j2, k), b_jk in transforms.items():
            if j2 != j or (i, k) not in transforms:
                continue
            loops.append(transform_difference(transforms[(i, k)], b_ij @ b_jk))
    if not loops:
        return {"count": 0}
    translation = np.array([item["translation_m"] for item in loops])
    rotation = np.array([item["rotation_deg"] for item in loops])
    return {
        "count": len(loops),
        "translation_rms_m": float(np.sqrt(np.mean(translation**2))),
        "translation_p95_m": float(np.quantile(translation, 0.95)),
        "rotation_rms_deg": float(np.sqrt(np.mean(rotation**2))),
        "rotation_p95_deg": float(np.quantile(rotation, 0.95)),
    }


def cmd_ground(args):
    stations = load_stations(args.frames, args.min_range, args.max_range)
    rows = []
    for timestamp, counter, _, xyz in stations:
        roi = xyz[(xyz[:, 2] >= args.z_min) & (xyz[:, 2] <= args.z_max)]
        if len(roi) < args.min_inliers:
            continue
        cloud = make_o3d_cloud(roi, args.voxel)
        plane, indexes = cloud.segment_plane(
            args.distance_threshold, 3, args.ransac_iterations
        )
        normal = np.asarray(plane[:3], float)
        norm = np.linalg.norm(normal)
        normal, distance = normal / norm, float(plane[3] / norm)
        if distance < 0:
            normal, distance = -normal, -distance
        points = np.asarray(cloud.points)[indexes]
        rms = float(np.sqrt(np.mean((points @ normal + distance) ** 2)))
        if len(indexes) >= args.min_inliers and rms <= args.max_rms:
            rows.append([timestamp, *normal, distance, len(indexes), rms, counter])
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(["time", "nx", "ny", "nz", "d", "inliers", "rms_m", "frame_counter"])
        writer.writerows(rows)
    print(json.dumps({"planes": len(rows), "output": str(output.resolve())}, indent=2))


def cmd_pairs(args):
    if len(args.voxels) != len(args.correspondences):
        raise ValueError("--voxels and --correspondences must have equal lengths")
    stations = load_stations(
        args.frames, args.min_range, args.max_range, args.z_min, args.z_max
    )
    body = read_poses(args.body)
    if len(stations) < args.min_stations:
        raise ValueError(f"need at least {args.min_stations} stations, got {len(stations)}")
    body_poses, body_dt = [], []
    for timestamp, _, _, xyz in stations:
        if len(xyz) < args.min_roi_points:
            raise ValueError(f"station at {timestamp} has only {len(xyz)} ROI points")
        pose, dt = nearest_pose(body, timestamp + args.time_offset)
        body_poses.append(pose)
        body_dt.append(dt)
    body_poses = np.asarray(body_poses)
    split = [split_holdout(station[3], args.holdout_fraction, i)
             for i, station in enumerate(stations)]
    rng = np.random.default_rng(args.seed)
    accepted_a, accepted_b, accepted_meta, reports = [], [], [], []
    accepted_transforms = {}
    for i in range(len(stations)):
        for j in range(i + args.min_gap, min(len(stations), i + args.max_gap + 1)):
            a_ij = inverse_transform(body_poses[i]) @ body_poses[j]
            translation = float(np.linalg.norm(a_ij[:2, 3]))
            rotation = rotation_angle_deg(a_ij[:3, :3])
            if translation < args.min_translation and rotation < args.min_rotation:
                continue
            initial_b = a_ij.copy()  # X0=I; no measured extrinsic.
            target_fit, target_holdout = split[i]
            source_fit, source_holdout = split[j]
            forward = align_backend(args.backend, target_fit, source_fit, initial_b, args)
            heldout = symmetric_heldout_metrics(
                target_fit, target_holdout, source_fit, source_holdout,
                forward["transform"], args.evaluation_distance
            )
            hessian = hessian_metrics(forward["hessian"])
            reverse_answer = align_backend(
                args.backend, source_fit, target_fit, inverse_transform(initial_b), args
            )
            reverse = transform_difference(
                forward["transform"], inverse_transform(reverse_answer["transform"])
            )
            multistart = []
            for _ in range(args.multistart):
                perturb = np.r_[
                    rng.normal(0.0, args.multistart_translation_sigma, 3),
                    np.deg2rad(rng.normal(0.0, args.multistart_rotation_sigma, 3)),
                ]
                candidate = align_backend(
                    args.backend, target_fit, source_fit,
                    params_transform(perturb) @ initial_b, args
                )
                multistart.append(transform_difference(forward["transform"], candidate["transform"]))
            stable = [
                item["translation_m"] <= args.multistart_translation_tolerance
                and item["rotation_deg"] <= args.multistart_rotation_tolerance
                for item in multistart
            ]
            success_rate = float(np.mean(stable)) if stable else 1.0
            reasons = []
            if forward["converged"] is False:
                reasons.append("backend_not_converged")
            if heldout["inlier_ratio"] < args.min_inlier_ratio:
                reasons.append("heldout_inlier_ratio")
            if heldout["inlier_rmse_m"] is None or heldout["inlier_rmse_m"] > args.max_inlier_rmse:
                reasons.append("heldout_inlier_rmse")
            if hessian["effective_rank"] < 6:
                reasons.append("hessian_rank")
            if hessian["scaled_condition_number"] > args.max_hessian_condition:
                reasons.append("hessian_condition")
            if reverse["translation_m"] > args.reverse_translation_tolerance:
                reasons.append("forward_reverse_translation")
            if reverse["rotation_deg"] > args.reverse_rotation_tolerance:
                reasons.append("forward_reverse_rotation")
            if success_rate < args.min_multistart_success:
                reasons.append("multistart_instability")
            accepted = not reasons
            report = {
                "i": i, "j": j,
                "lidar_time_i": stations[i][0], "lidar_time_j": stations[j][0],
                "frame_counter_i": stations[i][1], "frame_counter_j": stations[j][1],
                "rtk_translation_m": translation, "rtk_rotation_deg": rotation,
                "nearest_rtk_dt_i_s": body_dt[i], "nearest_rtk_dt_j_s": body_dt[j],
                "initial_B_source": "X0=identity; B0=A (no measured extrinsic)",
                "B_ij_4x4": forward["transform"].tolist(),
                "backend": args.backend, "backend_converged": forward["converged"],
                "backend_iterations": forward["iterations"],
                "backend_num_inliers": forward["num_inliers"],
                "backend_objective": forward["objective"],
                "backend_elapsed_sec": forward["elapsed_sec"],
                "multiscale_stages": forward["stages"],
                "heldout_symmetric": heldout, "hessian": hessian,
                "forward_reverse": reverse,
                "multistart": {"runs": len(multistart), "success_rate": success_rate,
                               "deltas": multistart},
                "accepted": accepted, "rejection_reasons": reasons,
            }
            reports.append(report)
            print(f"{args.backend} {i:02d}->{j:02d} rmse={heldout['inlier_rmse_m']} "
                  f"ratio={heldout['inlier_ratio']:.3f} accepted={accepted}")
            if accepted:
                accepted_a.append(a_ij)
                accepted_b.append(forward["transform"])
                accepted_meta.append([i, j, stations[i][0], stations[j][0]])
                accepted_transforms[(i, j)] = forward["transform"]
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        output, A=np.asarray(accepted_a), B=np.asarray(accepted_b),
        meta=np.asarray(accepted_meta),
        station_times=np.asarray([item[0] for item in stations]),
        rtk_nearest_dt_s=np.asarray(body_dt), backend=np.asarray(args.backend),
    )
    quality = {
        "schema_version": 2,
        "backend": args.backend,
        "transform_convention": "B_ij=T_Li_Lj maps station j points into station i",
        "raw_point_field": "points_raw",
        "measured_extrinsic_used_as_initial": False,
        "stations": len(stations), "candidate_pairs": len(reports),
        "accepted_pairs": len(accepted_a),
        "parameters": vars(args),
        "accepted_loop_closure": loop_metrics(accepted_transforms),
        "pairs": reports,
    }
    quality["parameters"].pop("func", None)
    quality_path = Path(args.quality_json or output.with_suffix(".quality.json"))
    quality_path.write_text(json.dumps(quality, ensure_ascii=False, indent=2), encoding="utf-8")
    csv_path = Path(args.quality_csv or output.with_suffix(".quality.csv"))
    with csv_path.open("w", encoding="utf-8", newline="") as stream:
        fields = ["i", "j", "rtk_translation_m", "rtk_rotation_deg",
                  "heldout_inlier_ratio", "heldout_inlier_rmse_m",
                  "hessian_rank", "hessian_condition", "reverse_translation_m",
                  "reverse_rotation_deg", "multistart_success_rate", "accepted",
                  "rejection_reasons"]
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        for item in reports:
            writer.writerow({
                "i": item["i"], "j": item["j"],
                "rtk_translation_m": item["rtk_translation_m"],
                "rtk_rotation_deg": item["rtk_rotation_deg"],
                "heldout_inlier_ratio": item["heldout_symmetric"]["inlier_ratio"],
                "heldout_inlier_rmse_m": item["heldout_symmetric"]["inlier_rmse_m"],
                "hessian_rank": item["hessian"]["effective_rank"],
                "hessian_condition": item["hessian"]["scaled_condition_number"],
                "reverse_translation_m": item["forward_reverse"]["translation_m"],
                "reverse_rotation_deg": item["forward_reverse"]["rotation_deg"],
                "multistart_success_rate": item["multistart"]["success_rate"],
                "accepted": item["accepted"],
                "rejection_reasons": ";".join(item["rejection_reasons"]),
            })
    if len(accepted_a) < args.min_pairs:
        raise RuntimeError(f"only {len(accepted_a)} accepted pairs; need {args.min_pairs}")
    print(json.dumps({
        "backend": args.backend, "stations": len(stations),
        "candidate_pairs": len(reports), "accepted_pairs": len(accepted_a),
        "output": str(output.resolve()), "quality_json": str(quality_path.resolve()),
        "loop": quality["accepted_loop_closure"],
    }, indent=2))


def read_planes(path):
    planes = []
    with Path(path).open(encoding="utf-8-sig", newline="") as stream:
        for row in csv.DictReader(stream):
            normal = np.array([float(row[k]) for k in ("nx", "ny", "nz")])
            norm = np.linalg.norm(normal)
            normal, distance = normal / norm, float(row["d"]) / norm
            if distance < 0:
                normal, distance = -normal, -distance
            planes.append([*normal, distance])
    return np.asarray(planes)


def read_pairs(path):
    with np.load(path, allow_pickle=False) as data:
        return (np.asarray(data["A"], float), np.asarray(data["B"], float),
                np.asarray(data["meta"], float), len(data["station_times"]))


def calibration_residual(params, a_array, b_array, planes, args):
    x = params_transform(params)
    values = []
    for a_ij, b_ij in zip(a_array, b_array):
        error = inverse_transform(a_ij @ x) @ x @ b_ij
        values.extend((error[:3, 3] / args.translation_sigma).tolist())
        values.extend((so3_log(error[:3, :3]) / math.radians(args.rotation_sigma)).tolist())
    body_up = np.array([0.0, 0.0, 1.0])
    for plane in planes:
        normal_body = x[:3, :3] @ plane[:3]
        values.extend((np.cross(normal_body, body_up) / args.plane_normal_sigma).tolist())
        body_distance = plane[3] - float(normal_body @ x[:3, 3])
        values.append((body_distance - args.body_height) / args.plane_height_sigma)
    return np.asarray(values)


def pair_metrics(a_array, b_array, x):
    translation, rotation, rows = [], [], []
    for index, (a_ij, b_ij) in enumerate(zip(a_array, b_array)):
        predicted = inverse_transform(x) @ a_ij @ x
        delta = inverse_transform(b_ij) @ predicted
        t = float(np.linalg.norm(delta[:3, 3]))
        r = rotation_angle_deg(delta[:3, :3])
        translation.append(t); rotation.append(r)
        rows.append({"pair_index": index, "translation_m": t, "rotation_deg": r})
    translation, rotation = np.asarray(translation), np.asarray(rotation)
    def stats(values):
        return {
            "rms": float(np.sqrt(np.mean(values**2))),
            "median": float(np.median(values)),
            "p90": float(np.quantile(values, 0.90)),
            "p95": float(np.quantile(values, 0.95)),
            "max": float(np.max(values)),
        }
    return {"pairs": len(rows), "translation_m": stats(translation),
            "rotation_deg": stats(rotation), "per_pair": rows}


def solve_extrinsic(a_array, b_array, planes, args):
    rng = np.random.default_rng(args.seed)
    starts = [np.zeros(6)]
    for _ in range(args.solver_multistart - 1):
        starts.append(np.r_[
            rng.normal(0.0, args.start_translation_sigma, 3),
            np.deg2rad(rng.normal(0.0, args.start_rotation_sigma, 3)),
        ])
    candidates = []
    lower = np.r_[[-5.0] * 3, [-math.pi] * 3]
    upper = np.r_[[5.0] * 3, [math.pi] * 3]
    for start in starts:
        answer = least_squares(
            calibration_residual, np.clip(start, lower, upper),
            args=(a_array, b_array, planes, args),
            bounds=(lower, upper), loss="huber", f_scale=1.5,
            x_scale="jac", max_nfev=args.max_nfev,
        )
        candidates.append(answer)
    best = min(candidates, key=lambda item: item.cost)
    return best, candidates


def cmd_calibrate(args):
    a_array, b_array, meta, stations = read_pairs(args.pairs)
    planes = read_planes(args.ground_planes)
    best, candidates = solve_extrinsic(a_array, b_array, planes, args)
    x = params_transform(best.x)
    residual = calibration_residual(best.x, a_array, b_array, planes, args)
    absolute = np.abs(residual)
    weights = np.ones_like(residual)
    weights[absolute > 1.5] = 1.5 / absolute[absolute > 1.5]
    weighted_jacobian = best.jac * np.sqrt(weights)[:, None]
    singular = np.linalg.svd(weighted_jacobian, compute_uv=False)
    condition = float(singular[0] / max(singular[-1], 1e-15))
    dof = max(len(residual) - 6, 1)
    covariance = np.linalg.pinv(weighted_jacobian.T @ weighted_jacobian) * float(
        np.sum(weights * residual**2) / dof
    )
    sigma = np.sqrt(np.maximum(np.diag(covariance), 0.0))
    candidate_summary = []
    for item in candidates:
        candidate_x = params_transform(item.x)
        candidate_summary.append({
            "cost": float(item.cost), "success": bool(item.success),
            **transform_difference(x, candidate_x),
        })
    bootstrap = []
    rng = np.random.default_rng(args.seed + 1)
    for _ in range(args.bootstrap):
        indexes = rng.integers(0, len(a_array), len(a_array))
        answer = least_squares(
            calibration_residual, best.x,
            args=(a_array[indexes], b_array[indexes], planes, args),
            loss="huber", f_scale=1.5, x_scale="jac", max_nfev=args.max_nfev,
        )
        bootstrap.append(np.r_[answer.x[:3], rpy_deg(so3_exp(answer.x[3:]))])
    bootstrap = np.asarray(bootstrap)
    result = {
        "schema_version": 2,
        "success": bool(best.success),
        "message": best.message,
        "convention": "T_body_lidar maps raw LiDAR points into rear-axle body frame",
        "equation": "A_ij X = X B_ij",
        "measured_extrinsic_used_as_initial": False,
        "translation_m": x[:3, 3].tolist(),
        "rotation_rpy_deg_xyz": rpy_deg(x[:3, :3]),
        "quaternion_xyzw": rotation_to_quat(x[:3, :3]).tolist(),
        "matrix_4x4": x.tolist(),
        "estimation": {"stations": stations, "pairs": len(a_array),
                       "residuals": pair_metrics(a_array, b_array, x)},
        "ground": {
            "planes": len(planes),
            "body_origin_height_above_ground_m": args.body_height,
            "formula": "d_lidar - (R_X n_lidar)^T t_X - body_height",
        },
        "linearized_one_sigma": {
            "translation_m": sigma[:3].tolist(),
            "rotation_deg": np.rad2deg(sigma[3:]).tolist(),
            "warning": "conditional local estimate; bootstrap is the primary stability check",
        },
        "weighted_jacobian_condition_number": condition,
        "solver_multistart": {
            "runs": len(candidates), "candidates_relative_to_best": candidate_summary,
        },
        "bootstrap": {
            "runs": len(bootstrap),
            "order": ["x_m", "y_m", "z_m", "roll_deg", "pitch_deg", "yaw_deg"],
            "std": np.std(bootstrap, axis=0, ddof=1).tolist() if len(bootstrap) > 1 else None,
            "p025": np.quantile(bootstrap, 0.025, axis=0).tolist() if len(bootstrap) else None,
            "p975": np.quantile(bootstrap, 0.975, axis=0).tolist() if len(bootstrap) else None,
        },
    }
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(result, ensure_ascii=False, indent=2))


def cmd_validate(args):
    result = json.loads(Path(args.extrinsic).read_text(encoding="utf-8-sig"))
    x = np.asarray(result["matrix_4x4"], float)
    a_array, b_array, meta, stations = read_pairs(args.pairs)
    metrics = pair_metrics(a_array, b_array, x)
    for row, pair_meta in zip(metrics["per_pair"], meta):
        row.update({"i": int(pair_meta[0]), "j": int(pair_meta[1])})
    report = {
        "role": "auxiliary check only; first-batch RTK is sparse",
        "blind_with_respect_to_X": True,
        "note": "No AX residual was used to select these pairs",
        "stations": stations, "metrics": metrics,
    }
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2))


def build_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    ground = commands.add_parser("ground")
    ground.add_argument("--frames", required=True); ground.add_argument("--output", required=True)
    ground.add_argument("--min-range", type=float, default=1.0); ground.add_argument("--max-range", type=float, default=30.0)
    ground.add_argument("--z-min", type=float, default=-1.4); ground.add_argument("--z-max", type=float, default=-0.4)
    ground.add_argument("--voxel", type=float, default=0.08); ground.add_argument("--distance-threshold", type=float, default=0.025)
    ground.add_argument("--ransac-iterations", type=int, default=500); ground.add_argument("--min-inliers", type=int, default=500)
    ground.add_argument("--max-rms", type=float, default=0.025); ground.set_defaults(func=cmd_ground)

    pairs = commands.add_parser("pairs")
    pairs.add_argument("--backend", choices=["open3d", "small_gicp"], required=True)
    pairs.add_argument("--frames", required=True); pairs.add_argument("--body", required=True)
    pairs.add_argument("--output", required=True); pairs.add_argument("--quality-json"); pairs.add_argument("--quality-csv")
    pairs.add_argument("--time-offset", type=float, default=0.0)
    pairs.add_argument("--min-stations", type=int, default=30); pairs.add_argument("--min-pairs", type=int, default=25)
    pairs.add_argument("--min-gap", type=int, default=1); pairs.add_argument("--max-gap", type=int, default=5)
    pairs.add_argument("--min-translation", type=float, default=0.5); pairs.add_argument("--min-rotation", type=float, default=3.0)
    pairs.add_argument("--min-range", type=float, default=2.0); pairs.add_argument("--max-range", type=float, default=50.0)
    pairs.add_argument("--z-min", type=float, default=-0.60); pairs.add_argument("--z-max", type=float, default=5.0)
    pairs.add_argument("--min-roi-points", type=int, default=1000)
    pairs.add_argument("--holdout-fraction", type=float, default=0.20)
    pairs.add_argument("--voxels", nargs="+", type=float, default=[0.30, 0.15, 0.08])
    pairs.add_argument("--correspondences", nargs="+", type=float, default=[1.20, 0.50, 0.25])
    pairs.add_argument("--iterations", type=int, default=60); pairs.add_argument("--threads", type=int, default=8)
    pairs.add_argument("--evaluation-distance", type=float, default=0.25)
    pairs.add_argument("--min-inlier-ratio", type=float, default=0.35); pairs.add_argument("--max-inlier-rmse", type=float, default=0.16)
    pairs.add_argument("--max-hessian-condition", type=float, default=1e8)
    pairs.add_argument("--reverse-translation-tolerance", type=float, default=0.08)
    pairs.add_argument("--reverse-rotation-tolerance", type=float, default=0.50)
    pairs.add_argument("--multistart", type=int, default=2)
    pairs.add_argument("--multistart-translation-sigma", type=float, default=0.30)
    pairs.add_argument("--multistart-rotation-sigma", type=float, default=3.0)
    pairs.add_argument("--multistart-translation-tolerance", type=float, default=0.08)
    pairs.add_argument("--multistart-rotation-tolerance", type=float, default=0.50)
    pairs.add_argument("--min-multistart-success", type=float, default=0.50)
    pairs.add_argument("--seed", type=int, default=20260721); pairs.set_defaults(func=cmd_pairs)

    calibrate = commands.add_parser("calibrate")
    calibrate.add_argument("--pairs", required=True); calibrate.add_argument("--ground-planes", required=True)
    calibrate.add_argument("--output", required=True)
    calibrate.add_argument("--translation-sigma", type=float, default=0.05)
    calibrate.add_argument("--rotation-sigma", type=float, default=0.5)
    calibrate.add_argument("--plane-normal-sigma", type=float, default=0.02)
    calibrate.add_argument("--plane-height-sigma", type=float, default=0.03)
    calibrate.add_argument("--body-height", type=float, default=0.2335)
    calibrate.add_argument("--solver-multistart", type=int, default=12)
    calibrate.add_argument("--start-translation-sigma", type=float, default=1.0)
    calibrate.add_argument("--start-rotation-sigma", type=float, default=20.0)
    calibrate.add_argument("--bootstrap", type=int, default=100)
    calibrate.add_argument("--max-nfev", type=int, default=1000)
    calibrate.add_argument("--seed", type=int, default=20260721); calibrate.set_defaults(func=cmd_calibrate)

    validate = commands.add_parser("validate")
    validate.add_argument("--pairs", required=True); validate.add_argument("--extrinsic", required=True)
    validate.add_argument("--output", required=True); validate.set_defaults(func=cmd_validate)
    return parser


def main():
    args = build_parser().parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
