"""
Filename: calibration.py
Description: Calibrates distortion from images

Author: Tyler de Zeeuw
License: GPL-3.0
"""

# Built-in imports
import os
import glob

# External library imports
import cv2
import numpy as np


# BOARD CONFIGURATION
square_len = 0.040
marker_len  = 0.030
board_cols  = 6
board_rows  = 4

aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)

ids_grid = {}
for marker_id in range(board_cols * board_rows // 2):
    row = marker_id // 3
    col = (marker_id % 3) * 2 + (0 if row % 2 == 0 else 1)
    ids_grid[(row, col)] = marker_id

ids = []
for row in range(board_rows):
    for col in range(board_cols):
        if (row + col) % 2 == 0:
            ids.append(ids_grid[(row, col)])
ids = np.array(ids, dtype=np.int32)

board = cv2.aruco.CharucoBoard(
    (board_cols, board_rows),
    square_len,
    marker_len,
    aruco_dict,
    ids
)
board.setLegacyPattern(True)

detector_params = cv2.aruco.DetectorParameters()
detector_params.cornerRefinementMethod = cv2.aruco.CORNER_REFINE_SUBPIX


# Calibration routine
def calibrate_camera(image_pattern, camera_label):
    images = sorted(glob.glob(image_pattern))
    if not images:
        print(f"[{camera_label}] No images found.")
        return None, None, None

    print(f"[{camera_label}] Found {len(images)} images.")

    objpoints = []
    imgpoints = []
    image_size = None

    charuco_params = cv2.aruco.CharucoParameters()
    detector = cv2.aruco.CharucoDetector(board, charuco_params, detector_params)

    for path in images:
        img = cv2.imread(path)
        if img is None:
            continue
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        h, w = gray.shape
        if image_size is None:
            image_size = (w, h)
        elif (w, h) != image_size:
            print(f"[{camera_label}] WARNING: {path} has size {w}x{h}, expected {image_size}. Skipping.")
            continue

        charuco_corners, charuco_ids, _, _ = detector.detectBoard(img)
        if charuco_corners is None or len(charuco_corners) < 6:
            continue

        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 60, 0.0001)
        pts = charuco_corners.reshape(-1, 2)
        if len(pts) >= 2:
            dists = np.linalg.norm(pts[:, None, :] - pts[None, :, :], axis=-1)
            np.fill_diagonal(dists, np.inf)
            min_dist = np.min(dists)
            win_size = max(3, int(min_dist / 5))
        else:
            win_size = 5
        charuco_corners = cv2.cornerSubPix(gray, charuco_corners, (win_size, win_size), (-1, -1), criteria)

        obj_pts, img_pts = board.matchImagePoints(charuco_corners, charuco_ids.flatten())
        if obj_pts is None or len(obj_pts) < 6:
            continue

        objpoints.append(obj_pts)
        imgpoints.append(img_pts)

    if not objpoints:
        print(f"[{camera_label}] No valid views collected.")
        return None, None, None

    print(f"[{camera_label}] Collected {len(objpoints)} valid views.")
    ret, mtx, dist, _, _ = cv2.calibrateCamera(objpoints, imgpoints, image_size, None, None)
    print(f"[{camera_label}] RMS reprojection error: {ret:.4f} px")
    print(f"    Camera matrix:\n{mtx}")
    print(f"    Distortion coefficients: {dist.ravel()}")

    return mtx, dist, image_size


# Pose estimation
def estimate_board_pose(image_path, K, D):
    img = cv2.imread(image_path)
    if img is None:
        return None, None
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    charuco_params = cv2.aruco.CharucoParameters()
    charuco_params.cameraMatrix = K
    charuco_params.distCoeffs = D
    detector = cv2.aruco.CharucoDetector(board, charuco_params, detector_params)

    charuco_corners, charuco_ids, _, _ = detector.detectBoard(img)
    if charuco_corners is None or len(charuco_corners) < 4:
        return None, None
    obj_pts, img_pts = board.matchImagePoints(charuco_corners, charuco_ids.flatten())
    if obj_pts is None or len(obj_pts) < 4:
        return None, None

    success, rvec, tvec = cv2.solvePnP(obj_pts, img_pts, K, D, flags=cv2.SOLVEPNP_ITERATIVE)
    if not success:
        return None, None
    R, _ = cv2.Rodrigues(rvec)
    return R, tvec.reshape(3,1)


# Rotation averaging helpers
def rotmat_to_quat(R):
    q = np.empty(4)
    trace = np.trace(R)
    if trace > 0:
        s = 0.5 / np.sqrt(trace + 1.0)
        q[0] = 0.25 / s
        q[1] = (R[2, 1] - R[1, 2]) * s
        q[2] = (R[0, 2] - R[2, 0]) * s
        q[3] = (R[1, 0] - R[0, 1]) * s
    elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
        s = 2.0 * np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2])
        q[0] = (R[2, 1] - R[1, 2]) / s
        q[1] = 0.25 * s
        q[2] = (R[0, 1] + R[1, 0]) / s
        q[3] = (R[0, 2] + R[2, 0]) / s
    elif R[1, 1] > R[2, 2]:
        s = 2.0 * np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2])
        q[0] = (R[0, 2] - R[2, 0]) / s
        q[1] = (R[0, 1] + R[1, 0]) / s
        q[2] = 0.25 * s
        q[3] = (R[1, 2] + R[2, 1]) / s
    else:
        s = 2.0 * np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1])
        q[0] = (R[1, 0] - R[0, 1]) / s
        q[1] = (R[0, 2] + R[2, 0]) / s
        q[2] = (R[1, 2] + R[2, 1]) / s
        q[3] = 0.25 * s
    return q  # [w, x, y, z]

