#include <stdlib.h>
#include <string>
#include <fstream>
#include <algorithm>

// CUDA include
#ifdef __CUDACC__
#include "device_launch_parameters.h"
#include <cuda.h>
#include <cuda_runtime_api.h>
#include "cuda_runtime.h"
#endif

// OPENGL include
#include <GL/glut.h>
#include <GL/freeglut.h>

#include "flow.h"

// STIM include
#include <stim/visualization/gl_aaboundingbox.h>
#include <stim/parser/arguments.h>
#include <stim/visualization/camera.h>
#include <stim/visualization/colormap.h>
#include <stim/cuda/cudatools/error.h>


//********************parameter setting********************
// overall parameters
int vX, vY;
float dx, dy, dz;										// x, y and z image scaling(units/pixel)
std::string stackdir = "";								// directory where image stacks will be stored
stim::arglist args;										// create an instance of arglist
stim::gl_aaboundingbox<float> bb;						// axis-aligned bounding box object
stim::camera cam;										// camera object
unsigned num_edge;										// number of edges in the network
unsigned num_vertex;									// number of vertex in the network
std::vector<unsigned> pendant_vertex;					// list of pendant vertex index in GT
std::vector<std::string> menu_option = { "simulation", "build inlet/outlet", "manufacture" };
stim::flow<float> flow;									// flow object
float move_pace;										// camera moving parameter
float u;												// viscosity
float rou;												// density
float max_v;
float min_v;
int mods;												// special keyboard input
std::vector<unsigned char> color;						// velocity color map
std::vector<int> velocity_bar;							// velocity bar

// hard-coded parameters
float camera_factor = 1.2f;			// start point of the camera as a function of X and Y size
float orbit_factor = 0.01f;			// degrees per pixel used to orbit the camera
float zoom_factor = 10.0f;			// zooming factor
float border_factor = 20.0f;		// border
float radii_factor = 1.0f;			// radii changing factor
GLint subdivision = 20;				// slices and stacks
float default_radii = 5.0f;			// default radii of network vertex
float delta = 0.01f;				// small discrepancy
float eps = 20.0f;					// epsilon threshold
float max_pressure = 0.0f;			// maximum pressure that the channel can bear

// glut event parameters
int mouse_x;						// window x-coordinate
int mouse_y;						// window y-coordinate
bool LTbutton = false;				// true means down while false means up				

// simulation parameters
bool simulation = false;			// flag indicates simulation mode
bool color_bound = false;			// flag indicates velocity color map bound
bool to_select_pressure = false;	// flag indicates having selected a vertex to modify pressure
unsigned pressure_index;			// the index of vertex that is clicked

// build inlet/outlet parameters
bool build_inlet_outlet = false;	// flag indicates building inlets and outlets
bool modified_bridge = false;		// flag indicates having modified inlet/outlet connection

// manufacture parameters
bool manufacture = false;			// flag indicates manufacture mode


//********************helper function*********************
// get the network basic information
inline void get_background() {

	pendant_vertex = flow.get_boundary_vertex();
	num_edge = flow.edges();
	num_vertex = flow.vertices();

	// set the initial radii
	flow.init(num_edge, num_vertex);			// initialize flow object
	for (unsigned i = 0; i < num_edge; i++)
		flow.set_r(i, default_radii);
}

// convert from window coordinates to world coordinates
inline void window_to_world(GLdouble &x, GLdouble &y, GLdouble &z) {

	GLint    viewport[4];
	GLdouble modelview[16];
	GLdouble projection[16];
	GLdouble winX, winY;
	GLfloat  winZ;

	glGetIntegerv(GL_VIEWPORT, viewport);
	glGetDoublev(GL_MODELVIEW_MATRIX, modelview);
	glGetDoublev(GL_PROJECTION_MATRIX, projection);

	winX = (GLdouble)mouse_x;
	winY = viewport[3] - (GLdouble)mouse_y;
	glReadPixels((GLint)winX, (GLint)winY, (GLsizei)1, (GLsizei)1, GL_DEPTH_COMPONENT, GL_FLOAT, &winZ);
	gluUnProject(winX, winY, winZ, modelview, projection, viewport, &x, &y, &z);
}


