#include <chrono>
#include <iostream>
#include <mutex>
#include <math.h>
#include <string>
#include <string.h>
#include <thread>
#include <vector>

#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION

#include "stb_image.h"
#include "stb_image_write.h"
#include "getopt.h"

#define DEFAULT_THREAD_COUNT 4
#define DEFAULT_PIXEL_COUNT 1000
#define DEFAULT_AVG_PARAM 3

using std::cout;
using std::endl;
using std::mutex;
using std::string;
using std::thread;
using std::vector;

using namespace std::chrono;

typedef union img_func_param
{
	int i_param;
	float f_param;
} img_func_param;

// int w, int h, int x, int y, int channel_count, uint8_t *src, uint8_t *dst, img_func_param *param
typedef void (*img_func_ptr) (int, int, int, int, int, uint8_t*, uint8_t*, img_func_param*);

typedef struct thread_param
{
	int w;
	int h;
	int channel_count;
	int pixel_count;
	uint64_t *curr_pixel;
	std::mutex mutex;
	uint8_t *src;
	uint8_t *dst;
	img_func_ptr func;
	img_func_param *param;
} thread_param;

void average(int w, int h, int x, int y, int channel_count, uint8_t *src, uint8_t *dst, img_func_param *param)
{
	uint64_t pixel_idx;
	uint32_t value[4];
	int curr_x;
	int curr_y;
	int pixel_count;
	int kernel_size;
	int i;

	// Kernel size is positive odd integer
	kernel_size = param->i_param;
	if (kernel_size <= 0)
	{
		return;
	}

	// Round up
	if (kernel_size % 2 == 0)
	{
		++kernel_size;
	}

	pixel_count = 0;
	for (i = 0; i < 4; ++i)
	{
		value[i] = 0;
	}

	for (curr_y = y - (kernel_size / 2); curr_y <= y + (kernel_size / 2); ++curr_y)
	{
		if ((curr_y < 0) || (curr_y >= h))
		{
			continue;
		}
		for (curr_x = x - (kernel_size / 2); curr_x <= x + (kernel_size / 2); ++curr_x)
		{
			if ((curr_x < 0) || (curr_x >= w))
			{
				continue;
			}

			// Process kernel pixel
			pixel_idx = (curr_y * w + curr_x) * channel_count;
			for (i = 0; i < channel_count; ++i)
			{
				value[i] += src[pixel_idx + i];
			}
			++pixel_count;
		}
	}

	if (pixel_count == 0)
	{
		return;
	}

	pixel_idx = (y * w + x) * channel_count;
	for (i = 0; i < channel_count; ++i)
	{
		dst[pixel_idx + i] = value[i] / pixel_count;
	}
}

void laplacian(int w, int h, int x, int y, int channel_count, uint8_t *src, uint8_t *dst, img_func_param *param)
{
	uint64_t pixel_idx;
	int32_t value[4];
	int curr_x;
	int curr_y;
	int i;
	int value_matrix[3][3];
	int work_channel_count;
	bool copy_alpha;

	for (i = 0; i < 4; ++i)
	{
		value[i] = 0;
	}

	copy_alpha = false;
	work_channel_count = channel_count;
	if (channel_count %2 == 0)
	{
		--work_channel_count;
		copy_alpha = true;
	}

	// Laplacian weight matrix
	value_matrix[0][0] = 0;
	value_matrix[0][1] = -1;
	value_matrix[0][2] = 0;
	value_matrix[1][0] = -1;
	value_matrix[1][1] = 4;
	value_matrix[1][2] = -1;
	value_matrix[2][0] = 0;
	value_matrix[2][1] = -1;
	value_matrix[2][2] = 0;

	for (curr_y = y - 1; curr_y <= y + 1; ++curr_y)
	{
		if ((curr_y < 0) || (curr_y >= h))
		{
			continue;
		}
		for (curr_x = x - 1; curr_x <= x + 1; ++curr_x)
		{
			if ((curr_x < 0) || (curr_x >= w))
			{
				continue;
			}

			// Process kernel pixel
			pixel_idx = (curr_y * w + curr_x) * channel_count;
			for (i = 0; i < work_channel_count; ++i)
			{
				int mod_value;

				mod_value = src[pixel_idx + i];
				mod_value *= value_matrix[curr_x + 1 - x][curr_y + 1 - y];
				value[i] += mod_value;
			}
		}
	}

	pixel_idx = (y * w + x) * channel_count;
	for (i = 0; i < work_channel_count; ++i)
	{
		if (value[i] < 0)
		{
			value[i] = 0;
		}

		if (value[i] > 255)
		{
			value[i] = 255;
		}

		dst[pixel_idx + i] = value[i];
	}

	if (copy_alpha)
	{
		dst[pixel_idx + work_channel_count] = src[pixel_idx + work_channel_count];
	}
}

