//  ========================================================================
//  COSC422: Advanced Computer Graphics;  University of Canterbury (2019)
//
//  FILE NAME: ArmyPilot.cpp
//
//  Press key '1' to toggle 90 degs model rotation about x-axis on/off.
//  ========================================================================

#include <iostream>
#include <map>
#include <GL/freeglut.h>
#include <IL/il.h>
using namespace std;

#include <assimp/cimport.h>
#include <assimp/types.h>
#include <assimp/scene.h>
#include <assimp/postprocess.h>
#include "assimp_extras.h"

//----------Globals----------------------------
const aiScene* scene = NULL;
const aiScene* animation = NULL;
float angle = 0;
aiVector3D scene_min, scene_max, scene_center;
bool modelRotn = true;
std::map<int, int> texIdMap;

//------------Modify the following as needed----------------------
float materialCol[4] = { 0.9, 0.9, 0.9, 1 };   //Default material colour (not used if model's colour is available)
bool replaceCol = false;					   //Change to 'true' to set the model's colour to the above colour
float lightPosn[4] = { 0, 50, 50, 1 };         //Default light's position
bool twoSidedLight = false;					   //Change to 'true' to enable two-sided lighting

int tDuration; //Animation duration in ticks
int currTick = 0; //current tick
float timeStep=50; //Animation time step = 50 m.sec
float ambient[4] = { 0.2, 0.2, 0.2, 1.0 };  //Ambient light
float white[4] = { 1, 1, 1, 1 };			//Light's colour

struct meshInit
{
    int mNumVertices;
    aiVector3D* mVertices;
    aiVector3D* mNormals;
};

meshInit* initData;

/*
aiMatrix4x4 get_interpolated_rotation(double tick, aiNodeAnim *node) {
    aiQuatKey prevKey;
    aiQuatKey nextKey;
    for(int i = 0; i < node->mNumRotationKeys; i++) {
        if(tick >= node->mRotationKeys[i].mTime) {
            prevKey = node->mRotationKeys[i];
            continue;
        }

        nextKey = node->mRotationKeys[i];

        double timeDiff = nextKey.mTime - prevKey.mTime;
        double currTime = tick - prevKey.mTime;
        double timeLerp = currTime / timeDiff;

        aiQuaternion rotnOut;
        aiQuaternion::Interpolate(rotnOut, prevKey.mValue, nextKey.mValue, timeLerp);
        return aiMatrix4x4(rotnOut.GetMatrix());
    }
}

aiMatrix4x4 get_interpolated_position(double tick, aiNodeAnim *node) {
    aiVectorKey prevKey;
    aiVectorKey nextKey;
    for(int i = 0; i < node->mNumPositionKeys; i++) {
        if(tick >= node->mPositionKeys[i].mTime) {
            prevKey = node->mPositionKeys[i];
            continue;
        }
        nextKey = node->mPositionKeys[i];

        double timeDiff = nextKey.mTime - prevKey.mTime;
        double currTime = tick - prevKey.mTime;
        double timeLerp = currTime / timeDiff;

        aiVector3D posnOut;
        posnOut = prevKey.mValue + (nextKey.mValue - prevKey.mValue) * float(timeLerp);

        aiMatrix4x4 posnMat;
        aiMatrix4x4::Translation(prevKey.mValue, posnMat);
        return posnMat;
    }
}*/
bool loadAnimation(const char* fileName) {

    animation = aiImportFile(fileName, aiProcessPreset_TargetRealtime_MaxQuality);
	if(scene == NULL) exit(1);

	//printSceneInfo(scene);
	//printMeshInfo(scene);
	//printTreeInfo(scene->mRootNode);
	//printBoneInfo(scene);
	//printAnimInfo(scene);  //WARNING:  This may generate a lengthy output if the model has animation data
    tDuration = animation->mAnimations[0]->mDuration;

}

