import numpy as np
import ctypes as ct
libc = ct.cdll.LoadLibrary("libMatxFft.so")

w = 4096
h = 4096
cw = ct.c_uint(w)
ch = ct.c_uint(h)
length = w * h

class Complex(ct.Structure):
    _fields_ = [('real', ct.c_float),
                  ('imag', ct.c_float)]
                  
class ComplexD(ct.Structure):
    _fields_ = [('real', ct.c_double),
                  ('imag', ct.c_double)]

libc.Fft2C.restype = None
libc.Fft2C.argtypes = (ct.c_uint, ct.c_uint, ct.POINTER(Complex * length), ct.POINTER(Complex * length))

libc.IFft2C.restype = None
libc.IFft2C.argtypes = (ct.c_uint, ct.c_uint, ct.POINTER(Complex * length), ct.POINTER(Complex * length))

libc.Fft2CD.restype = None
libc.Fft2CD.argtypes = (ct.c_uint, ct.c_uint, ct.POINTER(ComplexD * length), ct.POINTER(ComplexD * length))

libc.IFft2CD.restype = None
libc.IFft2CD.argtypes = (ct.c_uint, ct.c_uint, ct.POINTER(ComplexD * length), ct.POINTER(ComplexD * length))

# prepare data
image = np.empty((w, h), dtype=np.complex64)
spectrum = np.empty((w, h), dtype=np.complex64)

for y in range(h):
	for x in range(w):
		image[x, y] = complex(y, x)	

# start CUDA     
libc.CreateCudaStream()

# Numpy -> Ctypes
image1 = image.reshape(length)
cImage = image1.ctypes.data_as(ct.POINTER(Complex * length)).contents
spectrum1 = spectrum.reshape(length)
cSpectrum = spectrum1.ctypes.data_as(ct.POINTER(Complex * length)).contents

#FFT -> IFFT(result -> image)
libc.Fft2C(cw, ch, ct.pointer(cImage), ct.pointer(cSpectrum))
libc.IFft2C(cw, ch, ct.pointer(cSpectrum), ct.pointer(cImage))
	
# end CUDA     
libc.DestroyCudaStream()
