#!/usr/bin/env python3
"""Scan body-left RPY corrections locally and validate them over every B pair.

This command is diagnostic only.  It never writes or replaces an extrinsic JSON.
"""
from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path

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

from rigorous_calibration import (
    inverse_transform, read_pairs, read_planes, rotation_angle_deg, rpy_deg,
)


def statistics(values):
    values = np.asarray(values, float)
    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)),
    }


def body_left_rpy(x, rpy_correction_deg):
    correction = np.eye(4)
    correction[:3, :3] = Rotation.from_euler(
        "xyz", np.asarray(rpy_correction_deg, float), degrees=True
    ).as_matrix()
    return correction @ x


def pair_delta(a_ij, b_ij, x):
    predicted = inverse_transform(x) @ a_ij @ x
    delta = inverse_transform(b_ij) @ predicted
    translation = np.asarray(delta[:3, 3], float)
    return {
        "translation_xyz_m": translation.tolist(),
        "translation_xyz_cm": (100.0 * translation).tolist(),
        "translation_norm_m": float(np.linalg.norm(translation)),
        "rotation_rpy_deg_xyz": rpy_deg(delta[:3, :3]),
        "rotation_angle_deg": rotation_angle_deg(delta[:3, :3]),
    }


def ground_metrics(planes, x, body_height):
    if len(planes) == 0:
        return None
    up = np.array([0.0, 0.0, 1.0])
    tilt_deg, height_m = [], []
    for plane in planes:
        normal_body = x[:3, :3] @ plane[:3]
        normal_body /= np.linalg.norm(normal_body)
        tilt_deg.append(math.degrees(math.atan2(
            np.linalg.norm(np.cross(normal_body, up)),
            float(np.clip(normal_body @ up, -1.0, 1.0)),
        )))
        height_m.append(
            float(plane[3] - normal_body @ x[:3, 3] - body_height)
        )
    return {
        "normal_tilt_deg": statistics(tilt_deg),
        "height_residual_m": statistics(height_m),
    }


def evaluate(label, correction, a_array, b_array, meta, x, pair_index,
             translation_scale, rotation_scale, planes, body_height):
    candidate_x = body_left_rpy(x, correction)
    per_pair = []
    translation, rotation, normalized = [], [], []
    for index, (a_ij, b_ij, pair_meta) in enumerate(zip(a_array, b_array, meta)):
        item = pair_delta(a_ij, b_ij, candidate_x)
        item.update({
            "pair_index": index,
            "i": int(pair_meta[0]),
            "j": int(pair_meta[1]),
        })
        t = item["translation_norm_m"]
        r = item["rotation_angle_deg"]
        translation.append(t)
        rotation.append(r)
        normalized.append(math.hypot(t / translation_scale, r / rotation_scale))
        per_pair.append(item)
    return {
        "label": label,
        "body_left_rpy_correction_deg_xyz": list(map(float, correction)),
        "candidate_extrinsic": {
            "translation_m": candidate_x[:3, 3].tolist(),
            "rotation_rpy_deg_xyz": rpy_deg(candidate_x[:3, :3]),
        },
        "all_pairs": {
            "count": len(per_pair),
            "translation_m": statistics(translation),
            "rotation_deg": statistics(rotation),
            "normalized_pair_score": statistics(normalized),
            "normalized_global_rms": float(np.sqrt(np.mean(np.asarray(normalized) ** 2))),
        },
        "selected_pair": per_pair[pair_index],
        "ground": ground_metrics(planes, candidate_x, body_height),
        "per_pair": per_pair,
    }


def candidate_grid(pitch_values, roll_values, yaw_values):
    answer = [("baseline", (0.0, 0.0, 0.0))]
    for pitch in pitch_values:
        answer.append((f"pitch_{pitch:+.3f}", (0.0, pitch, 0.0)))
    for pitch in (0.0, *pitch_values):
        for roll in roll_values:
            answer.append((
                f"pitch_{pitch:+.3f}_roll_{roll:+.3f}",
                (roll, pitch, 0.0),
            ))
    for yaw in yaw_values:
        answer.append((f"yaw_{yaw:+.3f}_diagnostic", (0.0, 0.0, yaw)))
    unique = []
    seen = set()
    for label, values in answer:
        key = tuple(round(float(value), 12) for value in values)
        if key not in seen:
            seen.add(key)
            unique.append((label, values))
    return unique