//-------Loads model data from file and creates a scene object----------
bool loadModel(const char* fileName)
{

	scene = aiImportFile(fileName, aiProcessPreset_TargetRealtime_MaxQuality);
	if(scene == NULL) exit(1);
	//printSceneInfo(scene);
	//printMeshInfo(scene);
	//printTreeInfo(scene->mRootNode);
	//printBoneInfo(scene);
	//printAnimInfo(scene);  //WARNING:  This may generate a lengthy output if the model has animation data
	get_bounding_box(scene, &scene_min, &scene_max);

    initData = new meshInit[scene->mNumMeshes];
    // New Code
    for (int i = 0; i < scene->mNumMeshes; i++) {
        aiMesh *mesh = scene->mMeshes[i];
        int numVert = mesh->mNumVertices;
        (initData + i)->mNumVertices = numVert;
        (initData + i)->mVertices = new aiVector3D[numVert];
        (initData + i)->mNormals = new aiVector3D[numVert];

        for (int i_vert = 0; i_vert < mesh->mNumVertices; i_vert++) {
            (initData + i)->mVertices[i_vert] = mesh->mVertices[i_vert];
            (initData + i)->mNormals[i_vert] = mesh->mNormals[i_vert];
        }
    }




    //node = scene->mRootNode->FindNode()

	return true;
}


void transformVertices()
{
    for (int i = 0; i < scene->mNumMeshes; i++) {
        aiMesh *mesh = scene->mMeshes[i];

        for (int i_bone = 0; i_bone < mesh->mNumBones; i_bone++) {
            aiBone *bone = mesh->mBones[i_bone];
            aiNode *node = scene->mRootNode->FindNode(bone->mName);
            aiMatrix4x4 boneTransform = bone->mOffsetMatrix;

            while(node != nullptr) {
                std::string check(node->mName.C_Str());
                if(check != "free3dmodel_skeleton") {
                    boneTransform = node->mTransformation * boneTransform;
                }
                node = node->mParent;

            }

            aiMatrix4x4 boneTransformTranspose = boneTransform;
            boneTransformTranspose.Transpose();

            for (int k = 0; k < bone->mNumWeights; k++) {
                int vid = (bone->mWeights[k]).mVertexId;

                aiVector3D vert = (initData + i)->mVertices[vid];
                aiVector3D norm = (initData + i)->mNormals[vid];

                mesh->mVertices[vid] = boneTransform * vert;
                mesh->mNormals[vid] = boneTransformTranspose * norm;

                //aiVertexWeight weight = bone->mWeights[i_weight];
                //mesh->mVertices[weight.mVertexId] = (boneTransform * initData[i].mVertices[weight.mVertexId]);
                //mesh->mNormals[weight.mVertexId] = (boneTransformTranspose * initData[i].mNormals[weight.mVertexId]);
            }
        }
    }
}

//-------------Loads texture files using DevIL library-------------------------------
void loadGLTextures(const aiScene* scene)
{
	/* initialization of DevIL */
	ilInit();
	if (scene->HasTextures())
	{
		std::cout << "Support for meshes with embedded textures is not implemented" << endl;
		return;
	}

	/* scan scene's materials for textures */
	/* Simplified version: Retrieves only the first texture with index 0 if present*/
	for (unsigned int m = 0; m < scene->mNumMaterials; ++m)
	{
		aiString path;  // filename

		if (scene->mMaterials[m]->GetTexture(aiTextureType_DIFFUSE, 0, &path) == AI_SUCCESS)
		{
			glEnable(GL_TEXTURE_2D);
			ILuint imageId;
			GLuint texId;
			ilGenImages(1, &imageId);
			glGenTextures(1, &texId);
			texIdMap[m] = texId;   //store tex ID against material id in a hash map
			ilBindImage(imageId); /* Binding of DevIL image name */
			ilEnable(IL_ORIGIN_SET);
			ilOriginFunc(IL_ORIGIN_LOWER_LEFT);

            std::string string(path.C_Str());
            std::string proper_path = string.substr(string.rfind("/") + 1, string.length());

			if (ilLoadImage((ILstring)proper_path.c_str()))   //if success
			{
				/* Convert image to RGBA */
				ilConvertImage(IL_RGBA, IL_UNSIGNED_BYTE);

				/* Create and load textures to OpenGL */
				glBindTexture(GL_TEXTURE_2D, texId);
				glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
				glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
				glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, ilGetInteger(IL_IMAGE_WIDTH),
					ilGetInteger(IL_IMAGE_HEIGHT), 0, GL_RGBA, GL_UNSIGNED_BYTE,
					ilGetData());
				glTexEnvi(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE);
				cout << "Texture:" << path.data << " successfully loaded." << endl;
				glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
				glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
				glPixelStorei(GL_UNPACK_SKIP_PIXELS, 0);
				glPixelStorei(GL_UNPACK_SKIP_ROWS, 0);
			}
			else
			{
				cout << "Couldn't load Image: " << path.data << endl;
			}
		}
	}  //loop for material

}

