import numpy as np
import pypocketfft

def _l2error(a,b):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))

nthreads=0
cmaxerr=0.
fmaxerr=0.
cmaxerrf=0.
fmaxerrf=0.
hmaxerr=0.
hmaxerrf=0.
def test():
    global cmaxerr, fmaxerr, hmaxerr, cmaxerrf, fmaxerrf, hmaxerrf
    ndim=np.random.randint(1,5)
    axlen=int((2**20)**(1./ndim))
    shape = np.random.randint(1,axlen,ndim)
    axes = np.arange(ndim)
    np.random.shuffle(axes)
    nax = np.random.randint(1,ndim+1)
    axes = axes[:nax]
    lastsize = shape[axes[-1]]
    fct = 1./np.prod(np.take(shape, axes))
    a=np.random.rand(*shape)-0.5 + 1j*np.random.rand(*shape)-0.5j
    b=pypocketfft.ifftn(pypocketfft.fftn(a,axes=axes,nthreads=nthreads),axes=axes,fct=fct,nthreads=nthreads)
    err = _l2error(a,b)
    if err > cmaxerr:
        cmaxerr = err
        print("cmaxerr:", cmaxerr, shape, axes)
    b=pypocketfft.irfftn(pypocketfft.rfftn(a.real,axes=axes,nthreads=nthreads),axes=axes,fct=fct,lastsize=lastsize,nthreads=nthreads)
    err = _l2error(a.real,b)
    if err > fmaxerr:
        fmaxerr = err
        print("fmaxerr:", fmaxerr, shape, axes)
    b=pypocketfft.ifftn(pypocketfft.fftn(a.astype(np.complex64),axes=axes,nthreads=nthreads),axes=axes,fct=fct,nthreads=nthreads)
    err = _l2error(a.astype(np.complex64),b)
    if err > cmaxerrf:
        cmaxerrf = err
        print("cmaxerrf:", cmaxerrf, shape, axes)
    b=pypocketfft.irfftn(pypocketfft.rfftn(a.real.astype(np.float32),axes=axes,nthreads=nthreads),axes=axes,fct=fct,lastsize=lastsize,nthreads=nthreads)
    err = _l2error(a.real.astype(np.float32),b)
    if err > fmaxerrf:
        fmaxerrf = err
        print("fmaxerrf:", fmaxerrf, shape, axes)
    b=pypocketfft.hartley(pypocketfft.hartley(a.real,axes=axes,nthreads=nthreads),axes=axes,fct=fct,nthreads=nthreads)
    err = _l2error(a.real,b)
    if err > hmaxerr:
        hmaxerr = err
        print("hmaxerr:", hmaxerr, shape, axes)
    b=pypocketfft.hartley(pypocketfft.hartley(a.real.astype(np.float32),axes=axes,nthreads=nthreads),axes=axes,fct=fct,nthreads=nthreads)
    err = _l2error(a.real.astype(np.float32),b)
    if err > hmaxerrf:
        hmaxerrf = err
        print("hmaxerrf:", hmaxerrf, shape, axes)

while True:
    test()