#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include <Windows.h>

typedef uint32_t COLOR32;

/*
convertRGBToYUV

Description:
	Convert an RGB color to YUV

Parameters:
	r		Red channel
	g		Green channel
	b		Blue channel
	y		Output Y channel
	u		Output U channel
	v		Output V channel
*/
void convertRGBToYUV(int r, int g, int b, int *y, int *u, int *v) {
	*y = (int) (0.2990 * r + 0.5870 * g + 0.1140 * b);
	*u = (int) (-0.1684 * r - 0.3316 * g + 0.5000 * b);
	*v = (int) (0.5000 * r - 0.4187 * g - 0.0813 * b);
}

/*
writeBitmap

Description:
	Write pixel data to a 32-bit bitmap file at a given path.

Parameters:
	px		Input pixel array
	width	Image width
	height	Image height
	path	Output image path

*/
void writeBitmap(COLOR32 *px, int width, int height, const char *path) {

	BITMAPFILEHEADER fh = { 0 };
	fh.bfType = 'B' | ('M' << 8);
	fh.bfSize = width * height * 4 + sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
	fh.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);

	BITMAPINFOHEADER hdr = { 0 };
	hdr.biSize = sizeof(hdr);
	hdr.biWidth = width;
	hdr.biHeight = height;
	hdr.biPlanes = 1;
	hdr.biBitCount = 32;
	hdr.biCompression = BI_RGB;

	FILE *fp = fopen(path, "wb");
	fwrite(&fh, sizeof(fh), 1, fp);
	fwrite(&hdr, sizeof(hdr), 1, fp);
	for (int y = height - 1; y >= 0; y--) {
		for (int x = 0; x < width; x++) {
			COLOR32 c = px[x + y * width];
			c = (c & 0xFF00FF00) | ((c >> 16) & 0xFF) | ((c & 0xFF) << 16);
			fwrite(&c, sizeof(c), 1, fp);
		}
	}
	fclose(fp);
}

/*
bezierInterpolate

Description:
	Interpolate a bezier curve from two given endpoints and a control point.

Parameters:
	x1		Endpoint 1 X
	y1		Endpoint 1 Y
	x2		Endpoint 2 X
	y2		Endpoint 2 Y
	cx		Control point X
	cy		Control point Y
	n		Step number along the bezier curve
	nMax	Numer of steps on the curve
	outX	Output X coordinate
	outY	Output Y coordinate
*/
void bezierInterpolate(int x1, int y1, int x2, int y2, int cx, int cy, int n, int nMax, int *outX, int *outY) {
	int nm2 = nMax * nMax;
	int nn = nMax - n;
	int lx1 = n * cx + nn * x1;
	int ly1 = n * cy + nn * y1;
	int lx2 = n * x2 + nn * cx;
	int ly2 = n * y2 + nn * cy;
	*outX = (lx1 * nn + lx2 * n) / nm2;
	*outY = (ly1 * nn + ly2 * n) / nm2;
}

/*
drawPixel

Description:
	Draw a pixel on a given image

Parameters:
	px		Image pixel array
	width	Image width
	height	Image height
	x		Pixel X
	y		Pixel Y
	c		Color to draw the pixel
*/
void drawPixel(COLOR32 *px, int width, int height, int x, int y, COLOR32 c) {
	if (x < 0 || y < 0) return;
	if (x >= width || y >= height) return;
	px[x + y * width] = c;
}

/*
drawLine

Description:
	Draw a line between two points with a given color

Parameters:
	px		Image pixel array
	width	Image width
	height	Image height
	x1		Endpoint 1 X
	y1		Endpoint 1 Y
	x2		Endpoint 2 X
	y2		Endpoint 2 Y
	color	Color of the line
*/
void drawLine(COLOR32 *px, int width, int height, int x1, int y1, int x2, int y2, COLOR32 color) {
	//find longest range
	int xRange = x2 - x1;
	if (xRange < 0) xRange = -xRange;
	int yRange = y2 - y1;
	if (yRange < 0) yRange = -yRange;

	//longest range gets iterated over
	if (xRange || yRange) {
		if (xRange > yRange) {
			int xStart = min(x1, x2);
			int xEnd = max(x1, x2);
			int yStart = xStart == x1 ? y1 : y2;
			int yEnd = xStart == x1 ? y2 : y1;
			yRange = yEnd - yStart;

			for (int x = xStart; x <= xEnd; x++) {
				int y = ((x - xStart) * yRange + xRange / 2) / xRange + yStart;
				drawPixel(px, width, height, x, y, color);
			}
		} else {
			int yStart = min(y1, y2);
			int yEnd = max(y1, y2);
			int xStart = yStart == y1 ? x1 : x2;
			int xEnd = yStart == y1 ? x2 : x1;
			xRange = xEnd - xStart;

			for (int y = yStart; y <= yEnd; y++) {
				int x = ((y - yStart) * xRange + yRange / 2) / yRange + xStart;
				drawPixel(px, width, height, x, y, color);
			}
		}
	}
}

