//============================================================================
// Name        : DisplayImage.cpp
// Author      : Tomas Bravenec
// Version     :
// Copyright   : Your copyright notice
// Description : OpenCV test in C++
//============================================================================

#include <iostream>
#include <stdio.h>
#include <opencv2/opencv.hpp>
#include "CUDA_Func_Library.cuh"
#include "stopwatch.h"

using namespace cv;

typedef unsigned char BYTE;

void usage()
{
	printf("usage: DisplayImage.out <Image_Path> <method> <type> <adjustment>\n");
	printf("method:\n");
	printf("\t\"-E\" - Edge detection\n\t\"-T\" - Thresholding\n\t\"-G\" - Grayscale\n\t\"-K\" - K-means clustering\n");
	printf("type: (Edge Detection filter)\n");
	printf("\t\"-0\" - Sobel\n\t\"-1\" - Prewitt\n\t\"-2\" - Kirsch\n");
	printf("type: (Thresholding type)\n");
	printf("\t\"-0\" - Simple\n\t\"-1\" - Dual binary\n\t\"-2\" - Dual Gray\n");
	printf("adjustment:\n");
	printf("\t\"/I\" - Color Inversion\n\t\"/0.XX\" - Threshold (XX - range from 00 to 99\n");
	printf("\t\t  Can be used twice for Thresholding method, type 1 and 2\n");
	printf("\t\"/CXX\" - Cluster Count (XX - range from 2 to 99)\n");
	printf("\t\"/MXX\" - Maximum iterations (XX - range from 0 to 999)\n");
}

char device = 0;
char method = -1;
int type = -1;
bool EDGEthreshold = false;
bool negative = false;
double threshold1 = -1; // higher
double threshold2 = -1; // lower
int threads = 16;
int clusterCount = -1;
int maxIterations = -1;