void brightness(int w, int h, int x, int y, int channel_count, uint8_t *src, uint8_t *dst, img_func_param *param)
{
	uint64_t pixel_idx;
	int i;
	int value_i;
	float value_f;
	float multiplier;
	bool copy_alpha;
	int work_channel_count;

	copy_alpha = false;
	work_channel_count = channel_count;
	if (channel_count %2 == 0)
	{
		--work_channel_count;
		copy_alpha = true;
	}

	multiplier = param->f_param;
	if (multiplier <= 0.0)
	{
		return;
	}

	pixel_idx = (y * w + x) * channel_count;
	for (i = 0; i < work_channel_count; ++i)
	{
		value_f = src[pixel_idx + i];
		value_f *= multiplier;
		value_i = value_f;

		if (value_i < 0)
		{
			value_i = 0;
		}

		if (value_i > 255)
		{
			value_i = 255;
		}

		dst[pixel_idx + i] = value_i;
	}

	if (copy_alpha)
	{
		dst[pixel_idx + work_channel_count] = src[pixel_idx + work_channel_count];
	}
}

void sobel(int w, int h, int x, int y, int channel_count, uint8_t *src, uint8_t *dst, img_func_param *param)
{
	uint64_t pixel_idx;
	int32_t value_x[4];
	int32_t value_y[4];
	int curr_x;
	int curr_y;
	int i;
	int value_matrix_x[3][3];
	int value_matrix_y[3][3];
	bool copy_alpha;
	int work_channel_count;

	for (i = 0; i < 4; ++i)
	{
		value_x[i] = 0;
		value_y[i] = 0;
	}

	copy_alpha = false;
	work_channel_count = channel_count;
	if (channel_count % 2 == 0)
	{
		--work_channel_count;
		copy_alpha = true;
	}

	// Gx weight matrix
	value_matrix_x[0][0] = 1;
	value_matrix_x[0][1] = 0;
	value_matrix_x[0][2] = -1;
	value_matrix_x[1][0] = 2;
	value_matrix_x[1][1] = 0;
	value_matrix_x[1][2] = -2;
	value_matrix_x[2][0] = 1;
	value_matrix_x[2][1] = 0;
	value_matrix_x[2][2] = -1;

	// Gy weight matrix
	value_matrix_y[0][0] = 1;
	value_matrix_y[0][1] = 2;
	value_matrix_y[0][2] = 1;
	value_matrix_y[1][0] = 0;
	value_matrix_y[1][1] = 0;
	value_matrix_y[1][2] = 0;
	value_matrix_y[2][0] = -1;
	value_matrix_y[2][1] = -2;
	value_matrix_y[2][2] = -1;

	for (curr_y = y - 1; curr_y <= y + 1; ++curr_y)
	{
		if ((curr_y < 0) || (curr_y >= h))
		{
			continue;
		}
		for (curr_x = x - 1; curr_x <= x + 1; ++curr_x)
		{
			if ((curr_x < 0) || (curr_x >= w))
			{
				continue;
			}

			// Process kernel pixel
			pixel_idx = (curr_y * w + curr_x) * channel_count;
			for (i = 0; i < work_channel_count; ++i)
			{
				int mod_value_x;
				int mod_value_y;

				mod_value_x = src[pixel_idx + i];
				mod_value_x *= value_matrix_x[curr_x + 1 - x][curr_y + 1 - y];

				mod_value_y = src[pixel_idx + i];
				mod_value_y *= value_matrix_y[curr_x + 1 - x][curr_y + 1 - y];

				value_x[i] += mod_value_x;
				value_y[i] += mod_value_y;
			}
		}
	}

	pixel_idx = (y * w + x) * channel_count;
	for (i = 0; i < work_channel_count; ++i)
	{
		int value;

		value = sqrt(value_x[i] * value_x[i] + value_y[i] * value_y[i]);

		if (value > 255)
		{
			value = 255;
		}

		dst[pixel_idx + i] = value;
	}

	if (copy_alpha)
	{
		dst[pixel_idx + work_channel_count] = src[pixel_idx + work_channel_count];
	}
}

