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-dir", default="output", help="Directory for converted images")
args = parser.parse_args()

input_path = Path(args.input)
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)

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

rgb = cv.cvtColor(image, cv.COLOR_BGR2RGB)
gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
hsv = cv.cvtColor(image, cv.COLOR_BGR2HSV)

height, width = image.shape[:2]
sample_y, sample_x = np.unravel_index(np.argmax(hsv[:, :, 1]), hsv[:, :, 1].shape)


def write_image(path, data):
    if not cv.imwrite(str(path), data):
        raise SystemExit(f"could not write output image: {path}")


rgb_path = output_dir / "scene-rgb-display.png"
gray_path = output_dir / "scene-gray.png"
hue_path = output_dir / "scene-hsv-hue.png"
preview_path = output_dir / "color-space-preview.png"

write_image(rgb_path, cv.cvtColor(rgb, cv.COLOR_RGB2BGR))
write_image(gray_path, gray)
write_image(hue_path, hsv[:, :, 0])

hue_display = cv.convertScaleAbs(hsv[:, :, 0], alpha=255 / 179)
hue_display = cv.applyColorMap(hue_display, cv.COLORMAP_HSV)
gray_display = cv.cvtColor(gray, cv.COLOR_GRAY2BGR)


def labeled_tile(tile, label):
    resized = cv.resize(tile, (360, 240), interpolation=cv.INTER_AREA)
    cv.rectangle(resized, (0, 0), (360, 34), (0, 0, 0), thickness=-1)
    cv.putText(
        resized,
        label,
        (12, 24),
        cv.FONT_HERSHEY_SIMPLEX,
        0.7,
        (255, 255, 255),
        2,
        cv.LINE_AA,
    )
    return resized


preview = cv.vconcat(
    [
        cv.hconcat(
            [
                labeled_tile(image, "BGR input"),
                labeled_tile(cv.cvtColor(rgb, cv.COLOR_RGB2BGR), "RGB display"),
            ]
        ),
        cv.hconcat(
            [
                labeled_tile(gray_display, "Grayscale"),
                labeled_tile(hue_display, "HSV hue"),
            ]
        ),
    ]
)
write_image(preview_path, preview)

print(f"input: {width}x{height} channels={image.shape[2]}")
print(f"sample pixel: x={sample_x} y={sample_y}")
print(f"sample BGR: {image[sample_y, sample_x].tolist()}")
print(f"sample RGB: {rgb[sample_y, sample_x].tolist()}")
print(f"sample gray: {int(gray[sample_y, sample_x])}")
print(f"sample HSV: {hsv[sample_y, sample_x].tolist()}")
print(f"wrote: {rgb_path}")
print(f"wrote: {gray_path}")
print(f"wrote: {hue_path}")
print(f"wrote: {preview_path}")