/*
drawBezier

Description:
	px		Image pixel array
	width	Image width
	height	Image height
	x1		Endpoint 1 X
	y1		Endpoint 1 Y
	x2		Endpoint 2 X
	y2		Endpoint 2 Y
	cx		Control point X
	cy		Control point Y
	c		Color of the bezier curve
*/
void drawBezier(COLOR32 *px, int width, int height, int x1, int y1, int x2, int y2, int cx, int cy, COLOR32 c) {
	int nSegs = 16;
	for (int i = 0; i < nSegs; i++) {
		int px1, py1, px2, py2;
		bezierInterpolate(x1, y1, x2, y2, cx, cy, i, nSegs, &px1, &py1);
		bezierInterpolate(x1, y1, x2, y2, cx, cy, i + 1, nSegs, &px2, &py2);
		drawLine(px, width, height, px1, py1, px2, py2, c);
	}
}

/*
drawRandomBezier

Description:
	Draw a bezier curve onto the destination image with random endpoints and
	a random color.

Parameters:
	px		Image pixel array
	width	Image width
	height	Image height
*/
void drawRandomBezier(COLOR32 *px, int width, int height, int *px1, int *py1, int *px2, int *py2, int *pcx, int *pcy, COLOR32 *pcol) {
	COLOR32 c = (rand() & 0xFF) | ((rand() & 0xFF) << 8) | ((rand() & 0xFF) << 16);
	int x1 = rand() % width, y1 = rand() % height;
	int x2 = rand() % width, y2 = rand() % height;
	int cx = rand() % width, cy = rand() % height;

	drawBezier(px, width, height, x1, y1, x2, y2, cx, cy, c);
	*px1 = x1; *py1 = y1;
	*px2 = x2; *py2 = y2;
	*pcx = cx; *pcy = cy;
	*pcol = c;
}

/*
colorDifference

Description:
	Compute the difference between two RGB colors

Parameters:
	r1		Color 1 red
	g1		Color 1 green
	b1		Colro 1 blue
	r2		Color 2 red
	g2		Color 2 green
	b2		Color 2 blue
*/
int __inline colorDifference(int r1, int g1, int b1, int r2, int g2, int b2) {
	int dr = r2 - r1;
	int dg = g2 - g1;
	int db = b2 - b1;
	int dy, du, dv;
	convertRGBToYUV(dr, dg, db, &dy, &du, &dv);
	return 4 * dy * dy + du * du + dv * dv;
}

/*
unpackAddColor

Description:
	Add an RGB color to a running total

Parameters:
	c		Color to add
	r		Destination red total
	g		Destination green total
	b		Destination blue total
*/
void __inline unpackAddColor(COLOR32 c, int *r, int *g, int *b) {
	*r += c & 0xFF;
	*g += (c >> 8) & 0xFF;
	*b += (c >> 16) & 0xFF;
}

