using PyCall
@pyimport nufft as nufft_fortran
@pyimport numpy as np

function nufftfreqs(M,df=1.0)
    """Compute the frequency range used in nufft for M frequency bins"""
    df * [-fld(M, 2): M - fld(M, 2)-1]
end

function nudftpy(x, y, M, df=1.0, iflag=1)
    """Non-Uniform Direct Fourier Transform. Using numpy"""
    sign = iflag < 0 ? -1 : 1
    (1 / length(x)) * np.dot(y, np.exp(sign * 1im * x*transpose(nufftfreqs(M, df))))
end

function nudft(x, y, M, df=1.0, iflag=1)
    """Non-Uniform Direct Fourier Transform. Using whole array operations"""
    sign = iflag < 0 ? -1 : 1
    (1 / length(x)) * *(transpose(y''), exp(sign * 1im * x*transpose(nufftfreqs(M, df))))
end

function nudft2(x::Vector, y::Vector, M::Int, df=1.0, iflag=1)
    """Non-Uniform Direct Fourier Transform. Using comprehensions"""
    freqs = nufftfreqs(M, df)
    sign = iflag < 0 ? -1 : 1
    n = size(x,1)
    m = size(freqs, 1)
    r = (1 / n) * [y[i]*exp(sign*1im*x[i]*freqs[j]) for i=1:n, j=1:m]
    sum(r,1)
end

function nudft3(x::Vector, y::Vector, M::Int, df=1.0, iflag=1)
    """Non-Uniform Direct Fourier Transform. Using comprehensions, summing different axis"""
    freqs = nufftfreqs(M, df)
    sign = iflag < 0 ? -1 : 1
    n = size(x,1)
    m = size(freqs, 1)
    r = (1 / n) * [y[i]*exp(sign*1im*x[i]*freqs[j]) for j=1:m, i=1:n]
    transpose(sum(r,2))
end


x = 100 * rand(1000)
y = sin(x)
Y0 = @time nudftpy(x, y, 1000) 
Y1 = @time nudft(x, y, 1000)
Y2 = @time nudft2(x, y, 1000)
Y3 = @time nudft3(x, y, 1000)
Yf = @time nufft_fortran.nufft1(x, y, 1000)
print([np.allclose(Y0, Yf), np.allclose(Y1, Yf),np.allclose(Y2, Yf),np.allclose(Y3, Yf)])

function _compute_grid_params(M, epsilon)
    # Choose Msp & tau from eps following Dutt & Rokhlin (1993)
    if epsilon <= 1E-33 || epsilon >= 1E-1
        error(@sprintf("eps = %f; must satisfy 1e-33 < eps < 1e-1.", epsilon))
    end
    ratio = epsilon > 1E-11 ? 2 : 3
    Msp = itrunc(-log(epsilon) / (pi * (ratio - 1) / (ratio - 0.5)) + 0.5)
    Mr = max(ratio * M, 2 * Msp)
    lambda_ = Msp / (ratio * (ratio - 0.5))
    tau = pi * lambda_ / M^2
    Msp, Mr, tau
end

function nufft_python(x, c, M, df=1.0, epsilon=1E-15, iflag=1):
    """Fast Non-Uniform Fourier Transform with Python"""
    Msp, Mr, tau = _compute_grid_params(M, epsilon)
    N = length(x)

    # Construct the convolved grid
    ftau = zeros(typeof(c[1]), Mr)
    Mr = size(ftau,1)
    hx = 2pi / Mr
    mm = [-Msp:Msp-1]
    for i=1:N
        xi = (x[i] * df) % (2 * pi)
        m = 1 + div(xi,hx)
        spread = exp(-0.25 * (xi - hx * (m + mm)).^2 / tau)
        ftau[1+mod(m + mm, Mr)] += c[i] * spread
    end
    # Compute the FFT on the convolved grid
    if iflag < 0
        Ftau = (1 / Mr) * fft(ftau)
    else
        Ftau = ifft(ftau)
    end
    Ftau = [Ftau[end-div(M,2)+1:end], Ftau[1:div(M,2)+M%2]]
    # Deconvolve the grid using convolution theorem
    k = nufftfreqs(M)
    (Ftau.*(1 / N) * sqrt(pi / tau)).* exp(tau * k.^2)
end

function test_nufft(nufft_func, M=1000, Mtime=100000)
    # Test vs the direct method
    print(repeat("-",30), "\n")
    print("testing ",nufft_func, "\n")
    x = 100 * rand(M + 1)
    y = sin(x)
    for df in [1.0, 2.0]
        for iflag in [1, -1]
            F1 = nudft(x, y, M, df, iflag)
            F2 = nufft_func(x, y, M, df, 1E-15, iflag)
            assert(all(x -> isapprox(x...), zip(F1, F2)))
        end
    end
    print("- Results match the DFT\n")
    
    # Time the nufft function
    x = 100 * rand(Mtime)
    y = sin(x)
    times = Float64[]
    for i = 1:5
        tic()
        F = nufft_func(x, y, Mtime)
        t1 = toq()
        push!(times,t1)
    end
    @printf("- Execution time (M=%d): %.2f sec\n",Mtime, median(times))
