// Copyright (C) 2014 ChaosForge Ltd
// http://chaosforge.org/
//
// This file is part of NV Libraries.
// For conditions of distribution and use, see copyright notice in nv.hh

#include "nv/formats/assimp_loader.hh"
#include <unordered_map>
#include <glm/gtx/transform.hpp>
#include "nv/io/std_stream.hh"
#include "nv/lib/assimp.hh"

using namespace nv;


nv::assimp_loader::assimp_loader( const string& a_ext, const mat4& a_rotate_transform, float a_scale, uint32 a_assimp_flags /*= 0 */ ) : m_ext( a_ext ), m_rotate_transform( a_rotate_transform ), m_scale( a_scale ), m_assimp_flags( a_assimp_flags ), m_mesh_count(0), m_scene( nullptr )
{
	if ( m_assimp_flags == 0 )
	{
		m_assimp_flags = ( 
			aiProcess_CalcTangentSpace				|  
			aiProcess_GenSmoothNormals				|  
			aiProcess_JoinIdenticalVertices			|   
			aiProcess_ImproveCacheLocality			|  
			aiProcess_LimitBoneWeights				|  
			aiProcess_RemoveRedundantMaterials      |  
			aiProcess_SplitLargeMeshes				|  
			aiProcess_Triangulate					|  
			aiProcess_GenUVCoords                   |  
			aiProcess_SortByPType                   |  
			aiProcess_FindDegenerates               |  
			aiProcess_FindInvalidData               |  
			0 );
	}
}


bool nv::assimp_loader::load( stream& source )
{
	if ( m_scene != nullptr ) aiReleaseImport( (const aiScene*)m_scene );
	m_scene = nullptr;
	m_mesh_count = 0;
	load_assimp_library();
	NV_LOG( nv::LOG_NOTICE, "AssImp loading file..." );
	int size = (int)source.size();
	char* data  = new char[ size ];
	source.read( data, size, 1 );
	const aiScene* scene = aiImportFileFromMemory( data, size, m_assimp_flags, m_ext.c_str() );

	if( !scene)
	{
		NV_LOG( nv::LOG_ERROR, aiGetErrorString() );
		return false;
	}
	m_scene      = scene;
	m_mesh_count = scene->mNumMeshes;
	NV_LOG( nv::LOG_NOTICE, "Loading successfull" );
	return true;
}

mesh_data* nv::assimp_loader::release_mesh_data( size_t index /*= 0 */ )
{
	if ( index >= m_mesh_count ) return nullptr;
	mesh_data* result = new mesh_data();
	const aiScene* scene = (const aiScene*)m_scene;
	const aiMesh*  mesh  = scene->mMeshes[ index ];

	mat3 scaled_rotatation = glm::mat3( glm::scale( m_scale, m_scale, m_scale ) * m_rotate_transform );
	vec3 vertex_offset     = glm::vec3(); 
	mat3 vertex_transform  = glm::mat3( glm::transpose( scaled_rotatation ) );
	mat3 normal_transform  = glm::mat3( glm::transpose( glm::inverse( scaled_rotatation ) ) );
	mat4 bone_transform    = glm::scale( 1.0f/m_scale, 1.0f/m_scale, 1.0f/m_scale ) * m_rotate_transform;

	bool skinned = mesh->mNumBones > 0;
	mesh_raw_channel* channel = nullptr;
	if ( skinned )
		channel = mesh_raw_channel::create< assimp_skinned_vtx >( mesh->mNumVertices );
	else
		channel = mesh_raw_channel::create< assimp_plain_vtx >( mesh->mNumVertices );

	result->add_channel( channel );
	for (unsigned int i=0; i<mesh->mNumVertices; i++)
	{
		vec3 v = vertex_transform * assimp_vec3_cast( mesh->mVertices[ i ] ) + vertex_offset;
		vec3 n = glm::normalize( normal_transform * assimp_vec3_cast( mesh->mNormals[ i ] ) );
		vec3 t = glm::normalize( normal_transform * assimp_vec3_cast( mesh->mTangents[ i ] ) );
		vec3 b = glm::normalize( normal_transform * assimp_vec3_cast( mesh->mBitangents[ i ] ) );
		vec2 s = assimp_st_cast( mesh->mTextureCoords[ 0 ][ i ] );

		glm::vec3 t_i = glm::normalize (t - n * glm::dot (n, t));
		float det = (glm::dot (glm::cross (n, t), b));
		det = (det < 0.0f ? -1.0f : 1.0f );
		nv::vec4 vt ( t_i[0], t_i[1], t_i[2], det );
		if ( skinned )
			((assimp_skinned_vtx*)channel->data)[i] = assimp_skinned_vtx( v, s, n, vt );
		else
			((assimp_plain_vtx*)channel->data)[i] = assimp_plain_vtx( v, s, n, vt );
	}

	if ( skinned )
	{
		assimp_skinned_vtx* vtx = (assimp_skinned_vtx*)channel->data;
		for (unsigned int m=0; m<mesh->mNumBones; m++)
		{
			aiBone* bone  = mesh->mBones[m];
			nv::mat4 offset = nv::assimp_mat4_cast( bone->mOffsetMatrix ) * bone_transform;
			for (unsigned int w=0; w<bone->mNumWeights; w++)
			{
				assimp_skinned_vtx& v = vtx[ bone->mWeights[w].mVertexId ];
				bool found = false;
				for (nv::uint32 i = 0 ; i < 4; ++i)
				{
					if ( v.boneweight[i] <= 0.0f ) 
					{
						v.boneindex[i] = m;
						v.boneweight[i] = bone->mWeights[w].mWeight;
						found = true;
						break;
					}
				}
				NV_ASSERT( found, "Too many weights!" );
			}
		}
	}

	mesh_raw_index_channel* ichannel = mesh_raw_index_channel::create( USHORT, mesh->mNumFaces * 3 );
	result->set_index_channel( ichannel );
	uint16* indices = (uint16*)ichannel->data;
	for (unsigned int i=0; i<mesh->mNumFaces; i++)
	{
		const aiFace* face = &mesh->mFaces[i];
		for (unsigned int j=0; j<face->mNumIndices; j++)
		{
			indices[ i*3 + j ] = (uint16)face->mIndices[j];
		}
	}

	return result;
}