//********************simulation function**********************
// initialize flow object
void flow_initialize() {

	stim::vec3<float> center = bb.center();

	for (unsigned i = 0; i < pendant_vertex.size(); i++) {
		if (flow.get_vertex(pendant_vertex[i])[0] <= center[0])
			flow.P[pendant_vertex[i]] = max_pressure - i * delta;	// should set minor discrepancy
		else
			flow.P[pendant_vertex[i]] = (i + 1) * delta;			// algorithm treat 0 as no initial pressure
	}
}

// find the stable flow state
void flow_stable_state() {
	
	flow.solve_flow(u);
	flow.get_color_map(max_v, min_v, color, pendant_vertex);
	color_bound = true;

	velocity_bar.resize(num_edge);
	for (unsigned i = 0; i < num_edge; i++)
		velocity_bar[i] = i;
	std::sort(velocity_bar.begin(), velocity_bar.end(), [&](int x, int y) {return abs(flow.v[x]) < abs(flow.v[y]); });
}


//********************glut function********************
// dynamically set menu
// @param num: number of current menu options
// @param range: range of option to be set from menu_option list
void glut_set_menu(int num, int range) {

	// remove last time menu options
	for (int i = 1; i < num + 1; i++)
		glutRemoveMenuItem(1);

	// set new menu options
	std::string menu_name;
	for (int i = 1; i < range + 1; i++) {
		menu_name = menu_option[i - 1];
		glutAddMenuEntry(menu_name.c_str(), i);
	}
}

// set up the squash transform to whole screen
void glut_projection() {

	glMatrixMode(GL_PROJECTION);					// load the projection matrix for editing
	glLoadIdentity();								// start with the identity matrix
	vX = glutGet(GLUT_WINDOW_WIDTH);				// use the whole screen for rendering
	vY = glutGet(GLUT_WINDOW_HEIGHT);
	glViewport(0, 0, vX, vY);						// specify a viewport for the entire window
	float aspect = (float)vX / (float)vY;			// calculate the aspect ratio
	gluPerspective(60, aspect, 0.1, 1000000);		// set up a perspective projection
}

// translate camera to origin
void glut_modelview() {

	glMatrixMode(GL_MODELVIEW);						// load the modelview matrix for editing
	glLoadIdentity();								// start with the identity matrix
	stim::vec3<float> eye = cam.getPosition();		// get the camera position (eye point)
	stim::vec3<float> focus = cam.getLookAt();		// get the camera focal point
	stim::vec3<float> up = cam.getUp();				// get the camera "up" orientation

	gluLookAt(eye[0], eye[1], eye[2], focus[0], focus[1], focus[2], up[0], up[1], up[2]);	// set up the OpenGL camera
}

