#!/usr/bin/env python3
"""
charuco_calibration.py — ChArUco Camera Calibration
OpenCV >= 4.7, Python >= 3.8
"""
import numpy as np
import cv2
import glob
import argparse
import json
import os
import sys
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(description="ChArUco Camera Calibration")
    parser.add_argument("-i", "--images", type=str, required=True,
                        help="Path to calibration images (dir or glob)")
    parser.add_argument("--squares_x", type=int, default=7,
                        help="Number of chessboard squares in X (default: 7)")
    parser.add_argument("--squares_y", type=int, default=5,
                        help="Number of chessboard squares in Y (default: 5)")
    parser.add_argument("--square_length", type=float, required=True,
                        help="Square side length in mm")
    parser.add_argument("--marker_length", type=float, default=None,
                        help="ArUco marker side length in mm (default: 0.6 * square_length)")
    parser.add_argument("--dictionary", type=str, default="DICT_6X6_250",
                        help="ArUco dictionary (default: DICT_6X6_250)")
    parser.add_argument("-o", "--out", type=str, default="calibration_results",
                        help="Output basename (default: calibration_results)")
    parser.add_argument("--fix_principal_point", action="store_true",
                        help="Fix principal point at image centre")
    parser.add_argument("--sample", type=str, default=None,
                        help="Sample image to undistort (optional)")
    parser.add_argument("--stereo", action="store_true",
                        help="Stereo calibration (requires left/right subdirs in --images)")
    return parser.parse_args()


def get_aruco_dict(name: str):
    """Resolve ArUco dictionary name to OpenCV object."""
    dict_map = {
        "DICT_4X4_50": cv2.aruco.DICT_4X4_50,
        "DICT_4X4_100": cv2.aruco.DICT_4X4_100,
        "DICT_4X4_250": cv2.aruco.DICT_4X4_250,
        "DICT_4X4_1000": cv2.aruco.DICT_4X4_1000,
        "DICT_5X5_50": cv2.aruco.DICT_5X5_50,
        "DICT_5X5_100": cv2.aruco.DICT_5X5_100,
        "DICT_5X5_250": cv2.aruco.DICT_5X5_250,
        "DICT_5X5_1000": cv2.aruco.DICT_5X5_1000,
        "DICT_6X6_50": cv2.aruco.DICT_6X6_50,
        "DICT_6X6_100": cv2.aruco.DICT_6X6_100,
        "DICT_6X6_250": cv2.aruco.DICT_6X6_250,
        "DICT_6X6_1000": cv2.aruco.DICT_6X6_1000,
        "DICT_7X7_50": cv2.aruco.DICT_7X7_50,
        "DICT_7X7_100": cv2.aruco.DICT_7X7_100,
        "DICT_7X7_250": cv2.aruco.DICT_7X7_250,
        "DICT_7X7_1000": cv2.aruco.DICT_7X7_1000,
        "DICT_ARUCO_ORIGINAL": cv2.aruco.DICT_ARUCO_ORIGINAL,
    }
    if name not in dict_map:
        print(f"Unknown dictionary {name}. Using DICT_6X6_250.")
        name = "DICT_6X6_250"
    return cv2.aruco.getPredefinedDictionary(dict_map[name])


def load_images(path_pattern: str):
    """Load images from directory or glob pattern."""
    path = Path(path_pattern)
    if path.is_dir():
        patterns = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.tiff"]
        files = []
        for p in patterns:
            files.extend(sorted(path.glob(p)))
            files.extend(sorted(path.glob(p.upper())))
    else:
        files = sorted(Path(".").parent.glob(path_pattern))
    if not files:
        print(f"No images found at: {path_pattern}")
        sys.exit(1)
    print(f"Found {len(files)} images")
    return files