// ------A recursive function to traverse scene graph and render each mesh----------
void render (const aiScene* sc, const aiNode* nd)
{
	aiMatrix4x4 m = nd->mTransformation;
	aiMesh* mesh;
	aiFace* face;
	aiMaterial* mtl;
	GLuint texId;
	aiColor4D diffuse;
	int meshIndex, materialIndex, textureIndex;

	aiTransposeMatrix4(&m);   //Convert to column-major order
	glPushMatrix();
	glMultMatrixf((float*)&m);   //Multiply by the transformation matrix for this node

	// Draw all meshes assigned to this node
	for (int n = 0; n < nd->mNumMeshes; n++)
	{
		meshIndex = nd->mMeshes[n];          //Get the mesh indices from the current node
		mesh = scene->mMeshes[meshIndex];    //Using mesh index, get the mesh object

		materialIndex = mesh->mMaterialIndex;  //Get material index attached to the mesh
		mtl = sc->mMaterials[materialIndex];
		if (replaceCol)
			glColor4fv(materialCol);   //User-defined colour
		else if (AI_SUCCESS == aiGetMaterialColor(mtl, AI_MATKEY_COLOR_DIFFUSE, &diffuse))  //Get material colour from model
			glColor4f(diffuse.r, diffuse.g, diffuse.b, 1.0);
		else
			glColor4fv(materialCol);   //Default material colour

		//Get the polygons from each mesh and draw them
		for (int k = 0; k < mesh->mNumFaces; k++)
		{
			face = &mesh->mFaces[k];
			GLenum face_mode;

			switch(face->mNumIndices)
			{
				case 1: face_mode = GL_POINTS; break;
				case 2: face_mode = GL_LINES; break;
				case 3: face_mode = GL_TRIANGLES; break;
				default: face_mode = GL_POLYGON; break;
			}

            if (mesh->HasTextureCoords(0)) {
                textureIndex = texIdMap[materialIndex];
                glBindTexture(GL_TEXTURE_2D, textureIndex);
            }

			glBegin(face_mode);

			for(int i = 0; i < face->mNumIndices; i++) {
				int vertexIndex = face->mIndices[i];

				if(mesh->HasVertexColors(0))
					glColor4fv((GLfloat*)&mesh->mColors[0][vertexIndex]);

				//Assign texture coordinates here
                if (mesh->HasTextureCoords(0)) {
                    glTexCoord2f(mesh->mTextureCoords[0][vertexIndex].x, mesh->mTextureCoords[0][vertexIndex].y);
                }

				if (mesh->HasNormals())
					glNormal3fv(&mesh->mNormals[vertexIndex].x);

				glVertex3fv(&mesh->mVertices[vertexIndex].x);
			}

			glEnd();
		}
	}

	// Draw all children
	for (int i = 0; i < nd->mNumChildren; i++)
		render(sc, nd->mChildren[i]);

	glPopMatrix();
}

