import static jcuda.driver.JCudaDriver.cuCtxCreate;
import static jcuda.driver.JCudaDriver.cuCtxSynchronize;
import static jcuda.driver.JCudaDriver.cuDeviceGet;
import static jcuda.driver.JCudaDriver.cuInit;
import static jcuda.driver.JCudaDriver.cuLaunchKernel;
import static jcuda.driver.JCudaDriver.cuMemAlloc;
import static jcuda.driver.JCudaDriver.cuMemFree;
import static jcuda.driver.JCudaDriver.cuMemcpyDtoH;
import static jcuda.driver.JCudaDriver.cuMemcpyHtoD;

import java.awt.Graphics2D;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import javax.imageio.ImageIO;
import javax.swing.SwingWorker;

import ij.IJ;
import jcuda.Pointer;
import jcuda.Sizeof;
import jcuda.driver.CUcontext;
import jcuda.driver.CUdevice;
import jcuda.driver.CUdeviceptr;
import jcuda.driver.CUfunction;
import jcuda.driver.JCudaDriver;
import jcuda.nvrtc.JNvrtc;

/**
 * This class will calculate the percentage of Color in an Image. 
 * Threads are GPU Device. 
 * 
 * @author Aqeel
 *
 */
public class ColorPercentageWithWorker extends SwingWorker<Void, Void>{
	
	public static String NAME = "Calculate Color Percentage";
	
	/**
	 * List of File Path of Images
	 */
	private List<File> imagesPaths;
	
	/**
	 * Producer will be responsible to Load the Image Data and Send it to Consumer. 
	 * 
	 */
	private BlockingQueue<File> producer;
	
	/**
	 * Consumer will be responsible to Take the Image Data and Perform Calculation on it. 
	 */
	private BlockingQueue<ImageData> consumer;
	
	/**
	 * Number Of Threads Responsible to Load Image Data. 
	 */
	private List<ImageLoaderTask> threads;
	
	/**
	 * Number of GPU Devices to Perform Action on Pixels
	 */
	private List<PercentageTask> consumerThreads;
	
	/**
	 * Simple Date Format. 
	 */
	private SimpleDateFormat sdf;
	
	
	private long timeUpdate = 0;
	protected long timeWork=0;
    protected long timeStart=0;
	
    //Constructor
	public ColorPercentageWithWorker(List<File> images) {
		this.imagesPaths = images;
		sdf = new SimpleDateFormat("dd:MM:yyyy hh.mm.ss");
		threads = new ArrayList<>();
		consumerThreads = new ArrayList<>();
		producer = new LinkedBlockingQueue<>(imagesPaths);
		consumer = new LinkedBlockingQueue<>();
	}
	
	@Override
	protected Void doInBackground() throws Exception {
		
		System.out.println(sdf.format(new Date()) + "  " + NAME + " started.");
		
		// Number of CPUs Available
		int threadCount = Runtime.getRuntime().availableProcessors();
		
		for (int i = 0; i < threadCount; i++)
			threads.add(new ImageLoaderTask());
		
		//TODO get the number of devices 
		for (int i = 0; i < 1; i++)
			consumerThreads.add(new PercentageTask(i));
		
		int cntAll = imagesPaths.size();
		
		//Start the CPU threads To load Image Data. 
		for (int i = 0; i < threads.size(); i++)
			threads.get(i).start();
		
		//Start the GPU threads To Perform Calculations. 
		for (int i = 0; i < consumerThreads.size(); i++)
			consumerThreads.get(i).start();
		
		int cntCheck = 0;
		
		System.out.println(sdf.format(new Date()) + "  " + NAME + " CPU Devices used " + threads.size());
		System.out.println(sdf.format(new Date()) + "  " + NAME + " Images count " + cntAll);
		
		// Run until cancel is not called or Object Queue is not Finished and Thread Completed his Work. 
		while (!isCancelled()) {

			long cTime = System.currentTimeMillis();
			if ((cTime - timeUpdate) > 40) {
				cntCheck = cntAll - producer.size();
				setProgress((int) Math.round((double) cntCheck / cntAll * 100));
				if (cntCheck > 0) {
					long est = getEstTimeToCompleteWork(cntAll, cntCheck);
				}
				
				timeUpdate = cTime;
			}
			
			if (producer.size() == 0 && consumer.size() == 0 && threadsIdle())
				break;
			
			//System.out.println("Consumer Size: "+ consumer.size() +" Producer Size: "+ queueImages.size());
		}
		
		System.out.println(sdf.format(new Date()) + "  " + NAME + " Completed ");
		return null;
	}
	
