import numpy as np

# Import from main directory
import os, sys
sys.path.append(os.path.abspath('..'))

import nfft
nfft.__version__

import pynfft
pynfft.__version__

"""
Explanation: Benchmarks
This notebook contains benchmarks against the pynfft package, which is a Python wrapper of Fortran code
End of explanation
"""


def make_forward_data(M, N):
    x = -0.5 + np.random.rand(M)
    f_hat = np.random.randn(N) + 1j * np.random.randn(N)
    return x, f_hat

"""
Explanation: Benchmarking the Forward Transform
Define some test data:
End of explanation
"""


def pynfft_forward(x, f_hat):
    M = len(x)
    N = len(f_hat)
    plan = pynfft.nfft.NFFT(N, M)
    plan.x = x
    plan.precompute()
    plan.f_hat = f_hat
    # Need copy because of bug in pynfft 1.x
    # See https://github.com/ghisvail/pyNFFT/issues/57
    return plan.trafo().copy()

x, f_hat = make_forward_data(1000, 100000)

out1 = nfft.nfft(x, f_hat)
out2 = pynfft_forward(x, f_hat)
np.allclose(out1, out2)

%timeit nfft.nfft(x, f_hat)
%timeit pynfft_forward(x, f_hat)

"""
Explanation: Define a utility function around pynfft:
End of explanation
"""


def make_adjoint_data(M):
    x = -0.5 + np.random.rand(M)
    f = np.random.randn(M) + 1j * np.random.randn(M)
    return x, f

"""
Explanation: Benchmarking the Adjoint Transform
Define some test data:
End of explanation
"""


def pynfft_adjoint(x, f, N):
    M = len(x)
    plan = pynfft.nfft.NFFT(N, M)
    plan.x = x
    plan.precompute()
    plan.f = f
    # Need copy because of bug in pynfft 1.x
    # See https://github.com/ghisvail/pyNFFT/issues/57
    return plan.adjoint().copy()

x, f = make_adjoint_data(1000)
N = 100000

out1 = nfft.nfft_adjoint(x, f, N)
out2 = pynfft_adjoint(x, f, N)

np.allclose(out1, out2)

%timeit nfft.nfft_adjoint(x, f, N, sigma=3)
%timeit pynfft_adjoint(x, f, N)

"""
Explanation: Define a utility function around pynfft:
End of explanation
"""


