#!/usr/bin/env python3
import argparse
from pathlib import Path

import cv2


STATUS_NAMES = {
    0: "OK",
    1: "ERR_NEED_MORE_IMGS",
    2: "ERR_HOMOGRAPHY_EST_FAIL",
    3: "ERR_CAMERA_PARAMS_ADJUST_FAIL",
}


def read_image(path):
    image = cv2.imread(str(path))
    if image is None:
        raise SystemExit(f"cannot_read: {path}")
    return image


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("images", nargs="+", type=Path)
    parser.add_argument("--output", default=Path("output/panorama.png"), type=Path)
    args = parser.parse_args()

    if len(args.images) < 2:
        raise SystemExit("need_at_least_two_images")

    images = [read_image(path) for path in args.images]

    stitcher = cv2.Stitcher_create(cv2.Stitcher_PANORAMA)
    status, panorama = stitcher.stitch(images)

    print(f"input_images: {len(images)}")
    print(f"stitcher_status: {status} ({STATUS_NAMES.get(status, 'UNKNOWN')})")

    if status != cv2.Stitcher_OK:
        raise SystemExit("stitch_failed")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    if not cv2.imwrite(str(args.output), panorama):
        raise SystemExit(f"cannot_write: {args.output}")

    height, width = panorama.shape[:2]
    print(f"saved_panorama: {args.output}")
    print(f"output_shape: {width}x{height}")


if __name__ == "__main__":
    main()