#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdbool.h>
#include <complex.h>
#include <pthread.h>

#include <gtk/gtk.h>
#include <cairo.h>

#define DEFAULT_THREADS 4
#define MAX_THREADS 32
#define STR(x) #x



/******* Program data structures *******/



int32_t thread_count;

struct threadInfo {
	int32_t index; // unique number from 0 to thread_count-1
	bool drawing; // marks if we are currently drawing for a given thread
	bool complete; // marks if the drawing that it was supposed to
};

/*
 * Concurrency model explained:
 * One reader thread (the GUI thread) of the pixel buffer, along with many
 * writers, who split the work of rendering roughly equally. The writers don't
 * do any work until the pixel map which they are meant to work on is marked as
 * available.
 *
 * Once it is marked as such, this means that the planeView structure is well
 * defined, that pixmap points to a memory region, which is allocated with
 * enough memory for the plain view, and they can all start writing without an
 * issue.
 */

pthread_mutex_t pixmapMutex;
pthread_cond_t pixmapCond;

struct threadInfo threads[MAX_THREADS];
// this will store an array with status info for all the threads

cairo_surface_t *surface = NULL;

bool pixmapAvailable = false;
unsigned char *pixmap;
struct planeView {
	double scale; // scale is represented as unit/pixel
	int width, height;
	double centerX, centerY;
};

// struct planeView mandelbrot = { 0.00000000001, 0, 0, 0.001643721971153, 0.822467633298876};
struct planeView mandelbrot = { 1, 0, 0, 0, 0};



/******* Worker thread code *******/



void draw_mandelbrot(struct threadInfo *info);

void color_from_iteration(int *r, int *g, int *b, double x0, double y0);
void color_lookup(int *r, int *g, int *b, double mu);

void *writer_thread(void *arg)
{
	struct threadInfo *info = arg;
	while (true) {
		pthread_mutex_lock(&pixmapMutex);
		while (!pixmapAvailable || info->complete) {
			pthread_cond_wait(&pixmapCond, &pixmapMutex);
		}
		pthread_mutex_unlock(&pixmapMutex);
		info->drawing = true;
		draw_mandelbrot(info);
		info->drawing = false;
	}
}

void draw_mandelbrot(struct threadInfo *info)
{
	int h = mandelbrot.height;
	int w = mandelbrot.width;
	int p = info->index, c = thread_count;
	int mod = h % c, div = h / c;
	int a = MIN(p, mod)*(div+1) + (MAX(p, mod)-mod)*div,
	    b = p < mod ? a + div + 1 : a + div;
	// [a..b) is the range of rows a given thread is meant to write pixels

	double centerX = mandelbrot.centerX;
	double centerY = mandelbrot.centerY;
	double scale = mandelbrot.scale;

	for (int i = a; i < b; i++) {
		double yRange = (double) i / h * 2.0 - 1.0;
		double y = centerY+yRange*scale;
		pthread_mutex_lock(&pixmapMutex);
		if (!pixmapAvailable) {
			pthread_mutex_unlock(&pixmapMutex);
			return;
			// after this, the thread will again enter the
			// loop in the writer_thread function, and
			// enter a wait state
		}
		pthread_mutex_unlock(&pixmapMutex);
		for (int j = 0; j < w; j++) {
			double xRange = (double) j / w * 2.0 - 1.0;
			double x = centerX+xRange*scale;

			int r, g, b;
			color_from_iteration(&r, &g, &b, x, y);
			pixmap[4*(i*w + j) + 2] = r;
			pixmap[4*(i*w + j) + 1] = g;
			pixmap[4*(i*w + j) + 0] = b;
		}
	}
	pthread_mutex_lock(&pixmapMutex);
	info->complete = true;
	bool allWritersComplete = true;
	for (int32_t i = 0; i < thread_count; i++) {
		if (!threads[i].complete) {
			allWritersComplete = false;
		}
	}
	if (allWritersComplete) {
		cairo_surface_mark_dirty(surface);
	}
	pthread_mutex_unlock(&pixmapMutex);
}

#define MAX_ITERATION 1000
#define ESC_RAD 20.0
void color_from_iteration(int *r, int *g, int *b, double x0, double y0)
{
	double x = 0;
	double y = 0;
	int iteration = 0;
	while (x*x + y*y <= ESC_RAD*ESC_RAD && iteration < MAX_ITERATION) {
		double xtmp = xtmp = x*x - y*y + x0;
		y = 2*x*y + y0;
		x = xtmp;
		iteration++;
	}
	// at this point iteration is the number of iterations required for the value to
	// escape. X and Y contain the first escaped value

	double mu = (iteration + 1 - log(log(sqrt(x*x+y*y)))/log(2));
	if (mu < 0)
		mu = 0;
	if (mu > MAX_ITERATION)
		mu = MAX_ITERATION;

	color_lookup(r, g, b, mu);

}

