import cv2
import numpy as np

from aare.daq.autofocus import calculate_focus_measure


def test_calculate_focus_measure_grayscale():
    # Create a focused spot (a sharp square)
    image = np.zeros((100, 100), dtype=np.uint8)
    image[40:60, 40:60] = 255

    # Calculate focus measure at the center
    fm = calculate_focus_measure(image, 50, 50, 20)
    assert fm > 0

    # Create a blurred spot
    blurred = cv2.GaussianBlur(image, (21, 21), 0)
    fm_blurred = calculate_focus_measure(blurred, 50, 50, 20)

    # Focused image should have higher variance of Laplacian
    assert fm > fm_blurred


def test_calculate_focus_measure_color():
    # Create a color image
    image = np.zeros((100, 100, 3), dtype=np.uint8)
    image[40:60, 40:60, 0] = 255  # Blue square

    fm = calculate_focus_measure(image, 50, 50, 20)
    assert fm > 0


def test_calculate_focus_measure_out_of_bounds():
    image = np.zeros((100, 100), dtype=np.uint8)
    # Even if radius goes out of bounds, np.ogrid and mask should handle it
    fm = calculate_focus_measure(image, 0, 0, 200)
    assert fm == 0  # All zeros, so variance is 0