#include <GL/glew.h>
#include <GLFW/glfw3.h>

#include <iostream>
#include <vector>
const double PI = 3.14159265358979323846;

#include "tensorfield.h"
#include <stim/math/quaternion.h>
//#include <stim/>

// Global variable to store the GLFW window
GLFWwindow* window;
tensorfield<float> T;
size_t z_slice = 0;
bool cout_frame = false;

float diffuse_intensity = 0.7f;
float ambient_intensity = 0.3f;

int glyph_type = 1;
float gamma = 3;
int glyph_resolution = 10;
float glyph_scale = 0.5;

//glyph vertices
bool glyph_calculated = false;

struct float3 {
	float x;
	float y;
	float z;
};

struct float2 {
	float u;
	float v;
};

std::vector<float3> vertices;
std::vector<float3> normals;
std::vector<float2> texcoords;
std::vector<float> sin_theta;
std::vector<float> cos_theta;
std::vector<float> sin_phi;
std::vector<float> cos_phi;

struct eigendecomposition {
	
	stim::vec3<float> v0;
	stim::vec3<float> v2;
	stim::vec3<float> lambda;
};

inline float sgn(float x) {
	if (x < 0)
		return -1;
	else if (x > 0)
		return 1;
	return 0;
}

stim::vec3<float> triangle_norm(stim::vec3<float> p0, stim::vec3<float> p1, stim::vec3<float> p2) {
	stim::vec3<float> a = p1 - p0;
	stim::vec3<float> b = p2 - p0;
	stim::vec3<float> n = a.cross(b);
	return n;
}


// returns the fractional anisotropy and stores the spherical, linear, and planar components in cs, cl, and cp
float get_anisotropy(stim::vec3<float> lambda, float& cs, float& cl, float& cp) {

	//calculate the denominator for each specific anisotropy
	float denom = lambda[0] + lambda[1] + lambda[2];

	cl = (lambda[2] - lambda[1]) / denom;
	cp = 2 * (lambda[1] - lambda[0]) / denom;
	cs = 3 * lambda[0] / denom;

	return 0;
}

//calculate the eigenvectors and eigenvalues for the input pixel (x, y, z)
//  eigenvalues are embedded in the length of the eigenvector
void get_pixel_eigen(eigendecomposition* e, size_t x, size_t y, size_t z) {

	e->v0[0] = T(x, y, z, 0, 0);
	e->v0[1] = T(x, y, z, 1, 0);
	e->v0[2] = T(x, y, z, 2, 0);
	
	//temporarily store v1 in order to get the eigenvalue
	//	(the eigenvector is redundant because it's the cross product of v0 and v1)
	stim::vec3<float> v1;
	v1[0] = T(x, y, z, 0, 1);
	v1[1] = T(x, y, z, 1, 1);
	v1[2] = T(x, y, z, 2, 1);

	e->v2[0] = T(x, y, z, 0, 2);
	e->v2[1] = T(x, y, z, 1, 2);
	e->v2[2] = T(x, y, z, 2, 2);

	e->lambda[0] = e->v0.len();
	e->lambda[1] = v1.len();
	e->lambda[2] = e->v2.len();
	e->v0 = e->v0 / e->lambda[0];
	e->v2 = e->v2 / e->lambda[2];
}

/// Create a rotation matrix to orient a glyph along the tensor direction. Orientations are based on the input vectors
///		v0, v1, and v2
void get_glyph_rotation_matrix(float* R, stim::vec3<float> v0, stim::vec3<float> v2) {

	stim::matrix_sq<float, 4> M;
	M(0, 0) = v0[0];
	M(1, 0) = v0[1];
	M(2, 0) = v0[2];
	M(3, 0) = 0;

	M(0, 2) = v2[0];
	M(1, 2) = v2[1];
	M(2, 2) = v2[2];
	M(3, 2) = 0;

	stim::vec3<float> a(v0[0], v0[1], v0[2]);
	stim::vec3<float> b(v2[0], v2[1], v2[2]);
	stim::vec3<float> c = a.cross(b);

	M(0, 1) = c[0];
	M(1, 1) = c[1];
	M(2, 1) = c[2];
	M(3, 1) = 0;

	M(0, 3) = 0;
	M(1, 3) = 0;
	M(2, 3) = 0;
	M(3, 3) = 1;

	memcpy(R, M.M, 16 * sizeof(float));
	return;
}


