import numpy as np
import pytest

from he11lib.noise import NoiseEstimator


def test_estimate_std_recovers_known_noise_on_flat_image():
    rng = np.random.default_rng(0)
    true_std = 0.05
    image = np.ones((200, 200)) * 10.0 + rng.normal(0, true_std, size=(200, 200))

    estimated = NoiseEstimator().estimate_std(image)
    assert estimated == pytest.approx(true_std, rel=0.15)


def test_estimate_std_recovers_known_noise_on_smooth_bump():
    rng = np.random.default_rng(1)
    x = np.linspace(-3, 3, 200)
    xx, yy = np.meshgrid(x, x)
    smooth = np.exp(-(xx**2 + yy**2))
    true_std = 0.01
    image = smooth + rng.normal(0, true_std, size=smooth.shape)

    estimated = NoiseEstimator().estimate_std(image)
    assert estimated == pytest.approx(true_std, rel=0.35)


def test_estimate_std_near_zero_for_noise_free_image():
    x = np.linspace(-3, 3, 100)
    xx, yy = np.meshgrid(x, x)
    smooth = np.exp(-(xx**2 + yy**2))

    estimated = NoiseEstimator().estimate_std(smooth)
    assert estimated < 1e-6


def test_weights_are_uniform_and_match_shape():
    rng = np.random.default_rng(2)
    image = np.ones((50, 50)) * 5.0 + rng.normal(0, 0.1, size=(50, 50))

    weights = NoiseEstimator().weights(image)
    assert weights.shape == image.shape
    assert np.allclose(weights, weights.flat[0])

    expected_std = NoiseEstimator().estimate_std(image)
    assert weights.flat[0] == pytest.approx(1.0 / expected_std**2, rel=1e-6)