import numpy as np
import cv2

def get_random_patch(image, patch_size):
    h, w = image.shape[1:3]
    x = np.random.randint(0, w - patch_size)
    y = np.random.randint(0, h - patch_size)
    return image[:, y:y+patch_size, x:x+patch_size]

def get_random_homography(max_rotation=30, max_translation=20):
    theta = np.random.uniform(-max_rotation, max_rotation) * np.pi / 180.0
    tx = np.random.uniform(-max_translation, max_translation)
    ty = np.random.uniform(-max_translation, max_translation)
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)
    H = np.array([
        [cos_theta, -sin_theta, tx],
        [sin_theta, cos_theta, ty],
        [0, 0, 1]
    ])
    return H

def apply_homography_to_image(image, H, output_size):
    warped = np.zeros_like(image)
    for c in range(image.shape[0]):
        warped[c] = cv2.warpPerspective(image[c], H, output_size, flags=cv2.INTER_NEAREST)
    return warped

def generate_training_pair(image, patch_size):
    patch = get_random_patch(image, patch_size)
    H = get_random_homography()
    warped_patch = apply_homography_to_image(patch, H, (patch_size, patch_size))
    return patch, warped_patch, H