nv::assimp_loader::~assimp_loader()
{
	if ( m_scene != nullptr ) aiReleaseImport( (const aiScene*)m_scene );
}

bool nv::assimp_loader::load_bones( size_t index, std::vector< assimp_bone >& bones )
{
	if ( m_scene == nullptr ) return false;
	const aiScene* scene = (const aiScene*)m_scene;
	const aiMesh*  mesh  = scene->mMeshes[ index ];
	if ( mesh->mNumBones == 0 ) return false;

	mat4 bone_transform    = glm::scale( 1.0f/m_scale, 1.0f/m_scale, 1.0f/m_scale ) * m_rotate_transform;

	for (unsigned int m=0; m<mesh->mNumBones; m++)
	{
		aiBone* bone   = mesh->mBones[m];
		mat4    offset = assimp_mat4_cast( bone->mOffsetMatrix ) * bone_transform;
		bones.emplace_back( bone->mName.data, offset );
	}
	return true;
}

void nv::assimp_loader::scene_report() const
{
	const aiScene* scene = (const aiScene*)m_scene;
	if ( scene == nullptr ) return;

	NV_LOG( nv::LOG_NOTICE, "------------------------" );
	NV_LOG( nv::LOG_NOTICE, "Texture   count - " << scene->mNumTextures );
	NV_LOG( nv::LOG_NOTICE, "Animation count - " << scene->mNumAnimations );
	NV_LOG( nv::LOG_NOTICE, "Material  count - " << scene->mNumMaterials );
	NV_LOG( nv::LOG_NOTICE, "Meshes    count - " << scene->mNumMeshes );
	NV_LOG( nv::LOG_NOTICE, "------------------------" );

	aiNode* root = scene->mRootNode;
	if (root)
	{
		NV_LOG( nv::LOG_NOTICE, "Root node  - " << root->mName.data );
		NV_LOG( nv::LOG_NOTICE, "  meshes   - " << root->mNumMeshes );
		NV_LOG( nv::LOG_NOTICE, "  children - " << root->mNumChildren );
	}
	else
	{
		NV_LOG( nv::LOG_NOTICE, "No root node!" );
	}
	NV_LOG( nv::LOG_NOTICE, "------------------------" );

	if ( scene->mNumMeshes > 0 )
	{
		for ( nv::uint32 mc = 0; mc < scene->mNumMeshes; mc++ )
		{
			aiMesh* mesh = scene->mMeshes[mc];

			NV_LOG( nv::LOG_NOTICE, "Mesh #"<<mc<<"   - " << mesh->mName.data );
			NV_LOG( nv::LOG_NOTICE, "  bones   - " << mesh->mNumBones );
			NV_LOG( nv::LOG_NOTICE, "  uvs     - " << mesh->mNumUVComponents[0] );
			NV_LOG( nv::LOG_NOTICE, "  verts   - " << mesh->mNumVertices );
			NV_LOG( nv::LOG_NOTICE, "  faces   - " << mesh->mNumFaces );

			// 			NV_LOG( nv::LOG_NOTICE, "Bones:" );
			// 			for (unsigned int m=0; m<mesh->mNumBones; m++)
			// 			{
			// 				aiBone* bone  = mesh->mBones[m];
			// 				NV_LOG( nv::LOG_DEBUG, bone->mName.C_Str() );
			// 			}
		}
	}
	else
	{
		NV_LOG( nv::LOG_NOTICE, "No meshes!" );
	}
	NV_LOG( nv::LOG_NOTICE, "------------------------" );


	// 	if ( scene->mNumMaterials > 0 )
	// 	{
	// 		for (unsigned int m=0; m < scene->mNumMaterials; m++)
	// 		{
	// 			int texIndex = 0;
	// 			aiReturn texFound = aiReturn_SUCCESS;
	// 			aiString path;	// filename
	//  			while (texFound == aiReturn_SUCCESS)
	//  			{
	// 				texFound = scene->mMaterials[m]->GetTexture(aiTextureType_DIFFUSE, texIndex, &path);
	// 				NV_LOG( nv::LOG_NOTICE, "  material - " << path.data );
	// 				texIndex++;
	// 			}
	// 		}
	// 	}
	// 	else
	// 	{
	// 		NV_LOG( nv::LOG_NOTICE, "No materials" );
	// 	}
	// 	NV_LOG( nv::LOG_NOTICE, "------------------------" );

}