// Perform the necessary updates when the user reshapes the window
void window_reshape() {
	int width, height;

	//get the context size
	glfwGetFramebufferSize(window, &width, &height);

	
	//create an OpenGL viewport
	glViewport(0, 0, width, height);

	//set the default orthographic view (assuming that the window and image aspect ratio are identical)
	float left = 0.0f;
	float right = (float)T.shape[2];
	float bottom = 0.0f;
	float top = (float)T.shape[1];

	//create an orthographic projection
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluOrtho2D(left, right, bottom, top);

	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
}

void init() {
	glEnable(GL_DEPTH_TEST);
	glEnable(GL_CULL_FACE);
	glFrontFace(GL_CCW);
}

void generate_glyph_points(int resolution = 5) {
	int sectorCount = 2 * resolution;
	int stackCount = resolution;

	int num_vertices = (sectorCount + 1) * (stackCount + 1);
	sin_theta.resize(num_vertices);
	sin_phi.resize(num_vertices);
	cos_theta.resize(num_vertices);
	cos_phi.resize(num_vertices);

	float theta, phi;                                     // vertex texCoord

	float sectorStep = 2 * (float)PI / sectorCount;
	float stackStep = (float)PI / stackCount;

	size_t i = 0;
	//calculate the vertex positions
	for (int phi_i = 0; phi_i <= stackCount; ++phi_i)
	{

		// add (sectorCount+1) vertices per stack
		// the first and last vertices have same position and normal, but different tex coords
		for (int theta_i = 0; theta_i <= sectorCount; ++theta_i)
		{
			
			// calculate the spherical coordinates of the vertex
			theta = (float)theta_i / sectorCount * 2 * (float)PI;
			phi = (float)phi_i / stackCount * (float)PI;

			//pre-compute the sine and cosine values used to create the superquadric
			sin_theta[i] = sinf(theta);
			sin_phi[i] = sinf(phi);
			cos_theta[i] = cosf(theta);
			cos_phi[i] = cosf(phi);
			i++;
		}
	}
}

inline float weird_exp(float x, float e) {
	return sgn(x) * powf(fabs(x), e);
}

stim::vec3<float> qx(size_t i, float alpha = 1, float beta = 1) {

	float sin_phi_beta = weird_exp(sin_phi[i], beta);
	float sin_theta_alpha = weird_exp(sin_theta[i], alpha);
	float cos_theta_alpha = weird_exp(cos_theta[i], alpha);
	float cos_phi_beta = weird_exp(cos_phi[i], beta);

	stim::vec3<float> p;
	p[0] = cos_phi_beta;
	p[1] = -sin_theta_alpha * sin_phi_beta;
	p[2] = cos_theta_alpha * sin_phi_beta;
	return p;
}

void render_triangle(stim::vec3<float> p0, stim::vec3<float> p1, stim::vec3<float> p2) {
	stim::vec3<float> n = triangle_norm(p0, p1, p2);

	glNormal3f(n[0], n[1], n[2]);
	glVertex3f(p0[0], p0[1], p0[2]);
	glVertex3f(p1[0], p1[1], p1[2]);
	glVertex3f(p2[0], p2[1], p2[2]);
}