void color_lookup(int *r, int *g, int *b, double mu)
{
	static const int table16[16][3] = {
		{ 66, 30, 15 },
		{ 25, 7, 26 },
		{ 9, 1, 47 },
		{ 4, 4, 73 },
		{ 0, 7, 100 },
		{ 12, 44, 138 },
		{ 24, 82, 177 },
		{ 57, 125, 209 },
		{ 134, 181, 229 },
		{ 211, 236, 248 },
		{ 241, 233, 191 },
		{ 248, 201, 95 },
		{ 255, 170, 0 },
		{ 204, 128, 0 },
		{ 153, 87, 0 },
		{ 106, 52, 3 }
	};
	// *r = 0; // (1000.0-i)/1000.0*256.0;
	// *g = 0; // (1000.0-i)/1000.0*256.0;
	// *b = ((double) MAX_ITERATION - i - 1) / MAX_ITERATION * 255.0 - log(log(2)) / log(2); // (i/MAX_ITERATION)
	int mod = (int) mu % 16;
	*r = table16[mod][0];
	*g = table16[mod][1];
	*b = table16[mod][2];
	if (mu == 0 || mu >= MAX_ITERATION - 50) {
		*r = *g = *b = 0;
	}
}
#undef MAX_ITERATION



/******* Main thread code *******/



void print_usage(FILE *stream);
static void app_activate(GApplication *app);

void blit_plane(GtkDrawingArea *da, cairo_t *cr, int width, int height, gpointer data);
void plane_resize(GtkWidget *widget);
void create_surface(GtkWidget *widget);

int main(int argc, char **argv)
{
	pthread_mutex_init(&pixmapMutex, NULL);
	pthread_cond_init(&pixmapCond, NULL);

	if (argc == 1) {
		thread_count = 4;
	} else if (argc == 2 && !strcmp(argv[1], "-h")) {
		print_usage(stdout);
		exit(0);
	} else if (argc == 2) {
		char *endptr;
		thread_count = strtol(argv[1], &endptr, 10);
		if (argv[1][0] == '\0' || *endptr != '\0' ||
				thread_count < 1 ||
				thread_count > MAX_THREADS) {
				// if argv[1] is not a valid decimal number, or
				// is not in the range
			fprintf(stderr, "Invalid number of threads, "
					"terminating application!\n");
			print_usage(stderr);
			exit(1);
		}
	} else {
		fprintf(stderr, "Wrong arguments, terminating application!\n");
		print_usage(stderr);
		exit(1);
	}

	int i, count;
	for (i = 0, count = 0; i < thread_count; i++) {
		threads[count].index = count;
		// we set this before the thread starts so it points to a valid
		// integer, if the thread fails to start, it will simply be
		// overwritten by the same value in the next run of the loop
		pthread_t id;
		if (pthread_create(&id, NULL, &writer_thread, &threads[count])) {
			pthread_detach(id);
			fprintf(stderr, "Failed to create thread %d\n", i);
			continue;
		}
		count++;
	}
	thread_count = count;
	// Update thread count to show the number of threads actually created,
	// not the number requested

	if (thread_count == 0) {
		fprintf(stderr, "Failed to start any threads, terminating "
				"application\n");
		exit(1);
	}

	// Main thread actually running GTK UI from here on out

	int stat = 0;

	GtkApplication *app;
	app = gtk_application_new("cool.bonsai.mandelbrot-visualizer",
			G_APPLICATION_DEFAULT_FLAGS);
	g_signal_connect(app, "activate", G_CALLBACK(app_activate), NULL);
	stat = g_application_run(G_APPLICATION(app), argc, argv);
	g_object_unref(app);

	return stat;
}

void print_usage(FILE *stream)
{
	char *err_msg = "Mandelbrot visualiser (visor):\n"
		"Usage:\n"
		"visor -h : Show this help\n"
		"visor [threads] : Run the GTK visualizer with a given number\n"
		"                  of threads (default " STR(DEFAULT_THREADS) ")\n"
		"                  [1.." STR(MAX_THREADS) "]\n";
	fprintf(stream, "%s", err_msg);
}

