#include <gtk/gtk.h>
#include <mpi.h>
#include <complex.h>

int procRank, commSz;
MPI_Datatype planeView_t;

MPI_Win window;

unsigned char *pixmap;

cairo_surface_t *surface = NULL;

#define SCR_DIMENSION 1000

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};

#define MAX_ITERATION 1000
#define ESC_RAD 20.0
int 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;
	}
}

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);

#undef MAX_ITERATION
}

void draw_mandelbrot(unsigned char* data, int height, int width, double planeCenterX, double planeCenterY, double scale)
{
	for (int i = 0; i < height; i++) {
		double yRange = (double) i / height * 2.0 - 1.0;
		double y = planeCenterY+yRange*scale;
		for (int j = 0; j < width; j++) {
			double xRange = (double) j / width * 2.0 - 1.0;
			double x = planeCenterX+xRange*scale;
			int r, g, b;
			color_from_iteration(&r, &g, &b, x, y);
			data[4*(i*width + j) + 2] = r;
			data[4*(i*width + j) + 1] = g;
			data[4*(i*width + j) + 0] = b;
		}
	}
}

void create_surface(GtkWidget *widget)
{
	cairo_t *cr;

	if (surface)
		cairo_surface_destroy(surface);

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

	cr = cairo_create(surface);

	cairo_surface_flush(surface);
	int h = cairo_image_surface_get_height(surface);
	int w = cairo_image_surface_get_width(surface);
	mandelbrot.height = h;
	mandelbrot.width = w;
	for (int p = 1; p < commSz; p++) {
		MPI_Send(&mandelbrot, 1, planeView_t, p, 0, MPI_COMM_WORLD);
	}
	pixmap = malloc(sizeof(unsigned char) * 4 * w * h);
	MPI_Win_create(pixmap, 4*w*h, 1,
			MPI_INFO_NULL, MPI_COMM_WORLD, &window);
	MPI_Win_fence(0, window);

	unsigned char *data = cairo_image_surface_get_data(surface);
	// draw_mandelbrot(data, h, w, mandelbrot.centerX, mandelbrot.centerY, mandelbrot.scale);

	
	MPI_Win_fence(0, window);
	memcpy(data, pixmap, 4*h*w);
	cairo_surface_mark_dirty(surface);
	// cairo_paint(cr);

	cairo_destroy(cr);
}

void plane_resize(GtkWidget *widget, int width, int height)
{
	if (!surface)
		create_surface(widget);
}

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

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), draw_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));
}



int main(int argc, char **argv)
{
	MPI_Init(NULL, NULL);

	MPI_Comm_rank(MPI_COMM_WORLD, &procRank);
	MPI_Comm_size(MPI_COMM_WORLD, &commSz);

	int blocklengths[5] = {1,1,1,1,1};
	const MPI_Aint displs[] = {0, sizeof(double), sizeof(double)+sizeof(int),
	sizeof(double)+2*sizeof(int), 2*sizeof(double)+2*sizeof(int)};
	MPI_Datatype types[5] = {MPI_DOUBLE, MPI_INT, MPI_INT, MPI_DOUBLE, MPI_DOUBLE};
	MPI_Type_create_struct(5, blocklengths, displs, types, &planeView_t);
	MPI_Type_commit(&planeView_t);

	if (commSz == 1) {
		printf("This application needs to be run with more than one"
				" process. Try using mpiexec.\n");
		MPI_Finalize();
		return -1;
	}

	if (procRank == 0) {
		int stat = 0;

		GtkApplication *app;
		app = gtk_application_new("com.github.ToshioCP.pr1",
				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);

		MPI_Abort(MPI_COMM_WORLD, 0);
		MPI_Finalize();
		return stat;
	}

	MPI_Recv(&mandelbrot, 1, planeView_t, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
	pixmap = malloc(sizeof(unsigned char) * 4 * mandelbrot.width * mandelbrot.height);
	MPI_Win_create(pixmap, 3*mandelbrot.width*mandelbrot.height, 1,
			MPI_INFO_NULL, MPI_COMM_WORLD, &window);
	MPI_Win_fence(0, window);

	int h = mandelbrot.height;
	int w = mandelbrot.width;
	int p = procRank - 1, c = commSz - 1;
	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;


	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;
		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;
		}
	}

	MPI_Put(pixmap + 4*a*w, 4*(b-a)*w, MPI_UNSIGNED_CHAR, 0, 4*a*w, 4*(b-a)*w, MPI_UNSIGNED_CHAR, window);
	MPI_Win_fence(0, window);


	MPI_Finalize();
}