#!/usr/bin/env python3
"""Compare dora calibration with reference calibration."""
import argparse
import numpy as np
from scipy.spatial.transform import Rotation


def load_calibration(path):
    """Load calibration from .npz file."""
    data = np.load(path, allow_pickle=True)
    result = {"T_cam2gripper": data["T_cam2gripper"]}
    if "camera_matrix" in data:
        result["camera_matrix"] = data["camera_matrix"]
    if "distortion" in data:
        result["distortion"] = data["distortion"]
    return result


def analyze_transform(T, name):
    """Print analysis of a 4x4 transformation matrix."""
    print(f"\n=== {name} ===")
    print(f"Translation (mm): [{T[0,3]*1000:.2f}, {T[1,3]*1000:.2f}, {T[2,3]*1000:.2f}]")
    R = T[:3, :3]
    rpy = Rotation.from_matrix(R).as_euler('xyz', degrees=True)
    print(f"RPY (deg): [{rpy[0]:.2f}, {rpy[1]:.2f}, {rpy[2]:.2f}]")
    print(f"Full matrix:\n{T}")


def analyze_camera(cam_matrix, name):
    """Print camera intrinsics analysis."""
    print(f"\n=== {name} Camera Intrinsics ===")
    fx, fy = cam_matrix[0, 0], cam_matrix[1, 1]
    cx, cy = cam_matrix[0, 2], cam_matrix[1, 2]
    print(f"fx={fx:.2f}, fy={fy:.2f}, cx={cx:.2f}, cy={cy:.2f}")

    # Sanity checks
    issues = []
    if fx < 100 or fy < 100:
        issues.append(f"Focal lengths too small (fx={fx:.1f}, fy={fy:.1f})")
    if cx < 100 or cy < 100:
        issues.append(f"Principal point too small (cx={cx:.1f}, cy={cy:.1f})")
    if abs(fx - fy) > 50:
        issues.append(f"Large focal length asymmetry: |fx-fy|={abs(fx-fy):.1f}")

    if issues:
        print("WARNINGS:")
        for issue in issues:
            print(f"  - {issue}")
    else:
        print("Intrinsics appear valid")


def main():
    parser = argparse.ArgumentParser(description="Compare dora vs reference calibration.")
    parser.add_argument("--dora", default="calibration.npz",
                        help="Path to dora calibration result")
    parser.add_argument("--reference", default="reference_calibration.npz",
                        help="Path to reference calibration result")
    parser.add_argument("--check-intrinsics", action="store_true",
                        help="Also compare camera intrinsics")
    args = parser.parse_args()

    try:
        dora_data = load_calibration(args.dora)
        dora = dora_data["T_cam2gripper"]
    except FileNotFoundError:
        print(f"Dora calibration not found: {args.dora}")
        return

    try:
        ref_data = load_calibration(args.reference)
        ref = ref_data["T_cam2gripper"]
    except FileNotFoundError:
        print(f"Reference calibration not found: {args.reference}")
        print("Analyzing dora calibration only:")
        analyze_transform(dora, "Dora Calibration")
        if args.check_intrinsics and "camera_matrix" in dora_data:
            analyze_camera(dora_data["camera_matrix"], "Dora")
        return

    analyze_transform(dora, "Dora Calibration")
    analyze_transform(ref, "Reference Calibration")

    # Check camera intrinsics if requested
    if args.check_intrinsics:
        if "camera_matrix" in dora_data:
            analyze_camera(dora_data["camera_matrix"], "Dora")
        if "camera_matrix" in ref_data:
            analyze_camera(ref_data["camera_matrix"], "Reference")

    # Compute differences
    print("\n=== DIFFERENCES ===")
    trans_diff = np.linalg.norm(dora[:3,3] - ref[:3,3]) * 1000
    print(f"Translation difference: {trans_diff:.2f} mm")

    R_dora = dora[:3,:3]
    R_ref = ref[:3,:3]
    R_diff = R_dora @ R_ref.T
    angle_diff = np.rad2deg(np.arccos(np.clip((np.trace(R_diff) - 1) / 2, -1, 1)))
    print(f"Rotation difference: {angle_diff:.2f} deg")

    # Thresholds for acceptable calibration
    if trans_diff < 5.0 and angle_diff < 2.0:
        print("\n[OK] Calibrations are within acceptable tolerance")
    else:
        print("\n[FAIL] Calibrations differ significantly!")
        if trans_diff >= 5.0:
            print(f"  Translation diff {trans_diff:.2f}mm exceeds 5mm threshold")
        if angle_diff >= 2.0:
            print(f"  Rotation diff {angle_diff:.2f} deg exceeds 2 deg threshold")


if __name__ == "__main__":
    main()