#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <string.h>
#include <stdbool.h>

#include <CImg.h>

#include <omp.h>

using namespace std;
using namespace cimg_library;

#define DELTA_X 1.0
#define DELTA_Y 1.0

#define WIDTH 680
#define HEIGHT 377

#define ITERATIONS 2000
#define EPSILON 0.00

#define OMEGA 1.0

//#define debug printf
#define debug //printf


float U = -17.0; // Primary energy
float Va = 100.0; // Accel. 
float Vf = -50.0;//8.18; // Focus
float Vd = 0.0;//0.605; // Deflection 
float Vw = 10.0; // Wenhault

float grid[2][WIDTH][HEIGHT];
bool fixed[WIDTH][HEIGHT];

void SimulateElectron(float x, float y, float theta, float vMax);


int main(int argc, char ** argv)
{
	float beta2 = pow(DELTA_X / DELTA_Y, 2);
	//printf("Beta2 is %f\n", beta2);
	omp_set_nested(1);
	
	for (int i = 0; i < WIDTH; ++i)
	{
		for (int j = 0; j < HEIGHT; ++j)
		{
			grid[0][i][j] = 0;
			fixed[i][j] = (i <= 1 || j <= 1 || i >= WIDTH-2 || j >= HEIGHT-2);
		}
	}

	debug("1. fixed[0][0] is %d\n", (int)fixed[0][0]);


	// Add the electrodes
	CImg<unsigned char> src(argv[1]);
	int w = src.width();
	int h = src.height();
	for (int i = 0; i < w; ++i)
	{
		for (int j = 0; j < h; ++j)
		{

			int ii = i * (WIDTH / w);
			int jj = j * (HEIGHT / h);
			//printf("At %d %d vs %d %d\n", i, j, ii, jj);
			
			int rgb[3];
			for (int d = 0; d < 3; ++d)
			{
				//printf("rgb[%d] is %d\n", d, rgb[d]);
				rgb[d] = (int)(src(i, j, 0, d));
			}

			
			
			// Determine the electrode based on colour
			if (rgb[0] == 255 && rgb[1] == 0 && rgb[2] == 0)
				grid[0][ii][jj] = Va; //red
			else if (rgb[0] == 0 && rgb[1] == 0 && rgb[2] == 255)
				grid[0][ii][jj] = Vf; // blue
			else if (rgb[0] == 0 && rgb[1] == 255 && rgb[2] == 0)
				grid[0][ii][jj] = Vw; //green 
			else if (rgb[0] == 255 && rgb[1] == 255 && rgb[2] == 0)
				grid[0][ii][jj] = Va - Vd; //yellow
			else if (rgb[0] == 0 && rgb[1] == 255 && rgb[2] == 255)
				grid[0][ii][jj] = Va + Vd; //light blue
			else if (rgb[0] == 255 && rgb[1] == 0 && rgb[2] == 255)
				grid[0][ii][jj] = U; // purple
			else
				grid[0][ii][jj] = 0.0;

			//if (fixed[ii][jj])
			//	printf("Fixed at %d %d\n", ii, jj);

			if (fixed[ii][jj] == false)
			{
				if (rgb[0] == 255 && rgb[1] == 255 && rgb[2] == 255)
					fixed[ii][jj] = false;
				else
					fixed[ii][jj] = true;
			}
			
		}
	}
	debug("2. fixed[0][0] is %d\n", (int)fixed[0][0]);

	for (int i = 0; i < WIDTH; ++i)
	{
		for (int j = 0; j < HEIGHT; ++j)
		{
			grid[1][i][j] = grid[0][i][j];
		}
	}

	

	int n = 1; int p = 0;
	int k = 0;
	float maxChange = 0.0;
	debug("3. fixed[0][0] is %d\n", (int)fixed[0][0]);

	for (k = 0; k < ITERATIONS; ++k)
	{
		maxChange = 0.0;
		
		//#pragma omp parallel for	
		for (int i = 0; i < WIDTH; ++i)
		{
			//#pragma omp parallel for
			for (int j = 0; j < HEIGHT; ++j)
			{
				//printf("4. fixed[0][0] is %d\n", (int)fixed[0][0]);

				if (fixed[i][j] == true)
				{

					grid[n][i][j] = grid[p][i][j];

					continue;
				}
				else
				{
					if (i <= 1 || j <= 1 || i >= WIDTH-2 || j >= HEIGHT-2)
					{
						debug("WTF at %d %d\n", i, j);
						fixed[i][j] = true;
						continue;
					}
					
				}
				
				grid[n][i][j] = (1.0 - OMEGA) * grid[p][i][j];
				if (isnan(grid[n][i][j]))
				{	
					debug("1. NaN at %d %d (n = %d), fixed is %d\n", i, j, n, (int)fixed[i][j]);

					exit(EXIT_FAILURE);
				}
				//printf("4. fixed[0][0] is %d\n", (int)fixed[0][0]);
				grid[n][i][j] += OMEGA * (grid[n][i-1][j] + grid[p][i+1][j] + beta2 * (grid[n][i][j-1] + grid[p][i][j+1]));
				if (isnan(grid[n][i][j]))
				{	
					debug("5. fixed[0][0] is %d\n", (int)fixed[0][0]);
					debug("2. NaN at %d %d (n = %d), fixed is %d\n", i, j, n, (int)fixed[i][j]);
					debug("6. fixed[0][0] is %d\n", (int)fixed[0][0]);
					exit(EXIT_FAILURE);
				}

				grid[n][i][j] = grid[n][i][j] / (2.0 * (1.0 + beta2));
				if (isnan(grid[n][i][j]))
				{	
					debug("3. NaN at %d %d (n = %d), fixed is %d\n", i, j, n, (int)fixed[i][j]);
					exit(EXIT_FAILURE);
				}

				float change = (grid[n][i][j] > grid[p][i][j]) ? grid[p][i][j] - grid[n][i][j] : grid[p][i][j] - grid[n][i][j];
				if (grid[p][i][j] != 0.00)
					change = change / grid[p][i][j];
				else
					change = 1.0;
				if (change >= maxChange)
					maxChange = change;
				

			}
		}
		
		//fprintf(stdout, "# Complete iteration %d; max change was %f\n", k, maxChange);

		
		
		n = (n == 1) ? 0 : 1;
		p = (p == 1) ? 0 : 1;

		if (maxChange < EPSILON)
			break;

	}
		char file[100];
		sprintf(file, "%s.dat", "main");
		FILE * out = fopen(file, "w");

		fprintf(out, "# Grid dimensions %d %d\n", WIDTH, HEIGHT);
		fprintf(out, "# dx = %f, dy = %f, beta2 = %f\n", DELTA_X, DELTA_Y, beta2);
		fprintf(out, "# Ran for %d iterations\n", k);
		fprintf(out, "# Largest change is %f\n", maxChange);
		for (int i = 0; i < WIDTH; ++i)
		{
			for (int j = 0; j < HEIGHT; ++j)
			{
				fprintf(out, "%f\t%f\t%f\t", i*DELTA_X, j*DELTA_Y, grid[n][i][j]);
				if (fixed[i][j])
					fprintf(out,"1\n");
				else
					fprintf(out,"0\n");
			}
			fprintf(out, "\n");
		}
		fprintf(out, "# DONE\n");
		if (out != stdout && out != stderr) fclose(out);

	
		//return 0;

		//#pragma omp parallel for
		int ox = 90;
		int oy = 190;
		for (float v = 1.0; v <= 1.0; v += 0.01)
		{


			SimulateElectron(ox, oy, -M_PI/4.0, v);
			SimulateElectron(ox, oy, +M_PI/4.0, v);
			SimulateElectron(ox, oy, 0.0, v);
			SimulateElectron(ox, oy, -M_PI/2.0, v);
			SimulateElectron(ox, oy, +M_PI/2.0, v);
			SimulateElectron(ox, oy, -M_PI/8.0, v);
			SimulateElectron(ox, oy, +M_PI/8.0, v);
			SimulateElectron(ox, oy, -M_PI/16.0, v);
			SimulateElectron(ox, oy, +M_PI/16.0, v);
			SimulateElectron(ox, oy, -M_PI/32.0, v);
			SimulateElectron(ox, oy, +M_PI/32.0, v);
			SimulateElectron(ox, oy, -M_PI/64.0, v);
			SimulateElectron(ox, oy, +M_PI/64.0, v);
			SimulateElectron(ox, oy, -M_PI/128.0,v);
			SimulateElectron(ox, oy, +M_PI/128.0, v);

		}	

		

}

