import numpy as np

from edge_aware_denoiser import add_gaussian_noise, denoise, quality_metrics


def test_methods():
    rng = np.random.default_rng(7)
    image = rng.random((64, 64, 3), dtype=np.float32)
    noisy = add_gaussian_noise(image, 10, seed=42)

    for method in ("Perona–Malik", "Total Variation", "Bilateral"):
        result, metadata = denoise(noisy, method=method, iterations=12)
        assert result.shape == image.shape
        assert 0 <= result.min() <= result.max() <= 1
        assert metadata.method == method


def test_metrics():
    image = np.zeros((32, 32), dtype=np.float32)
    assert set(quality_metrics(image, image)) == {
        "rmse",
        "psnr",
        "ssim",
        "edge_similarity",
    }


if __name__ == "__main__":
    test_methods()
    test_metrics()
    print("Tests passed.")