	/** 
	 * Estimate time to complete the process. 
	 * 
	 * @param cntAll total images list size. 
	 * @param cntCheck Number of images on which operation is performed. 
	 * @return Estimated Time. 
	 */
	public long getEstTimeToCompleteWork(int cntAll, int cntCheck) {
		return (cntAll - cntCheck) * ((timeWork + (System.currentTimeMillis() - timeStart)) / cntCheck);
	}
	
	//Thread will be ideal if queue is empty 
	private boolean threadsIdle() {
		for (ImageLoaderTask cr : threads)
			if (cr.getCheckState() == 1)
				return false;
		return true;
	}
	
	@Override
	protected void done() {
		try {
			get();
		} catch (InterruptedException e) {
			e.printStackTrace();
		} catch (ExecutionException e) {
			e.printStackTrace();
		}
	}
	
	public class ImageLoaderTask extends Thread{
		// 0 - not started, 1 - process image, 2 -
		// idle, 3 - stopped, 4 paused
		private int checkState = 0;
		
		//pause the thread
		AtomicBoolean pause = new AtomicBoolean(false);
		
		//unload the thread
		AtomicBoolean unloaded = new AtomicBoolean(false);
		
		private ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
		
		public void unload() {
			this.unloaded.set(true);
		}

		private void setCheckState(int checkState) {
			lock.writeLock().lock();
			this.checkState = checkState;
			lock.writeLock().unlock();
		}

		public int getCheckState() {
			lock.readLock().lock();
			int st = checkState;
			lock.readLock().unlock();
			return st;
		}
		
		@Override
		public void run() {
			
			File image = null;
			setCheckState(2);
			
			while (!this.isInterrupted()) {
				
				// If thread is pause then continue do nothing. 
				if (pause.get()) {
					setCheckState(4);
					try {
						sleep(100);
					} catch (Exception e) {
					}
					continue;
				}
				
				// Thread unload is true then break and complete the thread work. 
				if (unloaded.get()) break;
				
				// get object from Queue. 
				image = null;
				lock.writeLock().lock();
				try {
					image = producer.poll(10, TimeUnit.MILLISECONDS);
				} catch (Exception e) {
					lock.writeLock().lock();
					System.err.println("Interrupted Image Color Detection check");
					break;
				}
				
				if (image == null) {
					checkState = 2;
					lock.writeLock().unlock();
					try {
						sleep(100);
					} catch (Exception e) {
					}
					continue;
				} else {
					checkState = 1;
					lock.writeLock().unlock();
				}
				
				ImageData im;
				try {
					im = new ImageData(getBufferedImage(image));
					consumer.put(im);
				} catch (IOException | InterruptedException | NullPointerException e) {
					//e.printStackTrace();
				}
				
				setCheckState(2);
			}
			
			// thread is out from loop and stopped. 
			setCheckState(3);
		}
		
	}
	
	public class PercentageTask extends Thread{
		// 0 - not started, 1 - process image, 2 -
		// idle, 3 - stopped, 4 paused
		private int checkState = 0;
		
		//pause the thread
		private AtomicBoolean pause = new AtomicBoolean(false);
		
		//unload the thread
		private AtomicBoolean unloaded = new AtomicBoolean(false);
		
		private ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
		
		private CUfunction kernelFunction;
		private int deviceNo;
		
		public PercentageTask(int deviceNo) {
			this.deviceNo = deviceNo;
		}
		
		public void unload() {
			this.unloaded.set(true);
		}

		private void setCheckState(int checkState) {
			lock.writeLock().lock();
			this.checkState = checkState;
			lock.writeLock().unlock();
		}

		public int getCheckState() {
			lock.readLock().lock();
			int st = checkState;
			lock.readLock().unlock();
			return st;
		}
		
