from argparse import ArgumentParser
from pathlib import Path

import cv2 as cv


parser = ArgumentParser()
parser.add_argument("--query", default="input/scene-left.png", help="First image to match")
parser.add_argument("--train", default="input/scene-right.png", help="Second image to match")
parser.add_argument("--output", default="output/orb-matches.png", help="Path for the match visualization")
parser.add_argument("--max-features", type=int, default=1000)
parser.add_argument("--keep", type=int, default=30, help="Number of best matches to draw")
args = parser.parse_args()

query_path = Path(args.query)
train_path = Path(args.train)
output_path = Path(args.output)

query_gray = cv.imread(str(query_path), cv.IMREAD_GRAYSCALE)
train_gray = cv.imread(str(train_path), cv.IMREAD_GRAYSCALE)
query_color = cv.imread(str(query_path), cv.IMREAD_COLOR)
train_color = cv.imread(str(train_path), cv.IMREAD_COLOR)

if query_gray is None:
    raise SystemExit(f"could not read query image: {query_path}")
if train_gray is None:
    raise SystemExit(f"could not read train image: {train_path}")

orb = cv.ORB_create(nfeatures=args.max_features)
query_keypoints, query_descriptors = orb.detectAndCompute(query_gray, None)
train_keypoints, train_descriptors = orb.detectAndCompute(train_gray, None)

if query_descriptors is None:
    raise SystemExit(f"no ORB descriptors found in {query_path}")
if train_descriptors is None:
    raise SystemExit(f"no ORB descriptors found in {train_path}")

matcher = cv.BFMatcher(cv.NORM_HAMMING, crossCheck=True)
matches = matcher.match(query_descriptors, train_descriptors)
matches = sorted(matches, key=lambda match: match.distance)
kept_matches = matches[: args.keep]

if not kept_matches:
    raise SystemExit("no cross-checked ORB matches found")

match_view = cv.drawMatches(
    query_color,
    query_keypoints,
    train_color,
    train_keypoints,
    kept_matches,
    None,
    flags=cv.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS,
)

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

print(f"query keypoints: {len(query_keypoints)}")
print(f"train keypoints: {len(train_keypoints)}")
print(f"cross-checked matches: {len(matches)}")
print(f"drawn matches: {len(kept_matches)}")
print(f"best distance: {kept_matches[0].distance:.1f}")
print(f"wrote: {output_path}")