/*
imageDifference

Description:
	Compute the difference between two images.

Parameters:
	px1		Image 1 pixel array
	px2		Image 2 pixel array
	width	Image width
	height	Image height
	maxDiff	Maximum difference to count up to
*/
double imageDifference(COLOR32 *px1, COLOR32 *px2, int width, int height, double maxDiff) {
	int nPx = width * height;
	double diff = 0;
	for (int i = 0; i < nPx; i++) {
		int x = i % width;
		if (x == 0 || x == width - 1) continue;
		int y = i / width;
		if (y == 0 || y == height - 1) continue;

		COLOR32 cc1 = px1[i];
		COLOR32 cc2 = px2[i];
		COLOR32 cu1 = px1[i - width];
		COLOR32 cu2 = px2[i - width];
		COLOR32 cd1 = px1[i + width];
		COLOR32 cd2 = px2[i + width];
		COLOR32 cl1 = px1[i - 1];
		COLOR32 cl2 = px2[i - 1];
		COLOR32 cr1 = px1[i + 1];
		COLOR32 cr2 = px2[i + 1];
		int r1 = 0, g1 = 0, b1 = 0;
		int r2 = 0, g2 = 0, b2 = 0;

		unpackAddColor(cc1, &r1, &g1, &b1);
		unpackAddColor(cu1, &r1, &g1, &b1);
		unpackAddColor(cd1, &r1, &g1, &b1);
		unpackAddColor(cl1, &r1, &g1, &b1);
		unpackAddColor(cr1, &r1, &g1, &b1);

		unpackAddColor(cc2, &r2, &g2, &b2);
		unpackAddColor(cu2, &r2, &g2, &b2);
		unpackAddColor(cd2, &r2, &g2, &b2);
		unpackAddColor(cl2, &r2, &g2, &b2);
		unpackAddColor(cr2, &r2, &g2, &b2);

		r1 /= 5;
		g1 /= 5;
		b1 /= 5;
		r2 /= 5;
		g2 /= 5;
		b2 /= 5;
		
		diff += colorDifference(r1, g1, b1, r2, g2, b2);
		if (diff >= maxDiff) {
			return maxDiff;
		}
	}
	return diff;
}

/*
readBitmap24

Description:
	Read a 24-bit bitmap image to a pixel array

Parameters:
	path	File path of the input bitmap
	pWidth	Output image width
	pHeight	Output image height
*/
COLOR32 *readBitmap24(const char *path, int *pWidth, int *pHeight) {
	FILE *fp = fopen(path, "rb");
	if (fp == NULL) {
		*pWidth = *pHeight = 0;
		return NULL;
	}
	BITMAPFILEHEADER fh;
	BITMAPINFOHEADER bmi;
	fread(&fh, sizeof(fh), 1, fp);
	fread(&bmi, sizeof(bmi), 1, fp);
	if (bmi.biBitCount != 24) {
		*pWidth = *pHeight = 0;
		return NULL;
	}
	*pWidth = bmi.biWidth;
	*pHeight = bmi.biHeight;

	int width = bmi.biWidth, height = bmi.biHeight;
	int stride = (width * 3 + 3) & ~3;
	COLOR32 *px = (COLOR32 *) calloc(width * height, sizeof(COLOR32));
	unsigned char *row = (unsigned char *) calloc(stride, 1);
	for (int y = height - 1; y >= 0; y--) {
		fread(row, stride, 1, fp);
		COLOR32 *rDest = px + y * width;
		for (int i = 0; i < width; i++) {
			rDest[i] = (row[i * 3 + 0] << 16) | (row[i * 3 + 1] << 8) | (row[i * 3 + 2] << 0);
		}
	}
	free(row);
	fclose(fp);
	return px;
}

/*
allocImage

Description:
	Allocate a pixel array for an image

Parameters:
	width	Width of the image
	height	Height of the image
*/
COLOR32 *allocImage(int width, int height) {
	return (COLOR32 *) calloc(width * height, sizeof(COLOR32));
}

/*
imageAverage

Description:
	Compute the average color of an image

Parameters:
	px		Image pixel array
	width	Image width
	height	Image height
*/
COLOR32 imageAverage(COLOR32 *px, int width, int height) {
	int tr = 0, tg = 0, tb = 0;
	int nPx = width * height;
	for (int i = 0; i < nPx; i++) {
		COLOR32 c = px[i];
		tr += c & 0xFF;
		tg += (c >> 8) & 0xFF;
		tb += (c >> 16) & 0xFF;
	}
	tr /= nPx;
	tg /= nPx;
	tb /= nPx;
	return tr | (tg << 8) | (tb << 16);
}

int g_kill = 0;

/*
killerCallback

Description:
	Called in another thread that responds to keyboard input to stop the main
	thread computing

Parameters:
	lpParam	Unused
*/
DWORD CALLBACK killerCallback(LPVOID lpParam) {
	while (1) {
		getchar();
		g_kill = 1;
	}
	return 0;
}

typedef struct BEZIER_INFO_ {
	int x1;
	int y1;
	int x2;
	int y2;
	int cx;
	int cy;
	COLOR32 color;
} BEZIER_INFO;