void grayscale(int w, int h, int x, int y, int channel_count, uint8_t *src, uint8_t *dst, img_func_param *param)
{
	uint64_t pixel_idx;
	uint64_t out_pixel_idx;
	int i;
	bool copy_alpha;
	int work_channel_count;
	int out_channel_count;
	float value;
	float magnitudes[3] = {0.299, 0.587, 0.114};
	// BT.601

	copy_alpha = false;
	work_channel_count = channel_count;
	if (channel_count %2 == 0)
	{
		--work_channel_count;
		copy_alpha = true;
	}

	out_channel_count = channel_count;
	if (channel_count > 2)
	{
		out_channel_count -= 2;
	}

	pixel_idx = (y * w + x) * channel_count;
	out_pixel_idx = (y * w + x) * out_channel_count;

	if (work_channel_count > 1)
	{
		value = 0;
		for (int i = 0; i < work_channel_count; ++i)
		{
			value += magnitudes[i] * src[pixel_idx + i];
		}
		if (value > 255)
		{
			value = 255;
		}
		dst[out_pixel_idx] = value;
	}
	else
	{
		dst[out_pixel_idx] = src[pixel_idx];
	}

	if (copy_alpha)
	{
		dst[out_pixel_idx + out_channel_count - 1] = src[pixel_idx + work_channel_count];
	}
}

void usage(char *name)
{
	cout << "Usage:" << endl;
	cout << name << " [options] <input>" << endl;
	cout << "Image types: PNG, JPG, BMP, TGA" << endl << endl;

	cout << "Option list:" << endl;
	cout << "\tn: Number of threads (Default: " << DEFAULT_THREAD_COUNT << ")" << endl;
	cout << "\tx: Pixels per thread before syncronization (Default: " << DEFAULT_PIXEL_COUNT << ")" << endl;
	cout << "\tf: Choose filter" << endl;
	cout << "\tp: Filter parameter" << endl << endl;

	cout << "Filter list:" << endl;
	cout << "\taverage: Average filter. Parameter is aperture size (Default: " << DEFAULT_AVG_PARAM << ")" << endl;
	cout << "\tlaplacian: Laplacian filter" << endl;
	cout << "\tbrightness: Increase or decrease brightness via parameter" << endl;
	cout << "\tsobel: Sobel edge detection" << endl;
	cout << "\tgrayscale: Turn image to grayscale" << endl;
	// cout << "\tcolour: Extract most prominent colour and most prominent complementary colour" << endl;
}

void extract_filename(const string& filepath, string& path, string& extension)
{
	int path_idx;
	int ext_idx;

	path_idx = filepath.find_last_of("/\\") + 1;
	path = filepath.substr(0, path_idx);

	ext_idx = filepath.find_last_of('.', path_idx);
	extension = filepath.substr(ext_idx + 1);
}

string& to_lowercase(string& str)
{
	size_t size = str.size();
	size_t i;

	for (i = 0; i < size; ++i)
	{
		if ((str[i] >= 'A') && (str[i] <= 'Z'))
		{
			str[i] += 'a' - 'A';
		}
	}

	return str;
}

void thread_func(thread_param *param, int id)
{
	uint64_t start_pixel;
	uint64_t total_pixel_count;
	uint64_t total_pixel_worked;
	int work_pixel_count;
	int i;
	int x;
	int y;

	start_pixel = 0;
	total_pixel_count = param->w;
	total_pixel_count *= param->h;
	total_pixel_worked = 0;

	while (*(param->curr_pixel) < total_pixel_count)
	{
		param->mutex.lock();
		start_pixel = *(param->curr_pixel);
		if (*(param->curr_pixel) < total_pixel_count)
		{
			work_pixel_count = param->pixel_count;

			if ((total_pixel_count - *(param->curr_pixel)) < work_pixel_count)
			{
				work_pixel_count = total_pixel_count - *(param->curr_pixel);
			}
			*(param->curr_pixel) += work_pixel_count;
		}
		else
		{
			work_pixel_count = 0;
		}
		param->mutex.unlock();

		for (i = 0; i < work_pixel_count; ++i)
		{
			x = (start_pixel + i) % param->w;
			y = (start_pixel + i) / param->w;
			param->func(param->w, param->h, x, y, param->channel_count, param->src, param->dst, param->param);
		}
		total_pixel_worked += work_pixel_count;
	}

	// param->mutex.lock();
	// cout << "Thread " << id << " worked on " << total_pixel_worked << " pixels." << endl;
	// param->mutex.unlock();
}

int strcmp_autocomplete(const string& str1, const string& str2)
{
	int min_length = str1.size();

	if (str2.size() < min_length)
	{
		min_length = str2.size();
	}
	return strncmp(str1.c_str(), str2.c_str(), min_length);
}