int main(int argc, char** argv)
{
	// application usage
    if ( argc < 2 )
    {	
		usage();
        return -1;
    }
    
	// get settings from launch arguments
	if (argc > 2)
	{
		for (int i = 2; i < argc; i++)
		{
			if (argv[i][0] == '-')
			{
				switch (argv[i][1])
				{
					case 'e': case 'E': case 't': case 'T': case 'g': case 'G': case'k': case 'K':
						if(method == -1)
						{
							method = argv[i][1];
						}
						else
						{
							printf("Multiple use of parameter <method>\n\n");
							usage();
        					return -1;
						}
						break;
					case '0': case '1': case '2':
						if(type == -1)
						{
							type = argv[i][1] - 48;
						}
						else
						{
							printf("Multiple use of parameter <type>\n\n");
							usage();
        					return -1;
						}
						break;
					default:
						printf("Incorrect parameter <method> or <type>\n\n");
						usage();
        				return -1;
						break;
				}
			}
			else if (argv[i][0] == '/')
			{
				switch (argv[i][1])
				{
					case 'i': case 'I':
						negative = true;
						break;
					case 'T': case 't':
						EDGEthreshold = true;
						break;
					case 'c': case 'C':
						clusterCount = strtod(argv[i]+2, NULL);
						break;
					case 'm': case 'M':
						maxIterations = strtod(argv[i]+2, NULL);
						break;
					case '0':
						if(threshold1 == -1)
						{
							threshold1 = strtod(argv[i]+1, NULL);
							EDGEthreshold = true;
						}
						else if(threshold2 == -1)
						{
							threshold2 = strtod(argv[i]+1, NULL);
						}
						else
						{
							printf("Incorrect parameter <adjustment>\n\n");
							usage();
						}
						break;
					default:
						printf("Incorrect parameter <adjustment>\n\n");
						usage();
        				return -1;
						break;
				}
			}
		}
	}

	// parameter control
	double switchHelp;
	if(threshold1 < threshold2)
	{
		switchHelp = threshold1;
		threshold1 = threshold2;
		threshold2 = switchHelp;
	}

	// parameter control
	if((method == 'T' || method == 't') && (type == 0))
	{
		if(threshold1 == -1)
		{
			printf("Missing parameter <adjustment>\n\n");
			usage();
        	return -1;
		}
	}

	// parameter control
	if((method == 'T' || method == 't') && (type == 1 || type == 2))
	{
		if(threshold1 == -1 || threshold2 == -1)
		{
			printf("Missing parameter <adjustment>\n\n");
			usage();
        	return -1;
		}
	}

	// parameter control
	if ((method == 'E' || method == 'e') && (type == -1))
	{
		printf("Missing parameter <type>\n\n");
		usage();
        return -1;
	}

	// parameter control
	if ((method == 'k' || method == 'K'))
	{
		if (clusterCount < 2 || clusterCount > 99)
		{
			printf("Cluster count is out of allowed range!\n\n");
			usage();
        	return -1;
		}
		if (maxIterations == -1)
		{
			printf("Maximum iterations is not set. Using default 50.\n");
			maxIterations = 50;
		}
		if (maxIterations < 0 || maxIterations > 999)
		{
			printf("Maximum iterations is out of allowed range!\n\n");
			usage();
        	return -1;
		}
	}

    Mat image;
    image = imread( argv[1], 1 );

    if ( !image.data )
    {
        printf("No image data \n");
        return -1;
    }

    int width = image.cols, height = image.rows, stride = image.step;
    BYTE *rgbValues = image.data;
    BYTE *d_img, *d_gray, *d_output;
    cudaError_t cudaStatus;
    cudaDeviceProp prop;

    cudaGetDeviceProperties(&prop, device);

	// initialiye device
    cudaSetDevice(device);
    cudaFree(0);

    //
    // Start stopwatch
    //
    stopwatch sWatch;
    sWatch.start();

    int div = 1, offset = 0;
	unsigned long long bytes = stride * height;

	cudaKernelLauncher CUDA;
	int iterations;
    switch (method)
	{
	case 'e': case 'E':		// edge detection
		// count how many iterations have to be done to process image
		for (div = 1; div < height; div++)
		{
			if ((bytes / div) * 2 < prop.totalGlobalMem * 0.8)
			{
				break;
			}
		}
		// memory allocation for selected method
		cudaStatus = cudaMalloc(&d_img, stride * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}
		cudaStatus = cudaMalloc(&d_gray, width * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}
		cudaStatus = cudaMalloc(&d_output, width * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}

		// if image is too big, processing will be divided into cycles
		for (int i = 0; i < div; i++)
		{
			offset = stride * height / div * i;
			cudaStatus = cudaMemcpy(d_img, rgbValues + offset, stride * height / div * sizeof(BYTE), cudaMemcpyHostToDevice);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
			cudaStatus = CUDA.cudaRGB2Gray(d_img, d_gray, height / div, width, stride, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}

			cudaStatus = CUDA.cudaEdgeDetection(d_gray, d_output, height / div, width, type, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}

			// run tresholding, if neabled
			if (EDGEthreshold == true)
			{
				cudaStatus = CUDA.cudaThresholdEdge(d_output, height / div, width, threshold1, threads);
				if (cudaStatus != cudaSuccess)
				{
					goto cancel_CUDA;
				}
			}

			// run conversion to negative, if enabled
			if (negative == true)
			{
				cudaStatus = CUDA.cudaNegative(d_output, height / div, width, threads);
				if (cudaStatus != cudaSuccess)
				{
					goto cancel_CUDA;
				}
			}

			cudaStatus = CUDA.cudaGray2RGB(d_img, d_output, height / div, width, stride, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
			cudaStatus = cudaMemcpy(rgbValues + offset, d_img, stride * height / div * sizeof(BYTE), cudaMemcpyDeviceToHost);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
		}
		break;
	case 't': case 'T':		// thresholding
		// count how many iterations have to be done to process image
		for (div = 1; div < height; div++)
		{
			if ((bytes / div) * 2 < prop.totalGlobalMem * 0.8)
			{
				break;
			}
		}
		// memory allocation for selected method
		cudaStatus = cudaMalloc(&d_img, stride * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}
		cudaStatus = cudaMalloc(&d_output, width * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}

		// if image is too big, processing will be divided into cycles
		for (int i = 0; i < div; i++)
		{
			offset = stride * height / div * i;
			cudaStatus = cudaMemcpy(d_img, rgbValues + offset, stride * height / div * sizeof(BYTE), cudaMemcpyHostToDevice);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
			cudaStatus = CUDA.cudaRGB2Gray(d_img, d_output, height / div, width, stride, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}

			cudaStatus = CUDA.cudaThreshold(d_output, height / div, width, threshold1, threshold2, type, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}

			// run conversion to negative, if enabled
			if (negative == true)
			{
				cudaStatus = CUDA.cudaNegative(d_output, height / div, width, threads);
				if (cudaStatus != cudaSuccess)
				{
					goto cancel_CUDA;
				}
			}

			cudaStatus = CUDA.cudaGray2RGB(d_img, d_output, height / div, width, stride, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
			cudaStatus = cudaMemcpy(rgbValues + offset, d_img, stride * height / div * sizeof(BYTE), cudaMemcpyDeviceToHost);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
		}
		break;
	case 'k': case 'K':	// kmeans clustering
		// count how many iterations have to be done to process image
		for (div = 1; div < height; div++)
		{
			if ((bytes / div) * 4 < prop.totalGlobalMem * 0.8)
			{
				break;
			}
		}
		cudaStatus = cudaMalloc(&d_img, stride * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}
		cudaStatus = cudaMalloc(&d_output, width * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}

		cudaStatus = CUDA.cudaKmeans(rgbValues, div, clusterCount, maxIterations, &iterations, height, width, stride, threads);
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}

		break;
	case 'g': case 'G':	// grayscale
		// count how many iterations have to be done to process image
		for (div = 1; div < height; div++)
		{
			if ((bytes / div) * 2 < prop.totalGlobalMem * 0.8)
			{
				break;
			}
		}
		// memory allocation for selected method
		cudaStatus = cudaMalloc(&d_img, stride * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}
		cudaStatus = cudaMalloc(&d_output, width * height / div * sizeof(BYTE));
		if (cudaStatus != cudaSuccess)
		{
			goto cancel_CUDA;
		}

		// if image is too big, processing will be divided into cycles
		for (int i = 0; i < div; i++)
		{
			offset = stride * height / div * i;
			cudaStatus = cudaMemcpy(d_img, rgbValues + offset, stride * height / div * sizeof(BYTE), cudaMemcpyHostToDevice);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
			cudaStatus = CUDA.cudaRGB2Gray(d_img, d_output, height / div, width, stride, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}

			// run conversion to negative, if enabled
			if (negative == true)
			{
				cudaStatus = CUDA.cudaNegative(d_output, height / div, width, threads);
				if (cudaStatus != cudaSuccess)
				{
					goto cancel_CUDA;
				}
			}

			cudaStatus = CUDA.cudaGray2RGB(d_img, d_output, height / div, width, stride, threads);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
			cudaStatus = cudaMemcpy(rgbValues + offset, d_img, stride * height / div * sizeof(BYTE), cudaMemcpyDeviceToHost);
			if (cudaStatus != cudaSuccess)
			{
				goto cancel_CUDA;
			}
		}
		break;
	default:
		break;
	}

cancel_CUDA:
	// in case of error, print error code and quit application
	if (cudaStatus != cudaSuccess)
	{
		printf("Program exited with CUDA Error code %d\n", cudaStatus);
		return -1;
	}

	// free GPU memory
    cudaFree(d_img);
    cudaFree(d_gray);
	cudaFree(d_output);

	// stop stopwatch and print time into console
    sWatch.stop();
    printf("Processing time: %2.3f s\n", sWatch.getEllapsedTime());
	if (method == 'k' || method == 'K')
	{
		printf("Iterations: %d\n", iterations);
	}

	// deinitialiye CUDA device
    cudaDeviceReset();

	imwrite("segmented.bmp", image);

	// Show output image
    namedWindow("Display Image", WINDOW_NORMAL );
    imshow("Display Image", image);

    waitKey(0);
    destroyAllWindows();
    
    return 0;
}