// glut render function
void glut_render() {

	glEnable(GL_DEPTH_TEST);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
	glut_projection();
	glut_modelview();

	if (!simulation && !build_inlet_outlet || manufacture) {
		flow.glCylinder0();
	}
	else {
		flow.bounding_box();
		flow.glSolidSphere(max_pressure, subdivision);
		flow.mark_vertex();
		flow.glSolidCone(subdivision);
		flow.glSolidCylinder(subdivision, color);
		flow.glSolidCuboid();
	}

	if (build_inlet_outlet) {
		flow.line_bridge();
	}
	
	if (manufacture) {
		flow.glSolidCuboid();
		flow.tube_bridge(subdivision);
	}
	

	// render bars
	// bring up a pressure bar on left
	if (to_select_pressure) {
		
		glMatrixMode(GL_PROJECTION);									// set up the 2d viewport for mode text printing
		glPushMatrix();
		glLoadIdentity();
		vX = glutGet(GLUT_WINDOW_WIDTH);								// get the current window width
		vY = glutGet(GLUT_WINDOW_HEIGHT);								// get the current window height
		glViewport(0, 0, vX, vY);										// locate to left bottom corner
		gluOrtho2D(0, vX, 0, vY);										// define othogonal aspect

		glMatrixMode(GL_MODELVIEW);
		glPushMatrix();
		glLoadIdentity();

		glLineWidth(border_factor);
		glBegin(GL_LINES);
		glColor3f(0.0, 0.0, 1.0);										// blue to red
		glVertex2f(border_factor, border_factor);
		glColor3f(1.0, 0.0, 0.0);
		glVertex2f(border_factor, (vY - 2.0f * border_factor));
		glEnd();
		glFlush();

		// pressure bar text
		glColor3f(1.0f, 1.0f, 1.0f);
		glRasterPos2f(0.0f, vY - border_factor);
		std::stringstream ss_p;
		ss_p << "Pressure Bar";
		glutBitmapString(GLUT_BITMAP_HELVETICA_18, (const unsigned char*)(ss_p.str().c_str()));

		// pressure range text
		float step = vY - 3.0f * border_factor;
		step /= 10;
		for (unsigned i = 0; i < 11; i++) {
			glRasterPos2f((border_factor * 1.5f), (border_factor + i * step));
			std::stringstream ss_n;
			ss_n << (float)i * max_pressure / 10;
			glutBitmapString(GLUT_BITMAP_TIMES_ROMAN_10, (const unsigned char*)(ss_n.str().c_str()));
		}
		glPopMatrix();
		glMatrixMode(GL_PROJECTION);
		glPopMatrix();
	}

	// bring up a velocity bar on left
	if (simulation && !to_select_pressure) {
		
		glMatrixMode(GL_PROJECTION);									// set up the 2d viewport for mode text printing
		glPushMatrix();
		glLoadIdentity();
		vX = glutGet(GLUT_WINDOW_WIDTH);								// get the current window width
		vY = glutGet(GLUT_WINDOW_HEIGHT);								// get the current window height
		glViewport(0, 0, vX, vY);										// locate to left bottom corner
		gluOrtho2D(0, vX, 0, vY);										// define othogonal aspect

		glMatrixMode(GL_MODELVIEW);
		glPushMatrix();
		glLoadIdentity();

		float step = (vY - 3 * border_factor);
		step /= num_edge;
		for (unsigned i = 0; i < num_edge; i++) {
			glLineWidth(border_factor);
			glBegin(GL_LINES);
			glColor3f((float)color[velocity_bar[i] * 3 + 0] / 255, (float)color[velocity_bar[i] * 3 + 1] / 255, (float)color[velocity_bar[i] * 3 + 2] / 255);
			glVertex2f(border_factor, border_factor + i * step);
			glVertex2f(border_factor, border_factor + (i + 1) * step);
			glEnd();
		}
		glFlush();

		// pressure bar text
		glColor3f(1.0f, 1.0f, 1.0f);
		glRasterPos2f(0.0f, vY - border_factor);
		std::stringstream ss_p;
		ss_p << "Velocity range";
		glutBitmapString(GLUT_BITMAP_HELVETICA_18, (const unsigned char*)(ss_p.str().c_str()));

		// pressure range text
		step = vY - 3 * border_factor;
		step /= 10;
		for (unsigned i = 0; i < 11; i++) {
			glRasterPos2f(border_factor * 1.5f, border_factor + i * step);
			std::stringstream ss_n;
			ss_n << min_v + i * (max_v - min_v) / 10;
			glutBitmapString(GLUT_BITMAP_TIMES_ROMAN_10, (const unsigned char*)(ss_n.str().c_str()));
		}
		glPopMatrix();
		glMatrixMode(GL_PROJECTION);
		glPopMatrix();
	}
	
	glutSwapBuffers();
}

