from PIL import Image, ImageDraw
import numpy as np

from patch_editor import PatchDetector
from pathlib import Path


def create_test_images():
    # Create original image
    img1 = Image.new('RGB', (200, 200), color='white')
    draw1 = ImageDraw.Draw(img1)
    draw1.rectangle([50, 50, 150, 150], fill='blue')

    # Create edited image (shifted and scaled)
    # Shift by (10, 10) and scale by 1.1
    img2 = Image.new('RGB', (220, 220), color='white')
    draw2 = ImageDraw.Draw(img2)
    # Original rect was at 50,50 size 100x100
    # New rect should be at 60,60 size 110x110 roughly if we just draw it
    # But let's just make a similar image but shifted/scaled

    # Let's just rotate and translate the original image to create img2
    # Rotate 5 degrees, translate 10, 10
    img2 = img1.rotate(5, translate=(10, 10), fillcolor='white')
    # And resize it slightly
    img2 = img2.resize((210, 210))

    return img1, img2


def test_warping():
    img1, img2 = create_test_images()

    # Save to temp files
    img1.save('test_orig.png')
    img2.save('test_edit.png')

    print("Testing with allow_rescale=True")
    detector = PatchDetector(original_img_path=Path('test_orig.png'), edited_img_path=Path('test_edit.png'), allow_rescale=True)

    aligned_img1, aligned_img2 = detector.images

    print(f"Original size: {aligned_img1.size}")
    print(f"Aligned size: {aligned_img2.size}")

    assert aligned_img1.size == aligned_img2.size

    # Check if alignment worked reasonably well
    # The difference should be small in the center
    # But since we rotated, edges might be different.

    # Let's just check if it runs without error for now and produces same size image
    print("Warping successful (no crash, sizes match)")


if __name__ == "__main__":
    test_warping()