import pypocketfft
import pyfftw
import numpy as np
import pytest
from numpy.testing import assert_allclose

pmp = pytest.mark.parametrize

shapes = ((10,), (127,), (128, 128), (128, 129), (1, 129), (129,1), (32,17,39))

@pmp("shp", shapes)
def test_fftn(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    vmax = np.max(np.abs(pyfftw.interfaces.numpy_fft.fftn(a)))
    assert_allclose(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a), atol=2e-15*vmax, rtol=0)
    a=a.astype(np.complex64)
    assert_allclose(pyfftw.interfaces.numpy_fft.fftn(a), pypocketfft.fftn(a), atol=6e-7*vmax, rtol=0)

@pmp("shp", shapes)
def test_identity(shp):
    a=np.random.rand(*shp)-0.5 + 1j*np.random.rand(*shp)-0.5j
    vmax = np.max(np.abs(a))
    assert_allclose(pypocketfft.ifftn(pypocketfft.fftn(a))/a.size, a, atol=2e-15*vmax, rtol=0)
    a=a.astype(np.complex64)
    assert_allclose(pypocketfft.ifftn(pypocketfft.fftn(a))/a.size, a, atol=6e-7*vmax, rtol=0)

@pmp("shp", shapes)
def test_identity_r(shp):
    a=np.random.rand(*shp)-0.5
    b=a.astype(np.float32)
    vmax = np.max(np.abs(a))
    for ax in range(a.ndim):
        n = a.shape[ax]
        assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(a,ax),ax)/n, a, atol=2e-15*vmax, rtol=0)
        assert_allclose(pypocketfft.irfftn(pypocketfft.rfftn(b,ax),ax)/n, b, atol=6e-7*vmax, rtol=0)