end

test_nufft(nufft_python)
test_nufft(nufft_fortran.nufft1)

function nufft_numpy(x, y, M, df=1.0, epsilon=1E-15, iflag=1):
    """Fast Non-Uniform Fourier Transform"""
    Msp, Mr, tau = _compute_grid_params(M, epsilon)
    N = length(x)
    # Construct the convolved grid
    ftau = zeros(typeof(y[1]), Mr)
    hx = 2pi / Mr
    xmod = map(mod2pi, x*df)
    m = 1+ int(xmod/hx)
    mm = [-Msp:Msp-1]
    mpmm = broadcast(+, transpose(m), mm)
    spread = broadcast(*, exp(-0.25 * (transpose(xmod).- hx*mpmm).^ 2 / tau), transpose(y))
    for (i,s) in zip(map(xi->1+mod(xi, Mr), mpmm), spread)
        ftau[i] += s
    end
    # Compute the FFT on the convolved grid
    if iflag < 0
        Ftau = (1 / Mr) * fft(ftau)
    else
        Ftau = ifft(ftau)
    end
    Ftau = [Ftau[end-div(M,2)+1:end], Ftau[1:div(M,2)+M%2]]
    # Deconvolve the grid using convolution theorem
    k = nufftfreqs(M)
    (Ftau.*(1 / N) * sqrt(pi / tau)).* exp(tau * k.^2)
end

test_nufft(nufft_numpy)
test_nufft(nufft_fortran.nufft1)

function build_grid(x, c, tau, Msp, ftau)
    Mr = size(ftau,1)
    hx = 2pi / Mr
    for i=1:size(x,1)
        xi = mod2pi(x[i])
        m = 1 + int(xi/hx)
        for mm=-Msp:Msp-1
            ftau[1 + mod((m + mm) , Mr)] += c[i] * exp(-0.25 * (xi - hx * (m + mm))^2 / tau)
        end
    end
    ftau
end

function nufft_numba(x, c, M, df=1.0, eps=1E-15, iflag=1)
    """Fast Non-Uniform Fourier Transform from Python numba code"""
    Msp, Mr, tau = _compute_grid_params(M, eps)
    N = length(x)

    # Construct the convolved grid
    ftau = build_grid(x * df, c, tau, Msp, zeros(typeof(c[1]), Mr))

    # Compute the FFT on the convolved grid
    if iflag < 0
        Ftau = (1 / Mr) * fft(ftau)
    else
        Ftau = ifft(ftau)
    end
    Ftau = [Ftau[end-div(M,2)+1:end], Ftau[1:div(M,2)+M%2]]

    # Deconvolve the grid using convolution theorem
    k = nufftfreqs(M)
    (1 / N) * sqrt(pi / tau) .* exp(tau * k.^2).*Ftau
end

test_nufft(nufft_numba)
test_nufft(nufft_fortran.nufft1)

function build_grid_fast(x, c, tau, Msp, ftau, E3)
    Mr = size(ftau,1)
    hx = 2pi / Mr
    # precompute some exponents
    for j=0:Msp
        E3[j+1] = exp(-(pi * j / Mr)^2 / tau)
    end
    # spread values onto ftau
    for i=1:size(x,1)
        xi = mod2pi(x[i])
        m = 1 + int(xi/hx)
        xi = xi - hx * m
        E1 = exp(-0.25 * xi^2 / tau)
        E2 = exp((xi * pi) / (Mr * tau))
        E2mm = 1
        for mm=0:Msp-1
            ftau[1+mod((m + mm) , Mr)] += c[i] * E1 * E2mm * E3[mm+1]
            E2mm *= E2
            ftau[1+mod((m - mm - 1) , Mr)] += c[i] * E1 / E2mm *E3[mm+2]
        end
    end
    ftau
end

function nufft_numba_fast(x, c, M, df=1.0, eps=1E-15, iflag=1)
    """Fast Non-Uniform Fourier Transform from Python numba code"""
    Msp, Mr, tau = _compute_grid_params(M, eps)
    N = length(x)

    # Construct the convolved grid
    ftau = build_grid_fast(x * df, c, tau, Msp, 
    zeros(typeof(c[1]), Mr), zeros(typeof(x[1]), Msp+1) )

    # Compute the FFT on the convolved grid
    if iflag < 0
        Ftau = (1 / Mr) * fft(ftau)
    else
        Ftau = ifft(ftau)
    end
    Ftau = [Ftau[end-div(M,2)+1:end], Ftau[1:div(M,2)+M%2]]

    # Deconvolve the grid using convolution theorem
    k = nufftfreqs(M)
    (1 / N) * sqrt(pi / tau) .* exp(tau * k.^2).*Ftau
end

test_nufft(nufft_numba_fast)
test_nufft(nufft_fortran.nufft1)