"""Tests for image_noise.py."""

import numpy as np
import pytest
from image_noise import (add_gaussian, add_salt_pepper, add_speckle,
                          add_shot_noise, denoise_median, denoise_bilateral)


@pytest.fixture
def test_img():
    """Simple gradient test image."""
    return np.linspace(0, 1, 400).reshape(20, 20)


class TestAddGaussian:
    def test_range(self, test_img):
        noisy = add_gaussian(test_img, sigma=0.1, seed=42)
        assert noisy.min() >= 0 and noisy.max() <= 1

    def test_changes_image(self, test_img):
        noisy = add_gaussian(test_img, sigma=0.1, seed=1)
        assert not np.allclose(noisy, test_img)


class TestAddSaltPepper:
    def test_range(self, test_img):
        noisy = add_salt_pepper(test_img, prob=0.05, seed=42)
        assert noisy.min() == 0.0 and noisy.max() == 1.0

    def test_has_extreme_values(self, test_img):
        noisy = add_salt_pepper(test_img, prob=0.1, seed=7)
        assert np.any(noisy == 0.0) or np.any(noisy == 1.0)


class TestAddSpeckle:
    def test_range(self, test_img):
        noisy = add_speckle(test_img, sigma=0.2, seed=42)
        assert noisy.min() >= 0 and noisy.max() <= 1

    def test_dark_regions_less_absolute_noise(self, test_img):
        """Speckle is multiplicative: dark regions should vary less absolutely."""
        noisy = add_speckle(test_img, sigma=0.1, seed=1)
        # First column (dark) should have smaller variance than last (bright).
        var_dark = np.var(noisy[:, 0])
        var_bright = np.var(noisy[:, -1])
        assert var_bright > var_dark


class TestAddShotNoise:
    def test_range(self, test_img):
        noisy = add_shot_noise(test_img, peak_photons=200, seed=42)
        assert noisy.min() >= 0 and noisy.max() <= 1

    def test_low_photons_noisier(self, test_img):
        noisy_low = add_shot_noise(test_img, peak_photons=10, seed=1)
        noisy_high = add_shot_noise(test_img, peak_photons=1000, seed=1)
        assert np.std(noisy_low - test_img) > np.std(noisy_high - test_img)


class TestDenoiseMedian:
    def test_shape_preserved(self, test_img):
        noisy = add_salt_pepper(test_img, prob=0.1, seed=1)
        denoised = denoise_median(noisy, size=3)
        assert denoised.shape == test_img.shape

    def test_removes_salt_pepper(self):
        img = np.ones((10, 10)) * 0.5
        img[3, 3] = 0.0
        img[4, 4] = 1.0
        denoised = denoise_median(img, size=3)
        assert denoised[3, 3] == 0.5
        assert denoised[4, 4] == 0.5


class TestDenoiseBilateral:
    def test_shape_preserved(self, test_img):
        noisy = add_gaussian(test_img, sigma=0.1, seed=1)
        denoised = denoise_bilateral(noisy, sigma_spatial=2.0, sigma_intensity=0.1)
        assert denoised.shape == test_img.shape

    def test_preserves_edges_better_than_naive_mean(self):
        """Bilateral filter should preserve a sharp edge better than a box filter."""
        img = np.zeros((20, 20))
        img[:, :10] = 0.2
        img[:, 10:] = 0.8
        noisy = add_gaussian(img, sigma=0.05, seed=42)
        denoised = denoise_bilateral(noisy, sigma_spatial=3.0, sigma_intensity=0.15)

        # Edge steepness: difference across the boundary should be preserved.
        edge_grad = np.abs(np.mean(denoised[:, 10]) - np.mean(denoised[:, 9]))
        # Box filter would blur this; bilateral should preserve much of it.
        assert edge_grad > 0.3