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

n = 1000000
cn = ct.c_uint(n)

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.FftC.restype = None
libc.FftC.argtypes = (ct.c_uint, ct.POINTER(Complex * n), ct.POINTER(Complex * n))

libc.IFftC.restype = None
libc.IFftC.argtypes = (ct.c_uint, ct.POINTER(Complex * n), ct.POINTER(Complex * n))

libc.FftCD.restype = None
libc.FftCD.argtypes = (ct.c_uint, ct.POINTER(ComplexD * n), ct.POINTER(ComplexD * n))

libc.IFftCD.restype = None
libc.IFftCD.argtypes = (ct.c_uint, ct.POINTER(ComplexD * n), ct.POINTER(ComplexD * n))

# prepare data
data = np.empty((n), dtype=np.complex64)
spectrum = np.empty((n), dtype=np.complex64)

for i in range(n):
	data[i] = complex(i, 0)	

# start CUDA     
libc.CreateCudaStream()

# Numpy -> Ctypes
cData = data.ctypes.data_as(ct.POINTER(Complex * n)).contents
cSpectrum = spectrum.ctypes.data_as(ct.POINTER(Complex * n)).contents

# FFT -> IFFT	
libc.FftC(cn, ct.pointer(cData), ct.pointer(cSpectrum))
libc.IFftC(cn, ct.pointer(cSpectrum), ct.pointer(cData))

# end CUDA     
libc.DestroyCudaStream()