assimp_model* nv::assimp_loader::release_merged_model()
{
	if ( m_scene == nullptr || m_mesh_count == 0 ) return nullptr;
	assimp_model* model = new assimp_model;

	for ( size_t m = 0; m < m_mesh_count; ++m )
	{
		model->meshes.push_back( release_mesh_data(m) );
	}

	std::unordered_map< std::string, uint16 > names;
	for ( unsigned int m = 0; m < model->meshes.size(); ++m )
	{
		nv::sint16 translate[64];
		std::vector< assimp_bone > bones;
		load_bones( m, bones );
		for ( unsigned int b = 0; b < bones.size(); ++b )
		{

			nv::assimp_bone& bone = bones[b];
			auto iname = names.find( bone.name );
			if ( iname == names.end() )
			{
				NV_ASSERT( model->bones.size() < 64, "Too many bones to merge!" );
				nv::sint16 index = (nv::sint16)model->bones.size();
				model->bones.push_back( bone );
				names[ bone.name ] = index;
				translate[b] = index;
			}
			else
			{
				translate[b] = (nv::sint16)iname->second;
			}
		}
		if ( m > 0 )
		{
			mesh_data* mesh = model->meshes[m];
			nv::mesh_raw_channel* channel = mesh->get_channel_data()[0];
			nv::assimp_skinned_vtx* va = (nv::assimp_skinned_vtx*)channel->data;
			for ( unsigned v = 0; v < channel->count; ++v )
			{
				nv::assimp_skinned_vtx& vertex = va[v];

				for (nv::uint32 i = 0 ; i < 4; ++i)
				{
					if ( vertex.boneweight[i] > 0.0f ) 
					{
						vertex.boneindex[i] = translate[vertex.boneindex[i]];
					}
				}
			}
		}
	}
	return model;
}

assimp_animation* nv::assimp_loader::release_animation( size_t, bool pre_transform, const std::vector< assimp_bone >* bone_data )
{
	if ( m_scene == nullptr ) return nullptr;
	const aiScene* scene = (const aiScene*)m_scene;
	if ( scene->mRootNode == nullptr || scene->mAnimations[0] == nullptr ) return nullptr;
	assimp_animation* result = new assimp_animation;

	// need resize not to copy!
	result->nodes.resize( count_nodes( scene->mRootNode ) );

	const aiNode*      root = scene->mRootNode;
	const aiAnimation* anim = scene->mAnimations[0]; // if implemented, change in load_node also

	result->fps            = (float)anim->mTicksPerSecond;
	result->duration       = (float)anim->mDuration;
	result->pretransformed = pre_transform;

	load_node( result, root, 0, -1 );
	result->nodes[0].transform = glm::scale( m_scale, m_scale, m_scale ) * result->nodes[0].transform;

	if ( bone_data )
	{
		std::unordered_map< std::string, uint16 > names;
		for ( uint16 bi = 0; bi < bone_data->size(); ++bi )
		{
			names[ (*bone_data)[bi].name ] = bi;
		}

		for ( unsigned i = 0; i < result->nodes.size(); ++i )
		{
			assimp_animated_node_data& node = result->nodes[i];
			node.bone_id = -1;
			auto bi = names.find( node.name );
			if ( bi != names.end() )
			{
				node.bone_id = bi->second;
			}
			if ( node.parent_id != -1 ) 
			{
				result->nodes[ node.parent_id ].children.push_back( &node );
			}
		}
	}

	return result;
}

