from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

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


def load(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8-sig"))


def write(path: Path, document: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(document, ensure_ascii=False, indent=2), encoding="utf-8")


def inverse(t: np.ndarray) -> np.ndarray:
    result = np.eye(4)
    result[:3, :3] = t[:3, :3].T
    result[:3, 3] = -result[:3, :3] @ t[:3, 3]
    return result


def delta(a: np.ndarray, b: np.ndarray) -> dict:
    d = inverse(a) @ b
    return {
        "translation_m": float(np.linalg.norm(d[:3, 3])),
        "rotation_deg": float(np.linalg.norm(Rotation.from_matrix(d[:3, :3]).as_rotvec()) * 180.0 / math.pi),
        "delta_matrix_4x4": d.tolist(),
    }


def corrected(raw: dict, backend: str, reference_height: float) -> dict:
    return {
        "schema_version": 1,
        "success": bool(raw["success"]),
        "convention": "T_RTK_lidar maps raw LiDAR points into the RTK navigation frame",
        "equation": "A_RTK_ij X = X B_LiDAR_ij",
        "frames": {
            "RTK": {
                "origin": "GGA positioning reference point; confirm ANT1/reference antenna in receiver configuration",
                "x_axis": "horizontal projection of the rawHeading baseline direction reported by the receiver",
                "y_axis": "left",
                "z_axis": "up",
                "yaw_enu_deg": "90 - rawHeadingDeg",
            },
            "LiDAR": "raw LiDAR sensor frame",
        },
        "backend": backend,
        "measured_lidar_extrinsic_used_as_initial": False,
        "body_heading_offset_used": False,
        "body_antenna_lever_xy_used": False,
        "translation_m": raw["translation_m"],
        "rotation_rpy_deg_xyz": raw["rotation_rpy_deg_xyz"],
        "quaternion_xyzw": raw["quaternion_xyzw"],
        "matrix_4x4": raw["matrix_4x4"],
        "quality": {
            "stations": raw["estimation"]["stations"],
            "pairs": raw["estimation"]["pairs"],
            "residuals": raw["estimation"]["residuals"],
            "weighted_jacobian_condition_number": raw["weighted_jacobian_condition_number"],
            "linearized_one_sigma": raw["linearized_one_sigma"],
            "bootstrap": raw["bootstrap"],
        },
        "z_constraint": {
            "observable_from_planar_AX_XB": False,
            "method": "LiDAR ground planes plus externally supplied RTK reference-point height above ground",
            "rtk_reference_height_above_ground_m": reference_height,
            "warning": "z is conditional on the supplied RTK antenna height; it is not independently identified by planar Ackermann motion",
        },
        "important_limit": "AX residual and bootstrap quantify internal consistency, not independent centimetre-grade absolute certification",
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--result-root", type=Path, required=True)
    parser.add_argument("--reference-height", type=float, required=True)
    args = parser.parse_args()

    def solver_output(directory: str) -> Path:
        raw = args.result_root / directory / "extrinsic_raw.json"
        standard = args.result_root / directory / "extrinsic.json"
        return raw if raw.exists() else standard

    paths = {
        "open3d_gicp": solver_output("open3d_gicp"),
        "small_gicp": solver_output("small_gicp"),
        "consensus": solver_output("consensus"),
    }
    docs = {}
    for backend, path in paths.items():
        document = corrected(load(path), backend, args.reference_height)
        write(path.with_name("extrinsic_rtk_lidar.json"), document)
        docs[backend] = document

    open_t = np.asarray(docs["open3d_gicp"]["matrix_4x4"], float)
    small_t = np.asarray(docs["small_gicp"]["matrix_4x4"], float)
    final = dict(docs["consensus"])
    final["selection"] = {
        "recommended": True,
        "reason": "Uses only motion pairs accepted independently by both Open3D GICP and small_gicp",
        "open3d_vs_small_gicp": delta(open_t, small_t),
    }


    write(args.result_root / "final_T_RTK_lidar.json", final)
    summary = {
        "final": {
            "translation_m": final["translation_m"],
            "rotation_rpy_deg_xyz": final["rotation_rpy_deg_xyz"],
            "pairs": final["quality"]["pairs"],
            "translation_rms_m": final["quality"]["residuals"]["translation_m"]["rms"],
            "rotation_rms_deg": final["quality"]["residuals"]["rotation_deg"]["rms"],
            "condition_number": final["quality"]["weighted_jacobian_condition_number"],
        },
        "backend_difference": delta(open_t, small_t),
    }
    write(args.result_root / "summary.json", summary)
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()