typedef struct WORKSPACE_ {
	COLOR32 **images;
	BEZIER_INFO *beziers;
	COLOR32 *target;
	int width;
	int height;
	int nImages;
	int nImagesPerThread;
	HANDLE *hThreadCompletions;
} WORKSPACE;

typedef struct WORKER_STRUCT_ {
	WORKSPACE *workspace;
	HANDLE hEvent;
	int threadIndex;
	volatile int bestIndex;
	volatile double bestDistance;
} WORKER_STRUCT;

/*
workerthreadProc

Description:
	Function called to do work on the computations in another thread to run
	concurrently.

Parameters:
	lpParam	A pointer to the worker struct of parameters from the main thread
*/
DWORD CALLBACK workerThreadProc(LPVOID lpParam) {
	WORKER_STRUCT *worker = (WORKER_STRUCT *) lpParam;
	WORKSPACE *workspace = worker->workspace;
	int tid = worker->threadIndex;

	int startIndex = tid == 0 ? 1 : (tid * workspace->nImagesPerThread);
	int baseIndex = tid * workspace->nImagesPerThread;
	int nCanvas = workspace->nImagesPerThread;
	int width = workspace->width;
	int height = workspace->height;
	COLOR32 **images = workspace->images;
	COLOR32 *target = workspace->target;

	srand(tid + (int) time(0));
	while (1) {
		//wait for work
		WaitForSingleObject(worker->hEvent, INFINITE);

		//run calculations

		//draw tests
		for (int i = startIndex; i < baseIndex + nCanvas; i++) {
			COLOR32 *px = images[i];
			BEZIER_INFO *bInfo = workspace->beziers + i;
			volatile COLOR32 f = *px;
			drawRandomBezier(px, width, height, &bInfo->x1, &bInfo->y1, &bInfo->x2, &bInfo->y2, &bInfo->cx, &bInfo->cy, &bInfo->color);
		}

		//Find best match
		double bestDistance = 1e32;
		int bestImage = baseIndex;
		//if (tid == 0) __debugbreak();
		for (int i = baseIndex; i < baseIndex + nCanvas; i++) {
			double diff = imageDifference(target, images[i], width, height, bestDistance);
			if (diff < bestDistance) {
				bestDistance = diff;
				bestImage = i;
			}
		}

		worker->bestDistance = bestDistance;
		worker->bestIndex = bestImage;

		//signal completion
		SetEvent(workspace->hThreadCompletions[tid]);
	}
}

#define SVG_SWIZZLE(c) (((c)&0xFF00)|(((c)&0xFF)<<16)|(((c)>>16)&0xFF))

void writeBeziers(BEZIER_INFO *list, int nBeziers, COLOR32 bg, int width, int height, const char *path) {
	FILE *fp = fopen(path, "w");
	fprintf(fp, "<svg width=\"%d\" height=\"%d\" xmlns=\"http://www.w3.org/2000/svg\">", width, height);
	fprintf(fp, "<rect width=\"%d\" height=\"%d\" fill=\"#%06X\"/>", width, height, SVG_SWIZZLE(bg));
	fprintf(fp, "<g fill=\"none\" stroke-width=\"1\">");
	char cmdBufferAbsolute[64], cmdBufferRelative[64];
	for (int i = 0; i < nBeziers; i++) {
		COLOR32 col = list[i].color;
		col = (col & 0xFF00) | ((col & 0xFF) << 16) | ((col >> 16) & 0xFF);
		int x1 = list[i].x1, y1 = list[i].y1;
		int x2 = list[i].x2, y2 = list[i].y2;
		int cx = list[i].cx, cy = list[i].cy;

		int lenAbsolute = sprintf(cmdBufferAbsolute, "M%d %d Q%d %d %d %d", x1, y1, cx, cy, x2, y2);
		int lenRelative = sprintf(cmdBufferRelative, "M%d %d q%d %d %d %d", x1, y1, cx - x1, cy - y1, x2 - x1, y2 - y1);
		if (lenAbsolute <= lenRelative) {
			fprintf(fp, "<path d=\"%s\" stroke=\"#%06X\"/>",
				cmdBufferAbsolute, col);
		} else {
			fprintf(fp, "<path d=\"%s\" stroke=\"#%06X\"/>",
				cmdBufferRelative, col);
		}
	}
	fprintf(fp, "</g></svg>");
	fclose(fp);
}

