import numpy as np
import pytest

from he11lib.deconvolution import DiffusionDeconvolver


def gaussian_bump(n, sigma_px):
    coords = np.arange(n) - n // 2
    xx, yy = np.meshgrid(coords, coords)
    return np.exp(-(xx**2 + yy**2) / (2 * sigma_px**2))


def profile_std(image):
    n = image.shape[0]
    coords = np.arange(n) - n // 2
    weights = image[n // 2, :]
    weights = np.clip(weights, 0, None)
    mean = np.sum(coords * weights) / np.sum(weights)
    var = np.sum(weights * (coords - mean) ** 2) / np.sum(weights)
    return np.sqrt(var)


def test_blur_sigma_m_from_diffusivity_and_dwell_time():
    diffusivity = 1e-6  # m^2/s
    dwell_time = 0.5  # s
    deconvolver = DiffusionDeconvolver(thermal_diffusivity=diffusivity, dwell_time=dwell_time)

    expected_sigma = np.sqrt(2 * diffusivity * dwell_time)
    assert deconvolver.blur_sigma_m() == pytest.approx(expected_sigma)


def test_blur_widens_a_sharp_peak():
    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.5)
    pixel_scale = 2e-4
    sharp = gaussian_bump(101, sigma_px=3)

    blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale)

    assert profile_std(blurred) > profile_std(sharp)


def test_deconvolve_reduces_error_relative_to_blurred():
    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.3)
    pixel_scale = 2e-4
    sharp = gaussian_bump(101, sigma_px=4)

    blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale)
    deconvolved = deconvolver.deconvolve(blurred, pixel_scale=pixel_scale)

    error_blurred = np.sum((blurred - sharp) ** 2)
    error_deconvolved = np.sum((deconvolved - sharp) ** 2)
    assert error_deconvolved < error_blurred


def test_deconvolve_narrows_width_back_toward_original():
    deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.3)
    pixel_scale = 2e-4
    sharp = gaussian_bump(101, sigma_px=4)

    blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale)
    deconvolved = deconvolver.deconvolve(blurred, pixel_scale=pixel_scale)

    std_sharp = profile_std(sharp)
    std_blurred = profile_std(blurred)
    std_deconvolved = profile_std(deconvolved)

    assert std_sharp < std_deconvolved < std_blurred