img_func_ptr str_to_func(string& str)
{
	to_lowercase(str);

	if (0 == strcmp_autocomplete(str, "average"))
	{
		str = "average";
		return average;
	}

	if (0 == strcmp_autocomplete(str, "laplacian"))
	{
		str = "laplacian";
		return laplacian;
	}

	if (0 == strcmp_autocomplete(str, "brightness"))
	{
		str = "brightness";
		return brightness;
	}

	if (0 == strcmp_autocomplete(str, "sobel"))
	{
		str = "sobel";
		return sobel;
	}

	if (0 == strcmp_autocomplete(str, "grayscale"))
	{
		str = "grayscale";
		return grayscale;
	}

	return NULL;
}

void print_execution_time(time_point<high_resolution_clock> start, time_point<high_resolution_clock> stop)
{
	microseconds us;
	us = duration_cast<microseconds>(stop - start);

	cout << "Execution took " << us.count() << "us" << endl;
}

int main(int argc, char **argv)
{
	thread_param parameters;
	img_func_param func_param;
	uint64_t pixel;
	int thread_count;
	string func_param_str;
	string filter_str;
	string filename_src;
	string filename_dst;
	string path;
	string ext;
	vector<thread*> threads;
	int opt;
	int error;
	int out_channel_count;
	time_point<high_resolution_clock> t1, t2;

	thread_count = DEFAULT_THREAD_COUNT;
	parameters.pixel_count = DEFAULT_PIXEL_COUNT;
	pixel = 0;
	parameters.curr_pixel = &pixel;
	parameters.param = &func_param;

	// Process options
	opt = getopt(argc, argv, "n:x:f:p:");
	while (opt != -1)
	{
		switch (opt)
		{
		case 'n':
			thread_count = strtol(optarg, NULL, 0);
			break;

		case 'x':
			parameters.pixel_count = strtol(optarg, NULL, 0);
			break;

		case 'f':
			filter_str = optarg;
			break;

		case 'p':
			func_param_str = optarg;
			break;

		default:
			cout << "Unknown option '" << (char)opt << "'" << endl;
			usage(argv[0]);
			return -1;
		}
		opt = getopt(argc, argv, "n:x:f:p:");
	}

	if (optind >= argc)
	{
		cout << "No filename specified" << endl;
		usage(argv[0]);
		return -1;
	}
	filename_src = argv[optind];

	if (filter_str.empty())
	{
		cout << "No filter specified" << endl;
		usage(argv[0]);
		return -1;
	}

	parameters.func = str_to_func(filter_str);
	if (NULL == parameters.func)
	{
		cout << "Filter specified does not match" << endl;
		usage(argv[0]);
		return 0;
	}

	parameters.param->i_param = 0;
	if (filter_str == "average")
	{
		parameters.param->i_param = DEFAULT_AVG_PARAM;
		if (!func_param_str.empty())
		{
			parameters.param->i_param = strtol(func_param_str.c_str(), NULL, 0);
		}
	}

	if (filter_str == "brightness")
	{
		parameters.param->f_param = 1.0;
		if (!func_param_str.empty())
		{
			parameters.param->f_param = strtof(func_param_str.c_str(), NULL);
		}
	}

	// Load Input Image
	parameters.src = stbi_load(filename_src.c_str(), &(parameters.w), &(parameters.h), &(parameters.channel_count), 0);
	if (NULL == parameters.src)
	{
		cout << "Error opening file" << endl;
		return -1;
	}

	out_channel_count = parameters.channel_count;

	if ((filter_str == "grayscale") && (parameters.channel_count > 2))
	{
		out_channel_count = parameters.channel_count - 2;
	}

	// Allocate Output Image
	parameters.dst = new uint8_t[parameters.w * parameters.h * out_channel_count];
	if (NULL == parameters.dst)
	{
		cout << "Error allocating memory for output image" << endl;
		stbi_image_free(parameters.src);
		return -1;
	}

	extract_filename(filename_src, path, ext);
	filename_dst = path + filter_str + ".png";

	t1 = high_resolution_clock::now();

	// Execute Threads
	for (int i = 0; i < thread_count; ++i)
	{
		thread* t;
		t = new thread(thread_func, &parameters, i + 1);
		if (t != NULL)
		{
			threads.push_back(t);
		}
		else
		{
			cout << "Could not create thread " << i + 1 << endl;
		}
	}

	// Join Threads
	for (auto it = threads.begin(); it != threads.end(); ++it)
	{
		(*it)->join();
		delete (*it);
	}

	t2 = high_resolution_clock::now();

	print_execution_time(t1, t2);

	stbi_image_free(parameters.src);
	error = stbi_write_png(filename_dst.c_str(), parameters.w, parameters.h, out_channel_count,
		parameters.dst, parameters.w * out_channel_count);
	delete[] parameters.dst;
	if (0 == error)
	{
		cout << "Could not write file" << endl;
		return -1;
	}

	return 0;
}