//render a glyph (0 = ellipsoid, 1 = superquadric)
void render_glyph(eigendecomposition* e, int glyph_type = 0, int resolution = 5) {
	glMatrixMode(GL_MODELVIEW);
	glPushMatrix();

	stim::vec3<float> norm_lambda = e->lambda.norm();
	glScalef(glyph_scale, glyph_scale, glyph_scale);
	glScalef(norm_lambda[2], norm_lambda[1], norm_lambda[0]);

	if (glyph_type == 0) {
		GLUquadric* q = gluNewQuadric();
		gluSphere(q, 1, 2 * resolution, resolution);
	}
	else if (glyph_type == 1) {
		float radius = 0.5;
		int sectorCount = 2 * resolution;
		int stackCount = resolution;

		if (!glyph_calculated) {
			generate_glyph_points(resolution);
			glyph_calculated = true;
		}

		//get the anisotropy values
		float fa, cs, cl, cp, alpha, beta;
		fa = get_anisotropy(e->lambda, cs, cl, cp);
		
		if (cl >= cp) {
			alpha = powf(1 - cp, gamma);
			beta = powf(1 - cl, gamma);
		}
		else {
			alpha = powf(1 - cl, gamma);
			beta = powf(1 - cp, gamma);
		}

		//draw the sphere
		int k1, k2;
		float3 p[3];
		float2 s[3];
		float3 n;
		
		//draw the glyph
		glBegin(GL_TRIANGLES);
		for (int i = 0; i < stackCount; ++i)
		{
			k1 = i * (sectorCount + 1);     // beginning of current stack
			k2 = k1 + sectorCount + 1;      // beginning of next stack

			for (int j = 0; j < sectorCount; ++j, ++k1, ++k2)
			{
				// 2 triangles per sector excluding first and last stacks
				// k1 => k2 => k1+1
				if (i != 0)
				{

					stim::vec3<float> p0 = qx(k1, alpha, beta);
					stim::vec3<float> p1 = qx(k2, alpha, beta);
					stim::vec3<float> p2 = qx(k1 + 1, alpha, beta);
	
					render_triangle(p0, p1, p2);
				}

				// k1+1 => k2 => k2+1
				if (i != (stackCount - 1))
				{
					stim::vec3<float> p0 = qx(k1 + 1, alpha, beta);
					stim::vec3<float> p1 = qx(k2, alpha, beta);
					stim::vec3<float> p2 = qx(k2 + 1, alpha, beta);

					render_triangle(p0, p1, p2);
				}
			}
		}
		glEnd();
	}
	glPopMatrix();
}

void lighting() {

	GLfloat light_position[] = { 1, 1, -1.0, 0.0 };
	GLfloat light_ambient[] = { ambient_intensity, ambient_intensity, ambient_intensity, 1.0 };
	GLfloat light_diffuse[] = { diffuse_intensity, diffuse_intensity, diffuse_intensity, 1.0 };
	glShadeModel(GL_SMOOTH);

	glLightfv(GL_LIGHT0, GL_POSITION, light_position);
	glLightfv(GL_LIGHT0, GL_DIFFUSE, light_diffuse);
	glLightfv(GL_LIGHT0, GL_AMBIENT, light_ambient);

	glEnable(GL_LIGHTING);
	glEnable(GL_LIGHT0);

	glColorMaterial(GL_FRONT, GL_AMBIENT_AND_DIFFUSE);
	glEnable(GL_COLOR_MATERIAL);
	glEnable(GL_NORMALIZE);

}
void render() {
	size_t zi = z_slice;

	float dc = 1.0f / T.shape[1];
	float c[3];
	eigendecomposition e;
	float R[16];

	
	glMatrixMode(GL_MODELVIEW_MATRIX);

	lighting();
	
	for (size_t yi = 0; yi < T.shape[1]; yi++) {
		for (size_t xi = 0; xi < T.shape[2]; xi++) {
			get_pixel_eigen(&e, xi, yi, zi);
			
			c[0] = abs(e.v0[0]);
			c[1] = abs(e.v0[1]);
			c[2] = abs(e.v0[2]);

			glPushMatrix();
			
			
			glTranslatef((float)xi + 0.5f, (float)yi + 0.5f, 0);

			stim::vec3<float> v0(e.v0[0], e.v0[1], e.v0[2]);
			stim::vec3<float> v2(e.v2[0], e.v2[1], e.v2[2]);
			get_glyph_rotation_matrix(R, v0,v2);
			glMultMatrixf((GLfloat*)R);

			glColor3f(c[0], c[1], c[2]);
			render_glyph(&e, glyph_type, glyph_resolution);

			
			glPopMatrix();
		}
	}
	cout_frame = false;
	
}