		@Override
		public void run() {

			// Enable exceptions and omit all subsequent error checks
	        JCudaDriver.setExceptionsEnabled(true);
	        JNvrtc.setExceptionsEnabled(true);

	        // Initialize the driver and create a context for the first device.
	        cuInit(deviceNo);
	        CUdevice device = new CUdevice();
	        cuDeviceGet(device, deviceNo);
	        CUcontext context = new CUcontext();
	        cuCtxCreate(context, 0, device);
	        
			// Obtain the CUDA source code from the CUDA file
	        String cuFileName = "JCudaImageJColorPercentageKernel.cu";
	        String sourceCode = CudaUtils.readResourceAsString(cuFileName);
			if (sourceCode == null) {
				IJ.showMessage("Error", "Could not read the kernel source code");
			}
			
			kernelFunction = CudaUtils.createFunction(sourceCode, "convert_to_hsv");
			
			ImageData image = null;
			setCheckState(2);
			
			while (!this.isInterrupted()) {
				
				// If thread is pause then continue do nothing. 
				if (pause.get()) {
					setCheckState(4);
					try {
						sleep(100);
					} catch (Exception e) {
					}
					continue;
				}
				
				// Thread unload is true then break and complete the thread work. 
				if (unloaded.get()) break;
				
				// get object from Queue. 
				image = null;
				lock.writeLock().lock();
				try {
					image = consumer.poll(10, TimeUnit.MILLISECONDS);
				} catch (Exception e) {
					lock.writeLock().lock();
					System.err.println("Interrupted Image Color Detection check");
					break;
				}
				
				if (image == null) {
					checkState = 2;
					lock.writeLock().unlock();
					try {
						sleep(100);
					} catch (Exception e) {
					}
					continue;
				} else {
					checkState = 1;
					lock.writeLock().unlock();
				}
				
				execute(image.getPixcels(), image.getWidth(), image.getHeight());
				setCheckState(2);
				
				image.flush();
			}
			
			// thread is out from loop and stopped. 
			setCheckState(3);
		}
		
		public void execute(int pixels[], int w, int h) {
			// Allocate memory on the device, and copy the host data to the device
			int size = w * h * Sizeof.INT;
			CUdeviceptr pointer = new CUdeviceptr();
			cuMemAlloc(pointer, size);
			cuMemcpyHtoD(pointer, Pointer.to(pixels), size);

			int numElements = 9;
			int hostInputA[] = new int[numElements];
			for (int i = 0; i < numElements; i++) {
				hostInputA[i] = i;
			}

			// Allocate the device input data, and copy the
			// host input data to the device
			CUdeviceptr deviceInputA = new CUdeviceptr();
			cuMemAlloc(deviceInputA, numElements * Sizeof.INT);
			cuMemcpyHtoD(deviceInputA, Pointer.to(hostInputA), numElements * Sizeof.INT);

			// Allocate device output memory
			CUdeviceptr deviceOutput = new CUdeviceptr();
			cuMemAlloc(deviceOutput, numElements * Sizeof.INT);

			// Set up the kernel parameters: A pointer to an array
			// of pointers which point to the actual values.
			Pointer kernelParameters = Pointer.to(Pointer.to(pointer), Pointer.to(hostInputA), Pointer.to(new int[] { w }),
					Pointer.to(new int[] { h }), Pointer.to(deviceOutput));

			// Call the kernel function
			int blockSize = 16;
			int gridSize = (Math.max(w, h) + blockSize - 1) / blockSize;
			cuLaunchKernel(kernelFunction, gridSize, gridSize, 1, // Grid dimension
					blockSize, blockSize, 1, // Block dimension
					0, null, // Shared memory size and stream
					kernelParameters, null // Kernel- and extra parameters
			);
			cuCtxSynchronize();

			// Allocate host output memory and copy the device output
			// to the host.
			int hostOutput[] = new int[numElements];
			cuMemcpyDtoH(Pointer.to(hostOutput), deviceOutput, numElements * Sizeof.INT);

			// Clean up.
			cuMemFree(deviceInputA);
			cuMemFree(deviceOutput);
			cuMemFree(pointer);
		}
	}
	
	// Load the kernel File 
	public CUcontext loadKernal() {
		System.out.println(sdf.format(System.currentTimeMillis()) +" Loading Cuda Device and Code");
		
		// Enable exceptions and omit all subsequent error checks
        JCudaDriver.setExceptionsEnabled(true);
        JNvrtc.setExceptionsEnabled(true);

        // Initialize the driver and create a context for the first device.
        cuInit(0);
        CUdevice device = new CUdevice();
        cuDeviceGet(device, 0);
        CUcontext context = new CUcontext();
        cuCtxCreate(context, 0, device);
        
        // Create the kernel function
		return context;
	}
	
	public BufferedImage getBufferedImage(File file) throws IOException {
		BufferedImage image = ImageIO.read(file);
		int w = image.getWidth();
		int h = image.getHeight();
	    BufferedImage inputImage = new BufferedImage(w, h, BufferedImage.TYPE_INT_RGB);
		Graphics2D g = inputImage.createGraphics();
		g.drawImage(image, 0, 0, null);
		g.dispose();
		image.flush();
		image = null;
		return inputImage;
	}
	
}