def quat_to_rotmat(q):
    w, x, y, z = 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 average_rotations(rot_list):
    quats = [rotmat_to_quat(R) for R in rot_list]
    ref = quats[0]
    aligned = [q if np.dot(q, ref) >= 0 else -q for q in quats]
    mean_q = np.mean(aligned, axis=0)
    mean_q /= np.linalg.norm(mean_q)
    return quat_to_rotmat(mean_q)


def main():

    base_dir = os.path.dirname(os.path.abspath(__file__))

    # File paths
    #TODO: Unhardcode this or streamline it with the photo taker file
    cam1_pattern = os.path.join(base_dir, "Cam_1_fullres_*.jpg")
    cam2_pattern = os.path.join(base_dir, "Cam_2_fullres_*.jpg")
    phone_pattern = os.path.join(base_dir, "Cam_3_fullres_*.jpg")

    # 1. Calibrate each camera
    K1, D1, sz1 = calibrate_camera(cam1_pattern, "Camera 1")
    K2, D2, sz2 = calibrate_camera(cam2_pattern, "Camera 2")
    Kp, Dp, szp = calibrate_camera(phone_pattern, "Phone")

    if K1 is None or K2 is None or Kp is None:
        print("\nCalibration failed for one or more cameras. Exiting.")
        exit()

    # Save intrinsics
    np.save("camera1_matrix.npy", K1)
    np.save("dist1_coeffs.npy", D1)
    np.save("camera2_matrix.npy", K2)
    np.save("dist2_coeffs.npy", D2)
    np.save("phone_matrix.npy", Kp)
    np.save("phone_dist.npy", Dp)
    print("\nIntrinsics saved.")

    # 2. Estimate relative poses (Phone as world origin)
    R_cam1_list, t_cam1_list = [], []
    R_cam2_list, t_cam2_list = [], []

    #TODO: Is 20 calibration images really needed or can this be reduced without any loss?
    for i in range(1, 21):
        f1 = os.path.join(base_dir, f"Cam_1_fullres_{i}.jpg")
        f2 = os.path.join(base_dir, f"Cam_2_fullres_{i}.jpg")
        fp = os.path.join(base_dir, f"Cam_3_fullres_{i}.jpg")

        R1, t1 = estimate_board_pose(f1, K1, D1)
        R2, t2 = estimate_board_pose(f2, K2, D2)
        Rp, tp = estimate_board_pose(fp, Kp, Dp)

        if R1 is None or R2 is None or Rp is None:
            print(f"Frame {i}: board not seen by all three. Skipping.")
            continue

        # Cam1 -> Phone
        R_c1 = Rp @ R1.T
        t_c1 = tp - R_c1 @ t1
        R_cam1_list.append(R_c1)
        t_cam1_list.append(t_c1)

        # Cam2 -> Phone
        R_c2 = Rp @ R2.T
        t_c2 = tp - R_c2 @ t2
        R_cam2_list.append(R_c2)
        t_cam2_list.append(t_c2)

        print(f"Frame {i}: Cam1->Phone pos (mm): {t_c1.ravel()*1000}")

    if not R_cam1_list:
        print("\nNo common frames where all three cameras saw the board. Exiting.")
        exit()

    mean_R_cam1 = average_rotations(R_cam1_list)
    mean_t_cam1 = np.median(np.hstack(t_cam1_list), axis=1).reshape(3,1)

    mean_R_cam2 = average_rotations(R_cam2_list)
    mean_t_cam2 = np.median(np.hstack(t_cam2_list), axis=1).reshape(3,1)

    np.save("R_cam1_to_phone.npy", mean_R_cam1)
    np.save("t_cam1_to_phone.npy", mean_t_cam1)
    np.save("R_cam2_to_phone.npy", mean_R_cam2)
    np.save("t_cam2_to_phone.npy", mean_t_cam2)

    # Print final results
    print("\n=== FINAL EXTRINSICS (Phone as world origin) ===")
    cam1_pos = -mean_R_cam1.T @ mean_t_cam1
    cam2_pos = -mean_R_cam2.T @ mean_t_cam2
    print(f"Camera 1 -> Phone: t = [{mean_t_cam1[0,0]*1000:.1f}, {mean_t_cam1[1,0]*1000:.1f}, {mean_t_cam1[2,0]*1000:.1f}] mm")
    print(f"  => Camera 1 position (in Phone frame): [{cam1_pos[0,0]*1000:.1f}, {cam1_pos[1,0]*1000:.1f}, {cam1_pos[2,0]*1000:.1f}] mm")
    print(f"Camera 2 -> Phone: t = [{mean_t_cam2[0,0]*1000:.1f}, {mean_t_cam2[1,0]*1000:.1f}, {mean_t_cam2[2,0]*1000:.1f}] mm")
    print(f"  => Camera 2 position (in Phone frame): [{cam2_pos[0,0]*1000:.1f}, {cam2_pos[1,0]*1000:.1f}, {cam2_pos[2,0]*1000:.1f}] mm")

    # Orthogonality checks
    print("\nOrthogonality check R_cam1_to_phone @ R.T:")
    print(mean_R_cam1 @ mean_R_cam1.T)
    print("det(R) =", np.linalg.det(mean_R_cam1))
    print("Orthogonality check R_cam2_to_phone @ R.T:")
    print(mean_R_cam2 @ mean_R_cam2.T)
    print("det(R) =", np.linalg.det(mean_R_cam2))


if __name__ == "__main__":
    main()