void key_callback(GLFWwindow* window, int key, int scancode, int action, int mods)
{
	if (key == GLFW_KEY_RIGHT && action != GLFW_RELEASE)
		z_slice++;
	else if (key == GLFW_KEY_LEFT && action != GLFW_RELEASE)
		z_slice--;
	if (z_slice >= T.shape[0]) z_slice = 0;
	if (z_slice < 0) z_slice = T.shape[0] - 1;
}

void display_rotation_matrix(eigendecomposition* e) {
	float R[16];
	stim::vec3<float> v0(e->v0[0], e->v0[1], e->v0[2]);
	stim::vec3<float> v2(e->v2[0], e->v2[1], e->v2[2]);
	get_glyph_rotation_matrix(R, v0, v2);
	for (int r = 0; r < 4; r++) {
		for (int c = 0; c < 4; c++) {
			std::cout << R[r * 4 + c] << "     ";
		}
		std::cout << std::endl;
	}
}


void display_eigen(eigendecomposition* e) {
	std::cout << "v0 = (" << e->v0[0] << ", " << e->v0[1] << ", " << e->v0[2] << ")" << std::endl;
}

void mouse_button_callback(GLFWwindow* window, int button, int action, int mods) {
	//if the user clicks inside the window, display information about the tensor field
	if (button == GLFW_MOUSE_BUTTON_LEFT && action == GLFW_PRESS) {
		double xpos, ypos;
		glfwGetCursorPos(window, &xpos, &ypos);							//get the position of the mouse pointer
		int width, height;
		glfwGetFramebufferSize(window, &width, &height);
		int pixel_x = (int)(xpos / width * T.shape[2]);
		int pixel_y = T.shape[1] - (int)(ypos / height * (int)T.shape[1]) - 1;
		std::cout << "----------------------------" << std::endl;
		std::cout << "x: " << pixel_x << "   y: " << pixel_y << std::endl;
		std::cout << "----------------------------" << std::endl;
		eigendecomposition e;
		get_pixel_eigen(&e, pixel_x, pixel_y, z_slice);
		display_eigen(&e);
		std::cout << "===========" << std::endl;
		display_rotation_matrix(&e);
	}
}

void reshape_callback(GLFWwindow* window, int width, int height) {
	window_reshape();
	render();
}


int main(int argc, char** argv) {
	std::string filename(argv[1]);
	if (T.load_tira(filename) != TIRA_SUCCESS) {
		return -1;
	}

	/* Initialize the library */
	if (!glfwInit())
		return -1;


	/* Create a windowed mode window and its OpenGL context */
	window = glfwCreateWindow(700, 700, "Hello World", NULL, NULL);
	if (!window)
	{
		glfwTerminate();
		return -1;
	}

	/* Make the window's context current */
	glfwMakeContextCurrent(window);

	glfwSetKeyCallback(window, key_callback);
	glfwSetFramebufferSizeCallback(window, reshape_callback);
	glfwSetMouseButtonCallback(window, mouse_button_callback);

	GLenum err = glewInit();
	//deal with a GLEW initialization failure
	if (GLEW_OK != err)
		std::cout << "GLEW Error: " << glewGetErrorString(err) << std::endl;
	
	//set up the window viewport
	window_reshape();

	//initialize the OpenGL render details
	init();

	/* Loop until the user closes the window */
	while (!glfwWindowShouldClose(window))
	{
		/* Render here */
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		render();

		/* Swap front and back buffers */
		glfwSwapBuffers(window);

		/* Poll for and process events */
		glfwPollEvents();
	}

	glfwTerminate();
	return 0;
}