import pytest
import cv2
import numpy as np

from src.vision.preproc import thresholding_edge_detection, canny_edge_dection, edge_dection, Preprocessing_parameters, PREPROCESSING_MODE

FILES = [
    "9737717d-339f-44d0-9ea1-5d8398fbb418",
    "88004548-3fcd-4bcd-b329-dd75c8fc1917",
    "d132a3aa-ea93-4870-b504-87a97dbed2a3",
    "f778b048-c513-4faf-ae77-50ce30d59732"
]

def resolve_path(path: str) -> str:
    return (
        f"test/data/images/edges/{path}_gray.png",
        f"test/data/images/edges/{path}_thres.png",
        f"test/data/images/edges/{path}_canny.png"
    )

PREPROCESSING_PARAM = Preprocessing_parameters(
    low_threshold=130,
    high_threshold=150,
    kernel_size=19,
    threshold=3,
    gaussian_kernel=(3, 3)   
)

@pytest.mark.parametrize("file", FILES)
def test_edge_dection_thres(file):
    gray_file, thres_file, canny_file = resolve_path(file)
    
    gray_img = cv2.imread(gray_file, cv2.IMREAD_GRAYSCALE)
    ref_img = cv2.imread(thres_file, cv2.IMREAD_GRAYSCALE)
    edge_img = thresholding_edge_detection(gray_img, PREPROCESSING_PARAM)
    
    assert np.array_equal(edge_img, ref_img)
    
@pytest.mark.parametrize("file", FILES)
def test_edge_dection_canny(file):
    gray_file, thres_file, canny_file = resolve_path(file)
    
    gray_img = cv2.imread(gray_file, cv2.IMREAD_GRAYSCALE)
    ref_img = cv2.imread(canny_file, cv2.IMREAD_GRAYSCALE)
    print(gray_img.shape)
    edge_img = canny_edge_dection(gray_img, PREPROCESSING_PARAM)
    
    assert np.array_equal(edge_img, ref_img)
    
@pytest.mark.parametrize("mode,file", list(zip( 
    [PREPROCESSING_MODE.THRES]*len(FILES) + [PREPROCESSING_MODE.CANNY]*len(FILES),
    FILES*2
)))
def test_edge_dection(mode:str, file:str):
    gray_file, thres_file, canny_file = resolve_path(file)
    
    gray_img = cv2.imread(gray_file, cv2.IMREAD_GRAYSCALE)
    ref_thres_img = cv2.imread(thres_file, cv2.IMREAD_GRAYSCALE)
    ref_canny_img = cv2.imread(canny_file, cv2.IMREAD_GRAYSCALE)
    
    thres_img = edge_dection(gray_img, PREPROCESSING_MODE.THRES, PREPROCESSING_PARAM)
    canny_img = edge_dection(gray_img, PREPROCESSING_MODE.CANNY, PREPROCESSING_PARAM)
    
    assert np.array_equal(thres_img, ref_thres_img)
    assert np.array_equal(canny_img, ref_canny_img)
 
def test_edge_dection_invalid_mode():
    gray_img = np.zeros((10, 10), dtype=np.uint8)
    
    with pytest.raises(ValueError):
        edge_dection(gray_img, "INVALID_MODE", PREPROCESSING_PARAM)