def z_observability(a_array, x, test_shift_m):
    shift = np.eye(4)
    shift[2, 3] = test_shift_m
    shifted_x = shift @ x
    effects = []
    for a_ij in a_array:
        before = inverse_transform(x) @ a_ij @ x
        after = inverse_transform(shifted_x) @ a_ij @ shifted_x
        delta = inverse_transform(before) @ after
        effects.append((
            float(np.linalg.norm(delta[:3, 3])),
            rotation_angle_deg(delta[:3, :3]),
        ))
    effects = np.asarray(effects, float)
    maximum = np.max(effects, axis=0)
    return {
        "body_left_z_test_shift_m": test_shift_m,
        "max_predicted_motion_change_translation_m": float(maximum[0]),
        "max_predicted_motion_change_rotation_deg": float(maximum[1]),
        "numerically_unobservable": bool(maximum[0] < 1e-10 and maximum[1] < 1e-10),
        "note": "AX pairs cannot determine X.z when every A rotation preserves body Z; use ground/external height constraints.",
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pairs", required=True)
    parser.add_argument("--extrinsic", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--csv")
    parser.add_argument("--ground-planes")
    parser.add_argument("--pair-index", type=int, default=0)
    parser.add_argument("--pitch-values", nargs="+", type=float, default=[0.1, 0.2, 0.3])
    parser.add_argument("--roll-values", nargs="+", type=float, default=[-0.2, -0.1, 0.1])
    parser.add_argument("--yaw-values", nargs="+", type=float, default=[-0.2, 0.2])
    parser.add_argument("--translation-scale", type=float, default=0.05)
    parser.add_argument("--rotation-scale", type=float, default=0.5)
    parser.add_argument("--body-height", type=float, default=0.2335)
    args = parser.parse_args()

    a_array, b_array, meta, stations = read_pairs(args.pairs)
    if not 0 <= args.pair_index < len(a_array):
        raise IndexError(f"pair-index {args.pair_index} outside [0,{len(a_array)-1}]")
    with Path(args.extrinsic).open(encoding="utf-8-sig") as stream:
        x = np.asarray(json.load(stream)["matrix_4x4"], float)
    planes = read_planes(args.ground_planes) if args.ground_planes else np.empty((0, 4))

    candidates = [
        evaluate(
            label, correction, a_array, b_array, meta, x, args.pair_index,
            args.translation_scale, args.rotation_scale, planes, args.body_height,
        )
        for label, correction in candidate_grid(
            args.pitch_values, args.roll_values, args.yaw_values
        )
    ]
    baseline = candidates[0]
    baseline_scores = np.asarray([
        math.hypot(
            item["translation_norm_m"] / args.translation_scale,
            item["rotation_angle_deg"] / args.rotation_scale,
        )
        for item in baseline["per_pair"]
    ])
    base_global = baseline["all_pairs"]["normalized_global_rms"]
    for candidate in candidates:
        scores = np.asarray([
            math.hypot(
                item["translation_norm_m"] / args.translation_scale,
                item["rotation_angle_deg"] / args.rotation_scale,
            )
            for item in candidate["per_pair"]
        ])
        delta = scores - baseline_scores
        candidate["comparison_to_baseline"] = {
            "normalized_global_rms_change": float(
                candidate["all_pairs"]["normalized_global_rms"] - base_global
            ),
            "improved_pairs": int(np.sum(delta < -1e-12)),
            "worsened_pairs": int(np.sum(delta > 1e-12)),
            "unchanged_pairs": int(np.sum(np.abs(delta) <= 1e-12)),
            "median_per_pair_score_change": float(np.median(delta)),
            "global_consistency_signal": bool(
                candidate["all_pairs"]["normalized_global_rms"] < base_global
                and np.sum(delta < -1e-12) > np.sum(delta > 1e-12)
            ),
        }

    ranking = sorted(
        candidates,
        key=lambda item: item["all_pairs"]["normalized_global_rms"],
    )
    report = {
        "schema_version": 1,
        "diagnostic_only": True,
        "extrinsic_was_modified": False,
        "equation": "delta_ij = B_ij^-1 * (X^-1 * A_ij * X)",
        "correction_convention": "X_test = DeltaR_body * X; DeltaR uses fixed body xyz RPY axes",
        "component_frame": "delta translation/RPY components are in station-j LiDAR coordinates, not screen axes",
        "selection_rule": (
            "Never accept a correction from selected_pair alone. Require improvement over all "
            "refined pairs, directional consistency across pairs, acceptable ground constraints, "
            "and independent visual review. This script never overwrites X."
        ),
        "pairs_file": str(Path(args.pairs).resolve()),
        "extrinsic_file": str(Path(args.extrinsic).resolve()),
        "stations": stations,
        "pairs": len(a_array),
        "selected_pair_index": args.pair_index,
        "selected_pair_stations": [int(meta[args.pair_index, 0]), int(meta[args.pair_index, 1])],
        "normalization": {
            "translation_scale_m": args.translation_scale,
            "rotation_scale_deg": args.rotation_scale,
        },
        "z_observability": z_observability(a_array, x, 0.10),
        "ranking_by_all_pair_normalized_rms": [
            {
                "rank": rank,
                "label": item["label"],
                "body_left_rpy_correction_deg_xyz": item["body_left_rpy_correction_deg_xyz"],
                "normalized_global_rms": item["all_pairs"]["normalized_global_rms"],
                **item["comparison_to_baseline"],
            }
            for rank, item in enumerate(ranking, 1)
        ],
        "candidates": candidates,
    }

    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")
    csv_path = Path(args.csv) if args.csv else output.with_suffix(".csv")
    with csv_path.open("w", encoding="utf-8", newline="") as stream:
        fields = [
            "label", "roll_correction_deg", "pitch_correction_deg", "yaw_correction_deg",
            "selected_pair_translation_cm", "selected_pair_rotation_deg",
            "all_pair_translation_rms_m", "all_pair_rotation_rms_deg",
            "normalized_global_rms", "normalized_global_rms_change",
            "improved_pairs", "worsened_pairs", "global_consistency_signal",
            "ground_normal_tilt_rms_deg", "ground_height_rms_m",
        ]
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        for item in candidates:
            correction = item["body_left_rpy_correction_deg_xyz"]
            ground = item["ground"]
            comparison = item["comparison_to_baseline"]
            writer.writerow({
                "label": item["label"],
                "roll_correction_deg": correction[0],
                "pitch_correction_deg": correction[1],
                "yaw_correction_deg": correction[2],
                "selected_pair_translation_cm": item["selected_pair"]["translation_norm_m"] * 100.0,
                "selected_pair_rotation_deg": item["selected_pair"]["rotation_angle_deg"],
                "all_pair_translation_rms_m": item["all_pairs"]["translation_m"]["rms"],
                "all_pair_rotation_rms_deg": item["all_pairs"]["rotation_deg"]["rms"],
                "normalized_global_rms": item["all_pairs"]["normalized_global_rms"],
                "normalized_global_rms_change": comparison["normalized_global_rms_change"],
                "improved_pairs": comparison["improved_pairs"],
                "worsened_pairs": comparison["worsened_pairs"],
                "global_consistency_signal": comparison["global_consistency_signal"],
                "ground_normal_tilt_rms_deg": None if ground is None else ground["normal_tilt_deg"]["rms"],
                "ground_height_rms_m": None if ground is None else ground["height_residual_m"]["rms"],
            })

    print(json.dumps({
        "diagnostic_only": True,
        "selected_pair": baseline["selected_pair"],
        "z_observability": report["z_observability"],
        "top_all_pair_candidates": report["ranking_by_all_pair_normalized_rms"][:8],
        "output": str(output.resolve()),
        "csv": str(csv_path.resolve()),
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
