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

import cv2


METHODS = {
    "ccoeff-normed": cv2.TM_CCOEFF_NORMED,
    "ccorr-normed": cv2.TM_CCORR_NORMED,
    "sqdiff-normed": cv2.TM_SQDIFF_NORMED,
}


parser = argparse.ArgumentParser(description="Locate a template image inside a larger image with OpenCV.")
parser.add_argument("scene", type=Path)
parser.add_argument("template", type=Path)
parser.add_argument("output", type=Path)
parser.add_argument("--method", choices=METHODS, default="ccoeff-normed")
args = parser.parse_args()

scene_color = cv2.imread(str(args.scene), cv2.IMREAD_COLOR)
scene_gray = cv2.imread(str(args.scene), cv2.IMREAD_GRAYSCALE)
template_gray = cv2.imread(str(args.template), cv2.IMREAD_GRAYSCALE)

if scene_color is None or scene_gray is None:
    raise SystemExit(f"could not read scene image: {args.scene}")
if template_gray is None:
    raise SystemExit(f"could not read template image: {args.template}")
if template_gray.shape[0] > scene_gray.shape[0] or template_gray.shape[1] > scene_gray.shape[1]:
    raise SystemExit("template image must not be larger than the scene image")

method = METHODS[args.method]
result = cv2.matchTemplate(scene_gray, template_gray, method)
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)

if method == cv2.TM_SQDIFF_NORMED:
    top_left = min_loc
    score = min_val
    better = "lower"
else:
    top_left = max_loc
    score = max_val
    better = "higher"

height, width = template_gray.shape
bottom_right = (top_left[0] + width, top_left[1] + height)
cv2.rectangle(scene_color, top_left, bottom_right, (0, 0, 255), 3)
cv2.putText(
    scene_color,
    f"{args.method} {score:.3f}",
    (top_left[0], max(25, top_left[1] - 10)),
    cv2.FONT_HERSHEY_SIMPLEX,
    0.7,
    (0, 0, 255),
    2,
)

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

print(f"method: {args.method} ({better} score is better)")
print(f"scene size: {scene_gray.shape[1]}x{scene_gray.shape[0]}")
print(f"template size: {width}x{height}")
print(f"result map: {result.shape[1]}x{result.shape[0]}")
print(f"best score: {score:.3f}")
print(f"top-left: x={top_left[0]}, y={top_left[1]}")
print(f"bottom-right: x={bottom_right[0]}, y={bottom_right[1]}")
print(f"output: {args.output}")