import numpy as np
from discrete_hankel_transform import HankelTransform


def test_parseval_gaussian():
    """
    Verify Parseval identity using DHT quadrature weights.
    """

    N = 512
    r_max = 30.0
    sigma = 3.0

    ht = HankelTransform(N, order=0, r_max=r_max)

    # Gaussian test function
    psi_r = np.exp(-ht.r**2 / (2 * sigma**2)) * np.sqrt(2.0 / sigma**2)

    psi_k = ht.forward(psi_r)

    norm_r = np.sum(ht.measure * np.abs(psi_r)**2)
    norm_k = np.sum(ht.measure_k * np.abs(psi_k)**2)

    rel_error = abs(norm_r - norm_k) / norm_r

    # tolerance depends on grid size
    assert np.allclose(norm_k, 1.0)
    assert np.allclose(norm_r, 1.0)
    assert rel_error < 1e-8, f"Parseval failed: error={rel_error}"

def test_gaussian_normalization():
    """
    Check that a normalized Gaussian has unit norm under the DHT measure.
    """

    N = 512
    r_max = 30.0
    sigma = 2.0

    ht = HankelTransform(N, order=0, r_max=r_max)

    # normalized 2D radial Gaussian:
    # psi(r) = (1/(sqrt(pi)*sigma)) * exp(-r^2/(2 sigma^2))
    # so that ∫ |psi|^2 r dr = 1
    psi = np.sqrt(2.0 / sigma**2) * np.exp(-ht.r**2 / (2 * sigma**2))

    norm = np.sum(ht.measure * np.abs(psi)**2)

    assert abs(norm - 1) < 1e-8, f"Normalization failed: norm={norm}"

