from argparse import ArgumentParser
from pathlib import Path

import cv2 as cv
import numpy as np


RIGHT_ANGLE_MODES = {
    "cw90": cv.ROTATE_90_CLOCKWISE,
    "ccw90": cv.ROTATE_90_COUNTERCLOCKWISE,
    "180": cv.ROTATE_180,
}


parser = ArgumentParser()
parser.add_argument("--input", default="input/scene.png", help="Source image")
parser.add_argument("--output", default="output/scene-rotated.png", help="Rotated output image")
parser.add_argument(
    "--mode",
    choices=sorted(RIGHT_ANGLE_MODES),
    default="cw90",
    help="Right-angle rotation mode when --angle is not used",
)
parser.add_argument(
    "--angle",
    type=float,
    help="Arbitrary rotation angle in degrees; positive values rotate counter-clockwise",
)
parser.add_argument("--scale", type=float, default=1.0, help="Scale for --angle rotations")
parser.add_argument(
    "--expand",
    action="store_true",
    help="Expand the output canvas for --angle rotations instead of keeping the source size",
)
args = parser.parse_args()

input_path = Path(args.input)
output_path = Path(args.output)

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

source_height, source_width = image.shape[:2]

if args.angle is None:
    rotated = cv.rotate(image, RIGHT_ANGLE_MODES[args.mode])
    method = f"cv.rotate mode={args.mode}"
    matrix_text = None
else:
    center = (source_width / 2, source_height / 2)
    matrix = cv.getRotationMatrix2D(center, args.angle, args.scale)
    output_size = (source_width, source_height)

    if args.expand:
        cos = abs(matrix[0, 0])
        sin = abs(matrix[0, 1])
        new_width = int((source_height * sin) + (source_width * cos))
        new_height = int((source_height * cos) + (source_width * sin))
        matrix[0, 2] += (new_width / 2) - center[0]
        matrix[1, 2] += (new_height / 2) - center[1]
        output_size = (new_width, new_height)

    rotated = cv.warpAffine(
        image,
        matrix,
        output_size,
        flags=cv.INTER_LINEAR,
        borderMode=cv.BORDER_REPLICATE,
    )
    method = f"cv.warpAffine angle={args.angle:g} scale={args.scale:g} expand={args.expand}"
    matrix_text = np.array2string(matrix, precision=3, suppress_small=True)

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

written = cv.imread(str(output_path), cv.IMREAD_COLOR)
if written is None:
    raise SystemExit(f"could not read written image: {output_path}")

output_height, output_width = written.shape[:2]

print(f"method: {method}")
print(f"source: {source_width}x{source_height} channels={image.shape[2]}")
if matrix_text is not None:
    print("matrix:")
    print(matrix_text)
print(f"output: {output_width}x{output_height} channels={written.shape[2]}")
print(f"wrote: {output_path}")
print(f"verified: {output_path}")