from argparse import ArgumentParser
from pathlib import Path

import cv2
import numpy as np


parser = ArgumentParser()
parser.add_argument("--calibration", required=True, help="Path to calibration.npz")
parser.add_argument("--input", required=True, help="Distorted input image")
parser.add_argument("--output", required=True, help="Undistorted output image")
parser.add_argument(
    "--alpha",
    type=float,
    default=0.0,
    help="0 crops invalid pixels; 1 keeps more of the original field of view",
)
args = parser.parse_args()

data = np.load(args.calibration)
camera_matrix = data["camera_matrix"]
dist_coeffs = data["dist_coeffs"]

image = cv2.imread(args.input)
if image is None:
    raise SystemExit(f"could not read input image: {args.input}")

height, width = image.shape[:2]
new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(
    camera_matrix,
    dist_coeffs,
    (width, height),
    args.alpha,
    (width, height),
)

undistorted = cv2.undistort(image, camera_matrix, dist_coeffs, None, new_camera_matrix)
x, y, roi_width, roi_height = roi
if roi_width > 0 and roi_height > 0:
    undistorted = undistorted[y : y + roi_height, x : x + roi_width]

output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
if not cv2.imwrite(str(output_path), undistorted):
    raise SystemExit(f"could not write output image: {output_path}")

print(f"input: {width}x{height}")
print(f"camera matrix: {camera_matrix.shape[0]}x{camera_matrix.shape[1]}")
print(f"distortion coefficients: {dist_coeffs.size}")
print(f"roi: x={x} y={y} width={roi_width} height={roi_height}")
print(f"wrote: {output_path} ({undistorted.shape[1]}x{undistorted.shape[0]})")