void SimulateElectron(float x, float y, float theta, float vMax)
{
	char buffer[256];
	sprintf(buffer, "electron%f_%f.dat", theta, vMax);
	FILE * out = fopen(buffer, "w");
	
	float dt = 0.01;
	float p[2];
	float prev[2];
	
	float v[2];
	float a[2];
	p[0] = x; p[1] = y; prev[0] = x; prev[1] = y;
	v[0] = vMax * cos(theta);
	v[1] = vMax * sin(theta);		
	for (float t = 0.0; t < 2000.0; t += dt)
	{
		
		if ((int)p[0] <= 0 || (int)p[0] >= WIDTH-1 || (int)p[1] <= 0 || (int)p[1] >= HEIGHT-1)
			break;
		if (fixed[(int)p[0]][(int)p[1]] == true && grid[0][(int)p[0]][(int)p[1]] != U)
			break;

		fprintf(out, "%f\t%f\t%f\t%f\t%f\n", t, p[0], p[1], v[0], v[1]);
	
		float ddx = 0.5 * (grid[0][(int)(p[0])+1][(int)(p[1])] - grid[0][(int)(p[0])-1][(int)p[1]]);
		float ddy = 0.5 * (grid[0][(int)(p[0])][(int)(p[1])+1] - grid[0][(int)(p[0])][(int)(p[1])-1]);
		v[0] += ddx;
		v[1] += ddy;
		prev[0] = p[0]; prev[1] = p[1];
		p[0] += v[0] * dt; p[1] += v[1] * dt;
	}

	if (out != stdout && out != stderr) fclose(out);

}

//EOF