import cv2
import numpy as np
import time
import pycuda.autoinit
import pycuda.driver as cuda
from pycuda.compiler import SourceModule
import pycuda.gpuarray as gpuarray
def apply_threshold(img_src,img_width, img_height, img_dest, mythreshold):
    mod = SourceModule("""
        __global__ void ThresholdKernel(
            const int src_sizeX,  //< source image size. x: width,
            const unsigned char* src,   //< source image pointer
            const int dst_sizeX,  //< destination image size. x: width, y: height
            const int dst_sizeY,
            unsigned char* dst,         //< destination image pointer
            const int mythreshold) {
                int col = blockIdx.x * blockDim.x + threadIdx.x;
                int row = blockIdx.y * blockDim.y + threadIdx.y;
                if (dst_sizeX <= col || dst_sizeY <= row) return;
                auto src_val = src[row * src_sizeX + col];
                unsigned char dst_val = src_val > mythreshold ? 255 : 0;
                dst[row * dst_sizeX + col] = dst_val;
            }
    """)
    block_dim =(32,8,1)
    grid_dim_x = (img_width + block_dim[0] -1) // block_dim[0]
    grid_dim_y = (img_width + block_dim[1] -1) // block_dim[1]
    print(grid_dim_x,grid_dim_y)
    
    thresholdkernel = mod.get_function("ThresholdKernel")
    thresholdkernel(np.int32(img_width), img_src, np.int32(img_width),np.int32(img_height), 
                    img_dest,np.int32(mythreshold),
                    block = block_dim , grid = (grid_dim_x,grid_dim_y))
    
mythreshold = 128
img_path = "../images/lena_gray.png"
img = cv2.imread(img_path)
if img is None:
    print("Image not found")
    exit()
else:
    height,width,channels = img.shape
    print("Hegiht, width and channels",height,width,channels)
    print(type(width))
img_gpu = cuda.mem_alloc(img.nbytes)
cuda.memcpy_htod(img_gpu,img)
dtype=img.dtype
# dest_img=gpuarray.empty_like(img.shape,dtype=dtype)
dest_img = cuda.mem_alloc(img.nbytes)
apply_threshold(img_gpu,width,height,dest_img  ,mythreshold )
image_result= np.empty_like(img)
cuda.memcpy_dtoh(image_result,dest_img )
cv2.imshow("Original image",img)
cv2.imshow("Thresholded",image_result)
cv2.waitKey(0)
cv2.destroyAllWindows()

img = cv2.imread(img_path,cv2.IMREAD_GRAYSCALE)

grid_dim_x = (img_width + block_dim[0] -1) // block_dim[0]
grid_dim_y = (img_height + block_dim[1] -1) // block_dim[1]