#define MAXIMUM_ITERATIONS 0x4000

/*
main

Description:
	Main function
*/
int main(int argc, char **argv) {
	if (argc < 2) {
		puts("Usage: BezierTest <target bitmap> [iteration count]");
		return 1;
	}

	int width, height;
	COLOR32 *target = readBitmap24(argv[1], &width, &height);
	if (target == NULL) {
		puts("Error reading image file. Make sure it exists and is a valid 24-bit bitmap.");
		return 1;
	}
	COLOR32 avg = imageAverage(target, width, height);

	//get iteration count
	int nIterations = MAXIMUM_ITERATIONS;
	if (argc >= 3) {
		nIterations = atoi(argv[2]);
	}

	srand((int) time(0));
	int nCanvas = 128;
	COLOR32 **images = (COLOR32 **) calloc(nCanvas, sizeof(COLOR32 *));
	for (int i = 0; i < nCanvas; i++) {
		images[i] = allocImage(width, height);
		for (int j = 0; j < width * height; j++) {
			images[i][j] = avg;
		}
	}

	COLOR32 *best = images[0];
	int nMaxBezier = 4;
	int nIter = 0x1;
	DWORD tid;
	CreateThread(NULL, 0, killerCallback, NULL, 0, &tid);

	//create worker threads
	WORKSPACE workspace = { 0 };
	workspace.images = images;
	workspace.nImages = nCanvas;
	workspace.nImagesPerThread = 16;
	workspace.width = width;
	workspace.height = height;
	workspace.target = target;
	workspace.beziers = (BEZIER_INFO *) calloc(nCanvas, sizeof(BEZIER_INFO));

	//create buffer to hold result of beziers
	BEZIER_INFO *bezierList = (BEZIER_INFO *) calloc(nIterations, sizeof(BEZIER_INFO));
	int nBeziers = 0;

	int nThreads = nCanvas / workspace.nImagesPerThread;
	WORKER_STRUCT *workers = (WORKER_STRUCT *) calloc(nThreads, sizeof(WORKER_STRUCT));
	workspace.hThreadCompletions = (HANDLE *) calloc(nThreads, sizeof(HANDLE));
	
	for (int i = 0; i < nThreads; i++) {
		workers[i].hEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
		workers[i].threadIndex = i;
		workers[i].workspace = &workspace;
		workers[i].bestDistance = 1e32;
		workers[i].bestIndex = 0;
		CreateThread(NULL, 0, workerThreadProc, (LPVOID) &workers[i], 0, &tid);
		workspace.hThreadCompletions[i] = CreateEvent(NULL, TRUE, FALSE, NULL);
	}

	int k = 0;
	while(k < nIterations) {
		k++;

		//wake all threads
		for (int i = 0; i < nThreads; i++) {
			SetEvent(workers[i].hEvent);
		}

		WaitForMultipleObjects(nThreads, workspace.hThreadCompletions, TRUE, INFINITE);
		for (int i = 0; i < nThreads; i++) {
			ResetEvent(workspace.hThreadCompletions[i]);
		}

		//get best distance
		//__debugbreak();
		double bestDistance = 1e32;
		int bestImage = 0;
		for (int i = 0; i < nThreads; i++) {
			if (workers[i].bestDistance < bestDistance) {
				bestDistance = workers[i].bestDistance;
				bestImage = workers[i].bestIndex;
			}
		}
		if (k & 1)
			printf("%04X Distance: %f (%d)\n", k, bestDistance, bestImage);

		//copy best to all
		best = images[bestImage];
		for (int i = 0; i < nCanvas; i++) {
			if (i != bestImage) {
				memcpy(images[i], best, width * height * sizeof(COLOR32));
			}
		}
		
		//if best is nonzero, then that means a bezier was added
		if (bestImage) {
			BEZIER_INFO *toAdd = workspace.beziers + bestImage;
			BEZIER_INFO *dest = bezierList + nBeziers;
			memcpy(dest, toAdd, sizeof(BEZIER_INFO));
			nBeziers++;
		}

		if (g_kill) {
			writeBitmap(best, width, height, "out.bmp");
			writeBeziers(bezierList, nBeziers, avg, width, height, "out.svg");
			g_kill = 0;
		}
	}

	writeBitmap(best, width, height, "out.bmp");
	writeBeziers(bezierList, nBeziers, avg, width, height, "out.svg");
	return 0;
}