def detect_charuco_board(
    images, dictionary, board, squares_x, squares_y
):
    """Detect ChArUco corners in all images."""
    all_corners = []
    all_ids = []
    valid_indices = []
    detector_params = cv2.aruco.DetectorParameters()

    for idx, img_path in enumerate(images):
        img = cv2.imread(str(img_path))
        if img is None:
            print(f"  [SKIP] Cannot read: {img_path.name}")
            continue

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        marker_corners, marker_ids, _ = cv2.aruco.detectMarkers(
            gray, dictionary, parameters=detector_params
        )

        if marker_ids is None or len(marker_ids) < 4:
            print(f"  [SKIP] {img_path.name}: only {0 if marker_ids is None else len(marker_ids)} markers")
            continue

        # Refine marker detection
        cv2.aruco.refineDetectedMarkers(
            gray, board, marker_corners, marker_ids
        )

        # Interpolate chessboard corners from markers
        charuco_corners, charuco_ids = cv2.aruco.interpolateCornersCharuco(
            marker_corners, marker_ids, gray, board
        )

        if charuco_ids is None or len(charuco_ids) < 4:
            print(f"  [SKIP] {img_path.name}: too few ChArUco corners ({0 if charuco_ids is None else len(charuco_ids)})")
            continue

        all_corners.append(charuco_corners)
        all_ids.append(charuco_ids)
        valid_indices.append(idx)
        print(f"  [OK]   {img_path.name}: {len(charuco_ids)} corners, {len(marker_ids)} markers")

    if len(all_corners) < 5:
        print(f"\nERROR: Only {len(all_corners)} valid views (need >= 5). Add more images.")
        sys.exit(1)

    return all_corners, all_ids, valid_indices


def run_calibration(
    all_corners, all_ids, board, image_size, fix_principal_point
):
    """Run ChArUco calibration."""
    flags = 0
    if fix_principal_point:
        flags |= cv2.CALIB_FIX_PRINCIPAL_POINT

    print(f"\nRunning calibration with {len(all_corners)} views...")
    ret, mtx, dist, rvecs, tvecs = cv2.aruco.calibrateCameraCharuco(
        all_corners, all_ids, board, image_size, None, None, flags=flags
    )
    print(f"RMS reprojection error: {ret:.4f} pixels\n")
    return ret, mtx, dist, rvecs, tvecs


def save_results_json(ret, mtx, dist, rvecs, tvecs, image_size, out_path):
    """Save calibration results as JSON."""
    data = {
        "rms_error": float(ret),
        "image_size": list(image_size),
        "camera_matrix": mtx.tolist(),
        "distortion_coefficients": dist.tolist(),
        "per_view_errors": [],
    }
    with open(out_path, "w") as f:
        json.dump(data, f, indent=2)
    print(f"Saved: {out_path}")


def save_results_xml(ret, mtx, dist, image_size, out_path):
    """Save calibration results as OpenCV XML."""
    fs = cv2.FileStorage(out_path, cv2.FILE_STORAGE_WRITE)
    fs.write("rms_error", ret)
    fs.write("image_width", image_size[0])
    fs.write("image_height", image_size[1])
    fs.write("camera_matrix", mtx)
    fs.write("distortion_coefficients", dist)
    fs.release()
    print(f"Saved: {out_path}")


def compute_per_view_errors(all_corners, all_ids, rvecs, tvecs, mtx, dist, board):
    """Compute reprojection error for each view."""
    errors = []
    for i in range(len(all_corners)):
        img_points = all_corners[i].reshape(-1, 2)
        obj_points = board.matchImagePoints(all_corners[i], all_ids[i])

        proj_points, _ = cv2.projectPoints(
            obj_points, rvecs[i], tvecs[i], mtx, dist
        )
        proj_points = proj_points.reshape(-1, 2)
        error = np.sqrt(np.mean(np.sum((img_points - proj_points) ** 2, axis=1)))
        errors.append(float(error))
    return errors


def visualize_results(
    images, all_corners, all_ids, valid_indices, mtx, dist,
    valid_calib_indices, rvecs, tvecs, board, out_dir
):
    """Draw detected corners and save visualisation images."""
    viz_dir = out_dir / "visualisations"
    viz_dir.mkdir(parents=True, exist_ok=True)

    for i, img_idx in enumerate(valid_indices):
        if i >= len(valid_calib_indices):
            break
        img = cv2.imread(str(images[img_idx]))
        if img is None:
            continue

        # Draw detected ChArUco corners
        cv2.aruco.drawDetectedCornersCharuco(
            img, all_corners[i], all_ids[i]
        )

        # Draw axes
        cv2.drawFrameAxes(
            img, mtx, dist, rvecs[i], tvecs[i], 0.03
        )

        out_file = viz_dir / f"frame_{img_idx:03d}.jpg"
        cv2.imwrite(str(out_file), img)

    print(f"Visualisations saved to: {viz_dir}/")


