#!/usr/bin/env python3
"""
Usage:
    python register.py base.jpg test.jpg
Controls:
    • Drag  : translate (offset)
    • + / - : zoom in / out
    • ← / → : rotate CCW / CW
    • r     : reset transform
    • q     : quit and print scale, rotation, offset
"""
import argparse
import math
import sys
from pathlib import Path

import cv2
import numpy as np

# ─── argument parsing ────────────────────────────────────────────────────────────
parser = argparse.ArgumentParser(description="Interactive image registration")
parser.add_argument("base", help="reference (fixed) image")
parser.add_argument("test", help="moving image to align")
args = parser.parse_args()

# ─── image loading ───────────────────────────────────────────────────────────────
base = cv2.imread(str(Path(args.base)), cv2.IMREAD_COLOR)
test_orig = cv2.imread(str(Path(args.test)), cv2.IMREAD_UNCHANGED)
if base is None or test_orig is None:
    sys.exit("Error: could not load one or both images.")

h_base, w_base = base.shape[:2]
test_orig = test_orig.astype(np.float32) / 255.0

# ─── interactive state ──────────────────────────────────────────────────────────
offset = np.zeros(2, np.float32)  # x, y
scale = 2.0
angle = 90.0  # degrees
dragging, drag_start = False, np.zeros(2, np.float32)


def affine(a_deg, s, t):
    a = math.radians(a_deg)
    M = np.array(
        [
            [math.cos(a) * s, -math.sin(a) * s, t[0]],
            [math.sin(a) * s, math.cos(a) * s, t[1]],
        ],
        dtype=np.float32,
    )
    cx, cy = test_orig.shape[1] / 2, test_orig.shape[0] / 2
    M[:, 2] += np.dot(M[:, :2], [-cx, -cy]) + [cx, cy]
    return M


def on_mouse(evt, x, y, _flags, _data):
    global dragging, drag_start, offset
    if evt == cv2.EVENT_LBUTTONDOWN:
        dragging, drag_start = True, np.array([x, y], np.float32) - offset
    elif evt == cv2.EVENT_MOUSEMOVE and dragging:
        offset[:] = np.array([x, y], np.float32) - drag_start
    elif evt == cv2.EVENT_LBUTTONUP:
        dragging = False


cv2.namedWindow("register", cv2.WINDOW_NORMAL)
cv2.setMouseCallback("register", on_mouse)

while True:
    M = affine(angle, scale, offset)
    warped = cv2.warpAffine(
        test_orig,
        M,
        (w_base, h_base),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_TRANSPARENT,
    )
    view = np.clip(0.4 * base.astype(np.float32) / 255.0 + 0.6 * warped, 0, 1)
    cv2.imshow("register", (view * 255).astype(np.uint8))

    k = cv2.waitKey(15) & 0xFF
    if k == ord("q"):
        break
    elif k in (ord("+"), ord("=")):
        scale *= 1.01
    elif k in (ord("-"), ord("_")):
        scale /= 1.01
    elif k == 81:
        angle -= 1  # ←
    elif k == 83:
        angle += 1  # →
    elif k == ord("r"):
        offset[:] = 0
        scale = 1.0
        angle = 0.0
    print(f"scale: {scale:.5f}, rotation: {angle:.2f}° (positive = CW), offset: ({offset[0]:.2f}, {offset[1]:.2f}) px  (x, y)")

cv2.destroyAllWindows()
print("\n--- Alignment parameters ---")
print(f"scale   : {scale:.5f}")
print(f"rotation: {angle:.2f}°  (positive = CW)")
print(f"offset  : ({offset[0]:.2f}, {offset[1]:.2f}) px  (x, y)")