//--------------------OpenGL initialization------------------------
void initialise()
{

    glClearColor(1.0f, 1.0f, 1.0f, 1.0f);
	glEnable(GL_LIGHTING);
	glEnable(GL_LIGHT0);
	glEnable(GL_DEPTH_TEST);
	glEnable(GL_NORMALIZE);
	glLightfv(GL_LIGHT0, GL_AMBIENT, ambient);
	glLightfv(GL_LIGHT0, GL_DIFFUSE, white);
	glLightfv(GL_LIGHT0, GL_SPECULAR, white);
	if (twoSidedLight) glLightModeli(GL_LIGHT_MODEL_TWO_SIDE, 1);

	glEnable(GL_COLOR_MATERIAL);
	glColorMaterial(GL_FRONT_AND_BACK, GL_AMBIENT_AND_DIFFUSE);
	glMaterialfv(GL_FRONT_AND_BACK, GL_SPECULAR, white);
	glMaterialf(GL_FRONT_AND_BACK, GL_SHININESS, 50);
	glColor4fv(materialCol);
	loadModel("mannequin.fbx");			//<<<-------------Specify input file name here
    loadAnimation("run.fbx");

	loadGLTextures(scene);
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluPerspective(35, 1, 1.0, 1000.0);
}

void updateNodeMatrices(int tick)
{
    aiAnimation* anim = animation->mAnimations[0];
    aiMatrix4x4 matPos, matRot, matProd;
    aiMatrix3x3 matRot3;
    aiNode* nd;

    for (int i = 0; i < anim->mNumChannels; i++) {

        matPos = aiMatrix4x4(); //Identity
        matRot = aiMatrix4x4();
        aiNodeAnim* ndAnim = anim->mChannels[i]; //Channel

        // Position
        aiVector3D posn = (ndAnim->mPositionKeys[0]).mValue;
        for(int j = 0; j < ndAnim->mNumPositionKeys; j++) {
            if(ndAnim->mPositionKeys[j].mTime == tick) {
                posn = ndAnim->mPositionKeys[j].mValue;
                break;
            } else if(ndAnim->mPositionKeys[j-1].mTime < tick && tick <= ndAnim->mPositionKeys[j].mTime){
                aiVector3D posn1 = (ndAnim->mPositionKeys[j-1]).mValue;
                aiVector3D posn2 = (ndAnim->mPositionKeys[j]).mValue;
                float time1 = (ndAnim->mPositionKeys[j-1]).mTime;
                float time2 = (ndAnim->mPositionKeys[j]).mTime;
                float factor = (tick - time1) / (time2 - time1);
                posn = posn1 + factor * (posn2 - posn1);
                break;
            }
        }

        // Rotation
        aiQuaternion rotn = (ndAnim->mRotationKeys[0]).mValue;
        for(int j_rot = 0; j_rot < ndAnim->mNumRotationKeys; j_rot++) {
            if(ndAnim->mRotationKeys[j_rot].mTime == tick) {
                rotn = ndAnim->mRotationKeys[j_rot].mValue;
                break;
            } else if(ndAnim->mRotationKeys[j_rot-1].mTime < tick && tick <= ndAnim->mRotationKeys[j_rot].mTime){
                aiQuaternion rotn1 = (ndAnim->mRotationKeys[j_rot-1]).mValue;
                aiQuaternion rotn2 = (ndAnim->mRotationKeys[j_rot]).mValue;
                float time1 = (ndAnim->mRotationKeys[j_rot-1]).mTime;
                float time2 = (ndAnim->mRotationKeys[j_rot]).mTime;
                float factor = (tick - time1) / (time2 - time1);
                rotn.Interpolate(rotn, rotn1, rotn2, factor);
                break;
            }
        }

        matPos.Translation(posn, matPos);
        matRot3 = rotn.GetMatrix();
        matRot = aiMatrix4x4(matRot3);

        // applies everything

        matProd = matPos * matRot;
        nd = scene->mRootNode->FindNode(ndAnim->mNodeName);
        nd->mTransformation = matProd;
    }
    transformVertices();

}



