from __future__ import print_function, absolute_import, division
from ctypes import *

cu_device = c_int
cu_device_attribute = c_int     # enum
cu_context = c_void_p           # an opaque handle
cu_module = c_void_p            # an opaque handle
cu_jit_option = c_int           # enum
cu_jit_input_type = c_int       # enum
cu_function = c_void_p          # an opaque handle
cu_device_ptr = c_size_t        # defined as unsigned int on 32-bit
                                # and unsigned long long on 64-bit machine
cu_stream = c_void_p            # an opaque handle
cu_event = c_void_p
cu_link_state = c_void_p
cu_function_attribute = c_int


API_PROTOTYPES = {
# CUresult cuInit(unsigned int Flags);
'cuInit' :              (c_int, c_uint),

# CUresult cuDriverGetVersion ( int* driverVersion )
'cuDriverGetVersion':   (c_int, POINTER(c_int)),

# CUresult cuDeviceGetCount(int *count);
'cuDeviceGetCount':     (c_int, POINTER(c_int)),

# CUresult cuDeviceGet(CUdevice *device, int ordinal);
'cuDeviceGet':          (c_int, POINTER(cu_device), c_int),

# CUresult cuDeviceGetName ( char* name, int  len, CUdevice dev )
'cuDeviceGetName':      (c_int, c_char_p, c_int, cu_device),

# CUresult cuDeviceGetAttribute(int *pi, CUdevice_attribute attrib,
#                               CUdevice dev);
'cuDeviceGetAttribute': (c_int, POINTER(c_int), cu_device_attribute,
                         cu_device),

# CUresult cuDeviceComputeCapability(int *major, int *minor,
#                                    CUdevice dev);
'cuDeviceComputeCapability': (c_int, POINTER(c_int), POINTER(c_int),
                              cu_device),

# CUresult cuCtxCreate(CUcontext *pctx, unsigned int flags,
#                      CUdevice dev);
'cuCtxCreate':          (c_int, POINTER(cu_context), c_uint, cu_device),

# CUresult cuCtxGetDevice	(	CUdevice * 	device	 )
'cuCtxGetDevice':       (c_int, POINTER(cu_device)),

# CUresult cuCtxGetCurrent	(CUcontext *pctx);
'cuCtxGetCurrent':      (c_int, POINTER(cu_context)),

# CUresult cuCtxPushCurrent	(CUcontext pctx);
'cuCtxPushCurrent':      (c_int, cu_context),

# CUresult cuCtxPopCurrent	(CUcontext *pctx);
'cuCtxPopCurrent':      (c_int, POINTER(cu_context)),

# CUresult cuCtxDestroy(CUcontext pctx);
'cuCtxDestroy':         (c_int, cu_context),

# CUresult cuModuleLoadDataEx(CUmodule *module, const void *image,
#                             unsigned int numOptions,
#                             CUjit_option *options,
#                             void **optionValues);
'cuModuleLoadDataEx':   (c_int, cu_module, c_void_p, c_uint,
                        POINTER(cu_jit_option), POINTER(c_void_p)),

# CUresult cuModuleUnload(CUmodule hmod);
'cuModuleUnload':       (c_int, cu_module),

# CUresult cuModuleGetFunction(CUfunction *hfunc, CUmodule hmod,
#                              const char *name);
'cuModuleGetFunction':  (c_int, cu_function, cu_module, c_char_p),

# CUresult cuModuleGetGlobal ( CUdeviceptr* dptr, size_t* bytes, CUmodule
#                              hmod, const char* name )
'cuModuleGetGlobal': (c_int, POINTER(cu_device_ptr), POINTER(c_size_t),
                      cu_module, c_char_p),

# CUresult CUDAAPI cuFuncSetCacheConfig(CUfunction hfunc,
#                                       CUfunc_cache config);
'cuFuncSetCacheConfig': (c_int, cu_function, c_uint),

# CUresult cuMemAlloc(CUdeviceptr *dptr, size_t bytesize);
'cuMemAlloc':         (c_int, POINTER(cu_device_ptr), c_size_t),

# CUresult cuMemsetD8(CUdeviceptr dstDevice, unsigned char uc, size_t N)
'cuMemsetD8':         (c_int, cu_device_ptr, c_uint8, c_size_t),

# CUresult cuMemsetD8Async(CUdeviceptr dstDevice, unsigned char uc,
#                          size_t N, CUstream hStream);
'cuMemsetD8Async':    (c_int,
                       cu_device_ptr, c_uint8, c_size_t, cu_stream),

# CUresult cuMemcpyHtoD(CUdeviceptr dstDevice, const void *srcHost,
#                       size_t ByteCount);
'cuMemcpyHtoD':         (c_int, cu_device_ptr, c_void_p, c_size_t),

# CUresult cuMemcpyHtoDAsync(CUdeviceptr dstDevice, const void *srcHost,
#                            size_t ByteCount, CUstream hStream);
'cuMemcpyHtoDAsync':    (c_int, cu_device_ptr, c_void_p, c_size_t,
                        cu_stream),

# CUresult cuMemcpyHtoD(CUdeviceptr dstDevice, const void *srcHost,
#                       size_t ByteCount);
'cuMemcpyDtoD':         (c_int, cu_device_ptr, cu_device_ptr, c_size_t),

# CUresult cuMemcpyHtoDAsync(CUdeviceptr dstDevice, const void *srcHost,
#                            size_t ByteCount, CUstream hStream);
'cuMemcpyDtoDAsync':    (c_int, cu_device_ptr, cu_device_ptr, c_size_t,
                        cu_stream),


# CUresult cuMemcpyDtoH(void *dstHost, CUdeviceptr srcDevice,
#                       size_t ByteCount);
'cuMemcpyDtoH':         (c_int, c_void_p, cu_device_ptr, c_size_t),

# CUresult cuMemcpyDtoHAsync(void *dstHost, CUdeviceptr srcDevice,
#                            size_t ByteCount, CUstream hStream);
'cuMemcpyDtoHAsync':    (c_int, c_void_p, cu_device_ptr, c_size_t,
                        cu_stream),

# CUresult cuMemFree(CUdeviceptr dptr);
'cuMemFree':            (c_int, cu_device_ptr),

# CUresult cuStreamCreate(CUstream *phStream, unsigned int Flags);
'cuStreamCreate':       (c_int, POINTER(cu_stream), c_uint),

# CUresult cuStreamDestroy(CUstream hStream);
'cuStreamDestroy':      (c_int, cu_stream),

# CUresult cuStreamSynchronize(CUstream hStream);
'cuStreamSynchronize':  (c_int, cu_stream),

# CUresult cuLaunchKernel(CUfunction f, unsigned int gridDimX,
#                        unsigned int gridDimY,
#                        unsigned int gridDimZ,
#                        unsigned int blockDimX,
#                        unsigned int blockDimY,
#                        unsigned int blockDimZ,
#                        unsigned int sharedMemBytes,
#                        CUstream hStream, void **kernelParams,
#                        void ** extra)
'cuLaunchKernel':       (c_int, cu_function, c_uint, c_uint, c_uint,
                         c_uint, c_uint, c_uint, c_uint, cu_stream,
                         POINTER(c_void_p), POINTER(c_void_p)),

#  CUresult cuMemHostAlloc	(	void ** 	pp,
#                               size_t 	bytesize,
#                               unsigned int 	Flags
#                           )
'cuMemHostAlloc': (c_int, c_void_p, c_size_t, c_uint),

#  CUresult cuMemFreeHost	(	void * 	p	 )
'cuMemFreeHost': (c_int, c_void_p),

# CUresult cuMemHostRegister(void * 	p,
#                            size_t 	bytesize,
#                            unsigned int 	Flags)
'cuMemHostRegister':    (c_int, c_void_p, c_size_t, c_uint),

# CUresult cuMemHostUnregister(void * 	p)
'cuMemHostUnregister':  (c_int, c_void_p),

# CUresult cuMemHostGetDevicePointer(CUdeviceptr * pdptr,
#                                    void *        p,
#                                    unsigned int  Flags)
'cuMemHostGetDevicePointer': (c_int, POINTER(cu_device_ptr),
                              c_void_p, c_uint),

# CUresult cuMemGetInfo(size_t * free, size_t * total)
'cuMemGetInfo' : (c_int, POINTER(c_size_t), POINTER(c_size_t)),

# CUresult cuEventCreate	(	CUevent * 	phEvent,
#                               unsigned int 	Flags )
'cuEventCreate': (c_int, POINTER(cu_event), c_uint),

# CUresult cuEventDestroy	(	CUevent 	hEvent	 )
'cuEventDestroy': (c_int, cu_event),

# CUresult cuEventElapsedTime	(	float * 	pMilliseconds,
#                                   CUevent 	hStart,
#                                   CUevent 	hEnd )
'cuEventElapsedTime': (c_int, POINTER(c_float), cu_event, cu_event),

# CUresult cuEventQuery	(	CUevent 	hEvent	 )
'cuEventQuery': (c_int, cu_event),

# CUresult cuEventRecord	(	CUevent 	hEvent,
#                               CUstream 	hStream )
'cuEventRecord': (c_int, cu_event, cu_stream),

# CUresult cuEventSynchronize	(	CUevent 	hEvent	 )
'cuEventSynchronize': (c_int, cu_event),


# CUresult cuStreamWaitEvent	(	CUstream        hStream,
#                                   CUevent         hEvent,
#                                	unsigned int 	Flags )
'cuStreamWaitEvent': (c_int, cu_stream, cu_event, c_uint),

# CUresult 	cuPointerGetAttribute (void *data, CUpointer_attribute attribute, CUdeviceptr ptr)
'cuPointerGetAttribute': (c_int, c_void_p, c_uint, cu_device_ptr),

#    CUresult cuMemGetAddressRange	(	CUdeviceptr * 	pbase,
#                                        size_t * 	psize,
#                                        CUdeviceptr 	dptr
#                                        )
'cuMemGetAddressRange': (c_int,
                         POINTER(cu_device_ptr),
                         POINTER(c_size_t),
                         cu_device_ptr),

#    CUresult cuMemHostGetFlags	(	unsigned int * 	pFlags,
#                                   void * 	p )
'cuMemHostGetFlags': (c_int,
                      POINTER(c_uint),
                      c_void_p),

#   CUresult cuCtxSynchronize ( void )
'cuCtxSynchronize' : (c_int,),

#    CUresult
#    cuLinkCreate(unsigned int numOptions, CUjit_option *options,
#                 void **optionValues, CUlinkState *stateOut);
'cuLinkCreate': (c_int,
                 c_uint, POINTER(cu_jit_option),
                 POINTER(c_void_p), POINTER(cu_link_state)),

#    CUresult
#    cuLinkAddData(CUlinkState state, CUjitInputType type, void *data,
#                  size_t size, const char *name, unsigned
#                  int numOptions, CUjit_option *options,
#                  void **optionValues);
'cuLinkAddData': (c_int,
                  cu_link_state, cu_jit_input_type, c_void_p,
                  c_size_t, c_char_p, c_uint, POINTER(cu_jit_option),
                  POINTER(c_void_p)),

#    CUresult
#    cuLinkAddFile(CUlinkState state, CUjitInputType type,
#                  const char *path, unsigned int numOptions,
#                  CUjit_option *options, void **optionValues);

'cuLinkAddFile': (c_int,
                  cu_link_state, cu_jit_input_type, c_char_p, c_uint,
                  POINTER(cu_jit_option), POINTER(c_void_p)),

#    CUresult CUDAAPI
#    cuLinkComplete(CUlinkState state, void **cubinOut, size_t *sizeOut)
'cuLinkComplete': (c_int,
                   cu_link_state, POINTER(c_void_p), POINTER(c_size_t)),

#    CUresult CUDAAPI
#    cuLinkDestroy(CUlinkState state)
'cuLinkDestroy': (c_int, cu_link_state),


# cuProfilerInitialize ( const char* configFile, const char*
# outputFile, CUoutput_mode outputMode )
# 'cuProfilerInitialize': (c_int, c_char_p, c_char_p, cu_output_mode),

# cuProfilerStart ( void )
'cuProfilerStart': (c_int,),

# cuProfilerStop ( void )
'cuProfilerStop': (c_int,),

# CUresult cuFuncGetAttribute ( int* pi, CUfunction_attribute attrib,
#                              CUfunction hfunc )
'cuFuncGetAttribute': (c_int,
                       POINTER(c_int), cu_function_attribute, cu_function),

}