// register glut menu options
void glut_menu(int value) {

	int num = glutGet(GLUT_MENU_NUM_ITEMS);
	if (value == 1) {
		simulation = true;
		build_inlet_outlet = false;
		manufacture = false;
		modified_bridge = false;
		flow_initialize();
		flow_stable_state();					// main function of solving the linear system
		flow.print_flow();
		
		glut_set_menu(num, 2);
	}

	if (value == 2) {
		simulation = false;
		build_inlet_outlet = true;
		manufacture = false;
		if (!modified_bridge) {
			flow.set_main_feeder();
			flow.build_synthetic_connection(u);
		}

		glut_set_menu(num, 3);
	}

	if (value == 3) {
		simulation = false;
		build_inlet_outlet = false;
		manufacture = true;
		flow.check_synthetic_connection(u);
	}

	glutPostRedisplay();
}

// defines camera motion based on mouse dragging
void glut_motion(int x, int y) {

	mods = glutGetModifiers();
	if (LTbutton && mods == 0) {

		float theta = orbit_factor * (mouse_x - x);		// determine the number of degrees along the x-axis to rotate
		float phi = orbit_factor * (y - mouse_y);		// number of degrees along the y-axis to rotate

		cam.OrbitFocus(theta, phi);						// rotate the camera around the focal point
	}
	mouse_x = x;										// update the mouse position
	mouse_y = y;

	glutPostRedisplay();								// re-draw the visualization
}

// get click window coordinates
void glut_mouse(int button, int state, int x, int y) {

	mouse_x = x;
	mouse_y = y;
	if (button == GLUT_LEFT_BUTTON && state == GLUT_DOWN)
		LTbutton = true;
	else if (button == GLUT_LEFT_BUTTON && state == GLUT_UP)
		LTbutton = false;

	if (button == GLUT_LEFT_BUTTON && state == GLUT_DOWN && simulation && !to_select_pressure) {

		GLdouble posX, posY, posZ;
		window_to_world(posX, posY, posZ);			// get the world coordinates

		bool flag = flow.epsilon_vertex((float)posX, (float)posY, (float)posZ, eps, pressure_index);
		if (flag) {
			std::vector<unsigned>::iterator it = std::find(pendant_vertex.begin(), pendant_vertex.end(), pressure_index);
			if (it != pendant_vertex.end()) 		// if it is dangle vertex
				to_select_pressure = true;
		}
	}
	else if (button == GLUT_LEFT_BUTTON && state == GLUT_DOWN && simulation && to_select_pressure) {
		if (y > 2 * border_factor || y < vY - border_factor) {		// within the pressure bar range
			to_select_pressure = false;
			float tmp_pressure = (float)(vY - y - border_factor) / ((float)vY - border_factor) * max_pressure;
			flow.set_pressure(pressure_index, tmp_pressure);

			flow_stable_state();									// main function of solving the linear system
			flow.print_flow();
		}
	}
}

// define camera move based on mouse wheel move
void glut_wheel(int wheel, int direction, int x, int y) {
	
	mouse_x = x;
	mouse_y = y;

	GLdouble posX, posY, posZ;
	window_to_world(posX, posY, posZ);			// get the world coordinates

	if (!to_select_pressure) {
		bool flag = flow.epsilon_vertex((float)posX, (float)posY, (float)posZ, eps, pressure_index);
		if (flag && simulation) {
			float tmp_r;
			if (direction > 0) {				// increase radii
				tmp_r = flow.get_radius(pressure_index);
				tmp_r += radii_factor;
			}
			else {
				tmp_r = flow.get_radius(pressure_index);
				tmp_r -= radii_factor;
				if (tmp_r <= 0)
					tmp_r = default_radii;
			}
			flow.set_radius(pressure_index, tmp_r);
			flow_stable_state();
			flow.print_flow();
		}
		else {
			if (direction > 0)								// if it is button 3(up), move closer
				move_pace = zoom_factor;
			else											// if it is button 4(down), leave farther
				move_pace = -zoom_factor;

			cam.Push(move_pace);
		}
	}
	
	glutPostRedisplay();
}