def undistort_sample(sample_path, mtx, dist, out_dir):
    """Undistort a sample image."""
    img = cv2.imread(sample_path)
    if img is None:
        print(f"Cannot read sample image: {sample_path}")
        return

    h, w = img.shape[:2]
    new_mtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))

    # Method 1: cv.undistort
    undistorted = cv2.undistort(img, mtx, dist, None, new_mtx)
    x, y, w_roi, h_roi = roi
    if w_roi > 0 and h_roi > 0:
        undistorted = undistorted[y : y + h_roi, x : x + w_roi]

    out_path = out_dir / "undistorted_sample.jpg"
    cv2.imwrite(str(out_path), undistorted)

    # Method 2: remap (for video)
    mapx, mapy = cv2.initUndistortRectifyMap(
        mtx, dist, None, new_mtx, (w, h), cv2.CV_32FC1
    )
    remapped = cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR)
    remap_path = out_dir / "undistorted_remap_sample.jpg"
    cv2.imwrite(str(remap_path), remapped)

    print(f"Undistorted samples saved to: {out_dir}/")
    print(f"New camera matrix (with alpha=1):\n{new_mtx}")


def stereo_calibrate(images_dir: Path, board, square_length, dictionary, out_dir):
    """Run stereo ChArUco calibration."""
    left_dir = images_dir / "left"
    right_dir = images_dir / "right"

    if not left_dir.is_dir() or not right_dir.is_dir():
        print("ERROR: Stereo mode requires 'left/' and 'right/' subdirectories.")
        sys.exit(1)

    left_images = load_images(str(left_dir))
    right_images = load_images(str(right_dir))

    print(f"\nStereo calibration with {len(left_images)} left / {len(right_images)} right images")

    # Detect separately
    left_corners, left_ids, left_idx = detect_charuco_board(
        left_images, dictionary, board, board.getChessboardSize()[0], board.getChessboardSize()[1]
    )
    right_corners, right_ids, right_idx = detect_charuco_board(
        right_images, dictionary, board, board.getChessboardSize()[0], board.getChessboardSize()[1]
    )

    # Match pairs by index
    paired_indices = set(left_idx) & set(right_idx)
    if len(paired_indices) < 3:
        print(f"ERROR: Only {len(paired_indices)} matching stereo pairs (need >= 3).")
        sys.exit(1)

    paired_indices = sorted(paired_indices)
    left_matched = [left_corners[left_idx.index(i)] for i in paired_indices]
    left_ids_matched = [left_ids[left_idx.index(i)] for i in paired_indices]
    right_matched = [right_corners[right_idx.index(i)] for i in paired_indices]
    right_ids_matched = [right_ids[right_idx.index(i)] for i in paired_indices]

    # Get object points
    obj_points = []
    for i in range(len(paired_indices)):
        _, obj_pts, _ = board.matchImagePoints(left_corners_flat(left_matched[i], left_ids_matched[i]))
        obj_points.append(obj_pts)

    # Stereo calibration
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 1e-5)
    flags = cv2.CALIB_FIX_INTRINSIC
    ret, _, _, _, _, R, T, E, F = cv2.stereoCalibrate(
        obj_points,
        left_matched,
        right_matched,
        None, None, None, None,
        (0, 0),
        criteria=criteria,
        flags=flags,
    )

    print(f"\nStereo RMS: {ret:.4f} pixels")
    print(f"Rotation matrix:\n{R}")
    print(f"Translation vector:\n{T}")

    # Stereo rectify
    R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(
        None, None, None, None, (0, 0), R, T, alpha=0
    )

    # Save stereo results
    stereo_out = out_dir / "stereo_calibration.json"
    data = {
        "rms_error": float(ret),
        "rotation_matrix": R.tolist(),
        "translation_vector": T.tolist(),
        "essential_matrix": E.tolist(),
        "fundamental_matrix": F.tolist(),
    }
    with open(stereo_out, "w") as f:
        json.dump(data, f, indent=2)
    print(f"\nStereo results saved to: {stereo_out}")

    return ret, R, T, E, F


