from argparse import ArgumentParser
from pathlib import Path

import cv2 as cv
import numpy as np


parser = ArgumentParser()
parser.add_argument("--input", required=True, help="Source image")
parser.add_argument("--output", required=True, help="Annotated output image")
parser.add_argument("--threshold", type=float, default=0.01)
parser.add_argument("--block-size", type=int, default=2)
parser.add_argument("--aperture-size", type=int, default=3)
parser.add_argument("--k", type=float, default=0.04)
args = parser.parse_args()

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

gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
response = cv.cornerHarris(
    np.float32(gray),
    args.block_size,
    args.aperture_size,
    args.k,
)
response = cv.dilate(response, None)

max_response = float(response.max())
threshold_value = args.threshold * max_response
corner_mask = response > threshold_value
corner_pixels = int(np.count_nonzero(corner_mask))

components, labels, stats, centroids = cv.connectedComponentsWithStats(
    corner_mask.astype("uint8"),
    connectivity=8,
)
corner_groups = components - 1

annotated = image.copy()
annotated[corner_mask] = (0, 0, 255)
for centroid in centroids[1:]:
    x, y = centroid
    cv.circle(annotated, (int(round(x)), int(round(y))), 4, (0, 255, 0), 1)

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

height, width = gray.shape
print(f"image: {width}x{height}")
print(f"block size: {args.block_size}")
print(f"aperture size: {args.aperture_size}")
print(f"k: {args.k:.2f}")
print(f"threshold: {args.threshold:.3f} of max response")
print(f"corner pixels: {corner_pixels}")
print(f"corner groups: {corner_groups}")
print(f"wrote: {output_path}")