//----Timer callback for continuous rotation of the model about y-axis----
void update(int value)
{
	if (currTick < tDuration)
    {
        updateNodeMatrices(currTick);
        glutTimerFunc(timeStep, update, 0);
        currTick++;
    } else {
        currTick = 0;
        glutTimerFunc(timeStep, update, 0);
    }
    glutPostRedisplay();
}

//----Keyboard callback to toggle initial model orientation---
void keyboard(unsigned char key, int x, int y)
{
	//if(key == '1') modelRotn = !modelRotn;   //Enable/disable initial model rotation
	glutPostRedisplay();
}

void floor()
{
    /*
    glPushMatrix();
    glTranslatef(0, -1.2, -10);
    glScalef(0.3,1,0.3);
    	glEnable(GL_LIGHTING);
	float white[4] = {1., 1., 1., 1.};
	float black[4] = {0};

	glColor3f(0.2, 0.2, 0.2);   //The floor is gray in colour
	glNormal3f(0.0, 1.0, 0.0);*/
	//The floor is made up of several tiny squares on a 200x200 grid. Each square has a unit size.
    glDisable(GL_TEXTURE_2D);
    glNormal3f(0.0, 1.0, 0.0);
    glScalef(200,200,200);
	glBegin(GL_QUADS);
	for(int i = -150; i < 150; i++)
	{
		for(int j = -150;  j < 150; j++)
		{
            if((i+j) % 2 == 0) {
                glColor3f(1,1,1);
            } else {
                glColor3f(0,0,0.1);
            }
			glVertex3f(i, 0, j);
			glVertex3f(i, 0, j+1);
			glVertex3f(i+1, 0, j+1);
			glVertex3f(i+1, 0, j);
		}
	}
	glEnd();
    glEnable(GL_TEXTURE_2D);
    //glPopMatrix();


}

//------The main display function---------
//----The model is first drawn using a display list so that all GL commands are
//    stored for subsequent display updates.
void display()
{
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
	gluLookAt(0, 0, 3, 0, 0, -5, 0, 1, 0);
	glLightfv(GL_LIGHT0, GL_POSITION, lightPosn);

	glRotatef(angle, 0.f, 1.f ,0.f);  //Continuous rotation about the y-axis

	// scale the whole asset to fit into our view frustum
	float tmp = scene_max.x - scene_min.x;
	tmp = aisgl_max(scene_max.y - scene_min.y,tmp);
	tmp = aisgl_max(scene_max.z - scene_min.z,tmp);
	tmp = 1.f / tmp;


	float xc = (scene_min.x + scene_max.x)*0.5;
	float yc = (scene_min.y + scene_max.y)*0.5;
	float zc = (scene_min.z + scene_max.z)*0.5;
	// center the model

    glPushMatrix();
    glRotatef(-90,1,0,0);
    glTranslatef(0,0,-0.6);
    glScalef(tmp, tmp, tmp);
    glTranslatef(-xc, -yc, -zc);

    //Define shadow mat, multiply by thing below, translate up a bit
    render(scene, scene->mRootNode);
    glPopMatrix();
    glPushMatrix();
    glScalef(tmp, tmp, tmp);
    glTranslatef(-xc, -yc, -zc);

    floor();
    glPopMatrix();
	glutSwapBuffers();
}



int main(int argc, char** argv)
{
	glutInit(&argc, argv);
	glutInitDisplayMode(GLUT_RGB | GLUT_DOUBLE | GLUT_DEPTH);
	glutInitWindowSize(600, 600);
	glutCreateWindow("Mannequin");
	glutInitContextVersion (4, 2);
	glutInitContextProfile ( GLUT_CORE_PROFILE );

	initialise();
	glutDisplayFunc(display);
	glutTimerFunc(50, update, 0);
	glutKeyboardFunc(keyboard);
	glutMainLoop();

	aiReleaseImport(scene);
}