static void app_activate(GApplication *app)
{
	GtkWidget *win;
	GtkWidget *btn1;
	GtkWidget *btn2;
	GtkWidget *bigBox;
	GtkWidget *vertBox;
	GtkDrawingArea *da;

	win = gtk_application_window_new(GTK_APPLICATION(app));
	gtk_window_set_title(GTK_WINDOW(win), "Mandelbrot visualiser");
	gtk_window_set_default_size(GTK_WINDOW(win), 800, 600);

	btn1 = gtk_button_new_with_label("A");
	gtk_widget_set_size_request(GTK_WIDGET(btn1), 20, 20);
	gtk_widget_set_margin_bottom(GTK_WIDGET(btn1), 5);
	gtk_widget_set_margin_top(GTK_WIDGET(btn1), 5);
	gtk_widget_set_margin_start(GTK_WIDGET(btn1), 5);
	gtk_widget_set_margin_end(GTK_WIDGET(btn1), 5);
	btn2 = gtk_button_new_with_label("B");
	gtk_widget_set_size_request(GTK_WIDGET(btn2), 20, 20);
	gtk_widget_set_margin_bottom(GTK_WIDGET(btn2), 5);
	gtk_widget_set_margin_top(GTK_WIDGET(btn2), 5);
	gtk_widget_set_margin_start(GTK_WIDGET(btn2), 5);
	gtk_widget_set_margin_end(GTK_WIDGET(btn2), 5);

	vertBox = gtk_box_new(GTK_ORIENTATION_VERTICAL, 5);
	gtk_widget_set_valign(GTK_WIDGET(vertBox), GTK_ALIGN_START);
	gtk_box_append(GTK_BOX(vertBox), btn1);
	gtk_box_append(GTK_BOX(vertBox), btn2);

	bigBox = gtk_box_new(GTK_ORIENTATION_HORIZONTAL, 0);

	da = g_object_new(GTK_TYPE_DRAWING_AREA,
			"accessible-role", GTK_ACCESSIBLE_ROLE_IMG,
			NULL);
	gtk_widget_set_hexpand(GTK_WIDGET(da), true);
	gtk_drawing_area_set_content_width(GTK_DRAWING_AREA(da), 300);
	gtk_drawing_area_set_content_height(GTK_DRAWING_AREA(da), 200);
	gtk_drawing_area_set_draw_func(GTK_DRAWING_AREA(da), blit_plane, NULL, NULL);
	g_signal_connect(da, "resize", G_CALLBACK(plane_resize), NULL);

	// TODO: implement later drag = gtk_gesture_drag_new();

	gtk_box_append(GTK_BOX(bigBox), GTK_WIDGET(da));
	gtk_box_append(GTK_BOX(bigBox), vertBox);
	gtk_window_set_child(GTK_WINDOW(win), bigBox);

	gtk_window_present(GTK_WINDOW(win));
}

void blit_plane(GtkDrawingArea *da, cairo_t *cr, int width, int height, gpointer data)
{
	(void) da;
	(void) width;
	(void) height;
	(void) data;
	cairo_set_source_surface(cr, surface, 0, 0);
	cairo_paint(cr);
}

void plane_resize(GtkWidget *widget)
{
	if (!surface) {
		create_surface(widget);
	}
}

void create_surface(GtkWidget *widget)
{
	// Stop all current drawers
	pthread_mutex_lock(&pixmapMutex);
	pixmapAvailable = false; // Mark drawing area as unavailable
	pthread_mutex_unlock(&pixmapMutex);

	bool allWritersStopped;
	do { // Wait until all writing threads have been stopped
		allWritersStopped = true;
		for (int32_t i = 0; i < thread_count; i++) {
			if (threads[i].drawing) {
				allWritersStopped = false;
				break;
			}
		}
	} while (!allWritersStopped);

	// If all drawers hadn't completed, we need to call
	// cairo_surface_mark_dirty in the main thread
	bool allWritersHadCompleted = true;
	for (int32_t i = 0; i < thread_count; i++) {
		if (!threads[i].complete) {
			allWritersHadCompleted = false;
			break;
		}
	}
	if (!allWritersHadCompleted && surface != NULL) {
		cairo_surface_mark_dirty(surface);
	}

	for (int32_t i = 0; i < thread_count; i++) {
		threads[i].complete = false;
	}

	if (surface) {
		cairo_surface_destroy(surface);
	}

	surface = cairo_image_surface_create(CAIRO_FORMAT_RGB24,
			gtk_widget_get_width(widget),
			gtk_widget_get_height(widget));

	int h = cairo_image_surface_get_height(surface);
	int w = cairo_image_surface_get_width(surface);
	mandelbrot.height = h;
	mandelbrot.width = w;

	// Mark the surface as ready to be drawn by memory access
	// Then notify other threads to start drawing
	// One of these threads will make the corresponding call using
	// cairo_surface_mark_dirty, to notify they are all done drawing
	cairo_surface_flush(surface);
	pixmap = cairo_image_surface_get_data(surface);

	// This lock is unnecessary since all other threads are stopped already
	pthread_mutex_lock(&pixmapMutex);
	pixmapAvailable = true;
	pthread_mutex_unlock(&pixmapMutex);
	pthread_cond_broadcast(&pixmapCond);
}