def left_corners_flat(corners, ids):
    """Helper to flatten charuco corners and return (corners_flat, ids_flat, obj_points)."""
    return corners, ids


def main():
    args = parse_args()

    # Resolve paths
    images_dir = Path(args.images)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Board parameters
    squares_x = args.squares_x
    squares_y = args.squares_y
    square_length = args.square_length / 1000.0  # convert mm to metres
    marker_length = args.marker_length
    if marker_length is None:
        marker_length = square_length * 0.6
    else:
        marker_length = marker_length / 1000.0

    # Create dictionary and board
    dictionary = get_aruco_dict(args.dictionary)
    board = cv2.aruco.CharucoBoard(
        (squares_x, squares_y),
        square_length,
        marker_length,
        dictionary,
    )

    print(f"ChArUco board: {squares_x}x{squares_y}, "
          f"square={args.square_length}mm, marker={args.square_length * (0.6 if args.marker_length is None else args.marker_length / args.square_length):.1f}mm")
    print(f"Dictionary: {args.dictionary}")
    print()

    if args.stereo:
        stereo_calibrate(images_dir, board, args.square_length, dictionary, out_dir)
        return

    # Load and detect
    images = load_images(str(images_dir))
    all_corners, all_ids, valid_indices = detect_charuco_board(
        images, dictionary, board, squares_x, squares_y
    )

    # Get image size from first valid image
    img = cv2.imread(str(images[valid_indices[0]]))
    image_size = (img.shape[1], img.shape[0])

    # Run calibration
    ret, mtx, dist, rvecs, tvecs = run_calibration(
        all_corners, all_ids, board, image_size, args.fix_principal_point
    )

    # Per-view errors
    errors = compute_per_view_errors(
        all_corners, all_ids, rvecs, tvecs, mtx, dist, board
    )
    print("Per-view reprojection errors:")
    for i, (idx, err) in enumerate(zip(valid_indices, errors)):
        marker = "!" if err > 0.8 else ""
        print(f"  [{i:2d}] {images[idx].name}: {err:.3f} px {marker}")
    print()

    # Camera matrix summary
    print(f"Camera matrix:\n{mtx}")
    print(f"\nDistortion coefficients: {dist.ravel()}")
    fx = mtx[0, 0]
    fy = mtx[1, 1]
    cx = mtx[0, 2]
    cy = mtx[1, 2]
    print(f"\nFocal length: fx={fx:.1f}, fy={fy:.1f}")
    print(f"Principal point: cx={cx:.1f}, cy={cy:.1f}")
    print(f"Image size: {image_size[0]}x{image_size[1]}")

    # Sanity checks
    print(f"\n--- Sanity Checks ---")
    if fx > image_size[0] * 0.5 and fx < image_size[0] * 5:
        print("[OK] Focal length X looks reasonable")
    else:
        print("[WARN] Focal length X seems unusual")
    if abs(cx - image_size[0] / 2) < image_size[0] * 0.2:
        print("[OK] Principal point X is near centre")
    else:
        print("[WARN] Principal point X is far from centre")
    if abs(cy - image_size[1] / 2) < image_size[1] * 0.2:
        print("[OK] Principal point Y is near centre")
    else:
        print("[WARN] Principal point Y is far from centre")
    if ret < 0.5:
        print("[OK] RMS error is excellent (< 0.5 px)")
    elif ret < 0.8:
        print("[OK] RMS error is acceptable (< 0.8 px)")
    else:
        print("[WARN] RMS error is high (>= 0.8 px) — consider adding/improving images")
    print()

    # Save
    json_path = out_dir / f"{args.out}.json"
    xml_path = out_dir / f"{args.out}.xml"
    save_results_json(ret, mtx, dist, rvecs, tvecs, image_size, json_path)
    save_results_xml(ret, mtx, dist, image_size, xml_path)

    # Visualise
    visualize_results(
        images, all_corners, all_ids, valid_indices,
        mtx, dist, list(range(len(all_corners))),
        rvecs, tvecs, board, out_dir
    )

    # Undistort sample
    if args.sample:
        undistort_sample(args.sample, mtx, dist, out_dir)

    print(f"\nAll results saved to: {out_dir.resolve()}")


if __name__ == "__main__":
    main()