import cv2
import numpy as np
import glob
import os
import rawpy
from concurrent.futures import ThreadPoolExecutor, as_completed

# ============================================================
# CONFIG
# ============================================================

CHECKERBOARD = (10, 7)          # inner corners
SQUARE_SIZE = 25.0 / 1000.0    # 25 mm -> meters

IMAGE_PATTERNS = [
    "XPro2-12mm56-1mFokus_f8/*.JPG",
    "XPro2-12mm56-1mFokus_f8/*.RAF"
]


# Reject blurry images
MIN_SHARPNESS = 380.0

# Downscale for faster detection
MAX_WIDTH = 2000

# ============================================================
# PREPARE OBJECT POINTS
# ============================================================

objp = np.zeros((CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)
objp[:, :2] = np.mgrid[
    0:CHECKERBOARD[0],
    0:CHECKERBOARD[1]
].T.reshape(-1, 2)

objp *= SQUARE_SIZE


# ============================================================
# IMAGE PROCESSING
# ============================================================

blurry_images = []
failed_images = []

def process_image(fname):

    try:
        print(f"\nProcessing: {fname}")

        ext = os.path.splitext(fname)[1].lower()

        if ext == ".raf" or ext == ".RAF":

            with rawpy.imread(fname) as raw:

                rgb = raw.postprocess(
                    use_camera_wb=True,
                    no_auto_bright=True,
                    output_bps=8
                )

            img = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

        else:
            img = cv2.imread(fname)

        if img is None:
            print("  -> Could not read image")
            return None

        #gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        gray = img[:, :, 2]

        original_size = gray.shape[::-1]

        # ----------------------------------------------------
        # DOWNSCALE LARGE IMAGES
        # ----------------------------------------------------

        scale = 1.0

        if gray.shape[1] > MAX_WIDTH:
            scale = MAX_WIDTH / gray.shape[1]

            gray_small = cv2.resize(
                gray,
                None,
                fx=scale,
                fy=scale,
                interpolation=cv2.INTER_AREA
            )
        else:
            gray_small = gray

        # ----------------------------------------------------
        # SHARPNESS CHECK
        # ----------------------------------------------------

        sharpness = cv2.Laplacian(gray_small, cv2.CV_64F).var()

        print(f"  Sharpness: {sharpness:.1f}")

        if sharpness < MIN_SHARPNESS:
            print("  -> Rejected (too blurry)")
            failed_images.append(fname)
            return None

        # ----------------------------------------------------
        # FIND CHESSBOARD
        # ----------------------------------------------------

        flags = (
            cv2.CALIB_CB_NORMALIZE_IMAGE |
            cv2.CALIB_CB_EXHAUSTIVE
        )

        # Newer and more robust detector
        ret, corners = cv2.findChessboardCornersSB(
            gray_small,
            CHECKERBOARD,
            flags
        )

        if not ret:
            print("  -> No corners found")
            return None

        # ----------------------------------------------------
        # SCALE CORNERS BACK
        # ----------------------------------------------------

        if scale != 1.0:
            corners /= scale

        # ----------------------------------------------------
        # SUBPIX REFINEMENT
        # ----------------------------------------------------

        corners2 = cv2.cornerSubPix(
            gray,
            corners.astype(np.float32),
            (11, 11),
            (-1, -1),
            (
                cv2.TERM_CRITERIA_EPS +
                cv2.TERM_CRITERIA_MAX_ITER,
                30,
                1e-6
            )
        )

        print("  -> Success")

        return {
            "objpoints": objp,
            "imgpoints": corners2,
            "img_size": original_size
        }

    except Exception as e:
        print(f"  -> ERROR: {e}")
        return None


# ============================================================
# LOAD IMAGES
# ============================================================

images = []

for pattern in IMAGE_PATTERNS:
    images.extend(glob.glob(pattern))

print(f"Found images: {len(images)}")

if len(images) == 0:
    raise RuntimeError("No images found.")

# ============================================================
# PROCESS IMAGES IN PARALLEL
# ============================================================

results = []

with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:

    futures = [
        executor.submit(process_image, fname)
        for fname in images
    ]

    for future in as_completed(futures):

        result = future.result()

        if result is not None:
            results.append(result)

# ============================================================
# COLLECT VALID RESULTS
# ============================================================

if len(results) == 0:
    raise RuntimeError("No valid checkerboards detected.")

objpoints = [r["objpoints"] for r in results]
imgpoints = [r["imgpoints"] for r in results]
img_size = results[0]["img_size"]

print(f"\nValid images: {len(results)} / {len(images)}")

# ============================================================
# CALIBRATION
# ============================================================

print("\nRunning calibration...")

ret, K, D, rvecs, tvecs = cv2.calibrateCamera(
    objpoints,
    imgpoints,
    img_size,
    None,
    None
)

# ============================================================
# RESULTS
# ============================================================

print("\n=== Calibration Results ===")

print(f"RMS reprojection error: {ret:.4f}")

print("\nCamera Matrix:")
print(K)

print("\nDistortion Coefficients:")
print(D)

# ============================================================
# SAVE
# ============================================================

output_file = "calibration_XPro2-12mm56-1mFokus_f8.npz"

np.savez(
    output_file,
    camera_matrix=K,
    dist_coeffs=D
)

print(f"\nSaved calibration to: {output_file}")


# ------------------------------------------------------------
# DELETE COMMANDS
# ------------------------------------------------------------


if blurry_images:

    print("\nDelete blurry images:")

    delete_cmd = "rm " + " ".join(
        f'"{f}"' for f in blurry_images
    )

    print(delete_cmd)

if failed_images:

    print("\nDelete failed corner images:")

    delete_cmd = "rm " + " ".join(
        f'"{f}"' for f in failed_images
    )

    print(delete_cmd)