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

import cv2
import numpy as np


def odd_positive(value: str) -> int:
    size = int(value)
    if size < 1 or size % 2 == 0:
        raise argparse.ArgumentTypeError("kernel size must be a positive odd integer")
    return size


parser = argparse.ArgumentParser(description="Apply an OpenCV morphology operation to a binary mask.")
parser.add_argument("input_mask", type=Path)
parser.add_argument("output_mask", type=Path)
parser.add_argument("--operation", choices=("erode", "dilate", "open", "close"), default="open")
parser.add_argument("--kernel-size", type=odd_positive, default=5)
parser.add_argument("--kernel-shape", choices=("rect", "ellipse", "cross"), default="ellipse")
parser.add_argument("--iterations", type=int, default=1)
args = parser.parse_args()

gray = cv2.imread(str(args.input_mask), cv2.IMREAD_GRAYSCALE)
if gray is None:
    raise SystemExit(f"could not read mask: {args.input_mask}")

binary = np.where(gray > 0, 255, 0).astype(np.uint8)

shape_map = {
    "rect": cv2.MORPH_RECT,
    "ellipse": cv2.MORPH_ELLIPSE,
    "cross": cv2.MORPH_CROSS,
}
kernel = cv2.getStructuringElement(
    shape_map[args.kernel_shape],
    (args.kernel_size, args.kernel_size),
)

if args.operation == "erode":
    result = cv2.erode(binary, kernel, iterations=args.iterations)
elif args.operation == "dilate":
    result = cv2.dilate(binary, kernel, iterations=args.iterations)
else:
    operation_map = {
        "open": cv2.MORPH_OPEN,
        "close": cv2.MORPH_CLOSE,
    }
    result = cv2.morphologyEx(
        binary,
        operation_map[args.operation],
        kernel,
        iterations=args.iterations,
    )

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

foreground_before = int(np.count_nonzero(binary == 255))
foreground_after = int(np.count_nonzero(result == 255))
changed_pixels = int(np.count_nonzero(binary != result))
total_pixels = int(binary.size)
unique_values = " ".join(str(int(value)) for value in np.unique(result))

print(f"operation: {args.operation}")
print(f"kernel: {args.kernel_shape} {args.kernel_size}x{args.kernel_size}")
print(f"iterations: {args.iterations}")
print(f"foreground before: {foreground_before} / {total_pixels} ({foreground_before / total_pixels * 100:.2f}%)")
print(f"foreground after: {foreground_after} / {total_pixels} ({foreground_after / total_pixels * 100:.2f}%)")
print(f"changed pixels: {changed_pixels}")
print(f"unique values: {unique_values}")
print(f"output: {args.output_mask}")