nv::uint32 nv::assimp_loader::count_nodes( const void* node ) const
{
	const aiNode* ainode = (const aiNode*)node;
	nv::uint32 count = 1;
	for ( unsigned i = 0; i < ainode->mNumChildren; ++i )
	{
		count += count_nodes( ainode->mChildren[i] );
	}
	return count;
}

nv::uint32 nv::assimp_loader::load_node( assimp_animation* data, const void* vnode, sint32 this_id, sint32 parent_id )
{
	const aiScene* scene = (const aiScene*)m_scene;
	const aiNode*  node  = (const aiNode*)vnode;
	string name( node->mName.data );
	const aiAnimation* anim  = scene->mAnimations[0];
	const aiNodeAnim*  anode = nullptr;

	for ( unsigned i = 0 ; i < anim->mNumChannels ; i++ )
	{
		anode = anim->mChannels[i];
		if ( std::string( anode->mNodeName.data ) == name ) break;
		anode = nullptr;
	}

	assimp_animated_node_data& a_data = data->nodes[ this_id ];

	a_data.name      = name;
	a_data.parent_id = parent_id;
	a_data.bone_id = -1;
	a_data.transform = nv::assimp_mat4_cast( node->mTransformation );

	if (anode) 
	{
		if ( data->pretransformed )
		{
			a_data.keys = create_transformed_keys( anode, parent_id >= 0 ? data->nodes[ parent_id ].keys : nullptr );
		}
		else
		{
			a_data.keys = create_direct_keys( anode );
		}
	}

	nv::uint32 next = this_id + 1;
	for ( unsigned i = 0; i < node->mNumChildren; ++i )
	{
		next = load_node( data, node->mChildren[i], next, this_id );
	}
	return next;
}

key_animation_data* nv::assimp_loader::create_transformed_keys( const void* vnode, const key_animation_data* parent_keys )
{
	const aiNodeAnim* node = (const aiNodeAnim*)vnode;
	size_t max_keys = glm::max( node->mNumPositionKeys, node->mNumRotationKeys );
	nv::transform_vector* keys = new nv::transform_vector;
	for ( unsigned n = 0; n < max_keys; ++n )
	{
		size_t pn = glm::min( node->mNumPositionKeys - 1, n );
		size_t rn = glm::min( node->mNumRotationKeys - 1, n );
		nv::vec3 pos = nv::assimp_vec3_cast(node->mPositionKeys[pn].mValue);
		nv::quat rot = nv::assimp_quat_cast(node->mRotationKeys[rn].mValue);
		nv::transform ptr;
		if ( parent_keys )
		{
			const nv::transform_vector* pv = (const nv::transform_vector*)parent_keys;
			if ( pv && pv->size() > 0 ) ptr = pv->get( glm::min( n, pv->size()-1 ) );
		}

		nv::vec3 rot_pos = ptr.get_orientation() * pos;
		pos = ptr.get_position() + rot_pos;
		rot = glm::normalize( ptr.get_orientation() * rot );
		keys->insert( nv::transform( pos, rot ) );
	}
	return keys;
}

key_animation_data* nv::assimp_loader::create_direct_keys( const void* vnode )
{
	const aiNodeAnim* node = (const aiNodeAnim*)vnode;
	if ( node->mNumPositionKeys == 0 && node->mNumRotationKeys == 0 && node->mNumScalingKeys == 0 ) return nullptr;
	key_vectors_prs* keys = new key_vectors_prs;

	for ( unsigned np = 0; np < node->mNumPositionKeys; ++np )
	{
		keys->insert_position( (float)node->mPositionKeys[np].mTime, assimp_vec3_cast(node->mPositionKeys[np].mValue) );
	}
	for ( unsigned np = 0; np < node->mNumRotationKeys; ++np )
	{
		keys->insert_rotation( (float)node->mRotationKeys[np].mTime, assimp_quat_cast(node->mRotationKeys[np].mValue) );
	}
	if ( node->mNumScalingKeys > 0 )
	{
		nv::vec3 scale_vec0 = assimp_vec3_cast( node->mScalingKeys[0].mValue );
		float scale_value   = glm::length( glm::abs( scale_vec0 - nv::vec3(1,1,1) ) );
		if ( node->mNumScalingKeys > 1 || scale_value > 0.001 ) 
		{
			NV_LOG( nv::LOG_WARNING, "scale key significant!" );
			for ( unsigned np = 0; np < node->mNumRotationKeys; ++np )
			{
				keys->insert_scale( (float)node->mScalingKeys[np].mTime, assimp_vec3_cast(node->mScalingKeys[np].mValue) );
			}
		}
	}
	return keys;
}