// define keyboard inputs
void glut_keyboard(unsigned char key, int x, int y) {

	// register different keyboard operation
	switch (key) {

		// zooming
	case 'w':						// if keyboard 'w' is pressed, then move closer
		move_pace = zoom_factor;
		cam.Push(move_pace);
		break;
	case 's':						// if keyboard 's' is pressed, then leave farther
		move_pace = -zoom_factor;
		cam.Push(move_pace);
		break;

		// output image stack
	case 'm':
		if (manufacture) {
#ifdef __CUDACC__
			flow.make_image_stack(dx, dy, dz, stackdir);
#else
			std::cout << "You need to have a gpu to make image stack, sorry." << std::endl;
#endif
		}
		else if (build_inlet_outlet && !modified_bridge) {
			flow.modify_synthetic_connection(u, rou);
			modified_bridge = true;
		}
		break;
	}

	glutPostRedisplay();
}

// glut initialization
void glut_initialize() {

	int myargc = 1;
	char* myargv[1];
	myargv[0] = strdup("generate_network_network");

	glutInit(&myargc, myargv);
	glutInitDisplayMode(GLUT_DEPTH | GLUT_DOUBLE | GLUT_RGBA);
	glutInitWindowPosition(100, 100);							// set the initial window position
	glutInitWindowSize(1000, 1000);
	glutCreateWindow("3D flow simulation");

	glutDisplayFunc(glut_render);
	glutMouseFunc(glut_mouse);
	glutMotionFunc(glut_motion);
	glutMouseWheelFunc(glut_wheel);
	glutKeyboardFunc(glut_keyboard);

	glutCreateMenu(glut_menu);					// create a menu object 
	glut_set_menu(0, 1);
	glutAttachMenu(GLUT_RIGHT_BUTTON);			// register right mouse to open menu option

	stim::vec3<float> c = bb.center();			// get the center of the network bounding box
	// place the camera along the z-axis at a distance determined by the network size along x and y
	cam.setPosition(c + stim::vec<float>(0, 0, camera_factor * std::max(bb.size()[0], bb.size()[1])));
	cam.LookAt(c[0], c[1], c[2]);
}

int main(int argc, char* argv[]) {
	
	// add arguments
	args.add("help", "prints the help");
	args.add("network", "load network from .obj or .swc file");
	args.add("maxpress", "maximum allowed pressure in g / units / s^2, default 2 is for blood when units = um", "2", "real value > 0");
	args.add("viscosity", "set the viscosity of the fluid (in g / units / s), default .00001 is for blood when units = um", ".00001", "real value > 0");
	args.add("rou", "set the desity of the fluid (in g / units^3), default 1.06*10^-12 is for blood when units = um", ".00000000000106", "real value > 0");
	args.add("stackres", "spacing between pixel samples in each dimension(in units/pixel)", "1 1 1", "real value > 0");
	args.add("stackdir", "set the directory of the output image stack", "", "any existing directory (ex. /home/name/network)");
	
	args.parse(argc, argv);								// parse the command line

	// load network
	if (args["network"].is_set()) {						// load network from user 
		std::vector<std::string> tmp = stim::parser::split(args["network"].as_string(), '.');
		if ("obj" == tmp[1])
			flow.load_obj(args["network"].as_string());
		else if ("swc" == tmp[1])
			flow.load_swc(args["network"].as_string());
		else {
			std::cout << "Invalid file type" << std::endl;
			std::exit(1);
		}
	}
	get_background();

	// blood pressure in capillaries range from 15 - 35 torr
	// 1 torr = 133.3 Pa
	max_pressure = args["maxpress"].as_float();

	// normal blood viscosity range from 4 - 15 mPa·s(cP)
	// 1 Pa·s = 1 g / mm / s
	u = args["viscosity"].as_float();			// g / units / s

	// normally the blood density in capillaries: 1060 kg/m^3 = 1.06*10^-12 g/um^3
	rou = args["rou"].as_float();

	// get the vexel and image stack size
	dx = args["stackres"].as_float(0);
	dy = args["stackres"].as_float(1);
	dz = args["stackres"].as_float(2);

	// get the save directory of image stack
	if (args["stackdir"].is_set())
		stackdir = args["stackdir"].as_string();

	// glut main loop
	bb = flow.boundingbox();
	glut_initialize();
	glutMainLoop();
}