// Copyright (C) 2012-2013 ChaosForge / Kornel Kisielewicz
// 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/md5_loader.hh"

#include <glm/gtc/constants.hpp>
#include "nv/logging.hh"
#include "nv/io/std_stream.hh"
#include <cstring>

using namespace nv;

// based on http://tfc.duke.free.fr/coding/md5-specs-en.html

static void next_line( std::istream& stream )
{
	stream.ignore( std::numeric_limits<std::streamsize>::max(), '\n' );
}

static inline void discard( std::istream& stream, const std::string& token )
{
//	stream.ignore( std::numeric_limits<std::streamsize>::max(), ' ' );
	std::string discarded;
	stream >> discarded;
	assert( discarded == token );
}


static void remove_quotes( std::string& str )
{
	size_t n;
	while ( ( n = str.find('\"') ) != std::string::npos ) str.erase(n,1);
}

static void unit_quat_w( glm::quat& quat )
{
	float t = 1.0f - ( quat.x * quat.x ) - ( quat.y * quat.y ) - ( quat.z * quat.z );
	quat.w = ( t < 0.0f ? 0.0f : -sqrtf(t) );
}

bool md5_loader::load( stream& source )
{
	std_stream sstream( &source );
	std::string command;

	sstream >> command;
	while ( !sstream.eof() )
	{
		if ( command == "MD5Version" )
		{
			sstream >> m_md5_version;
			assert( m_md5_version == 10 );
		}
		else if ( command == "commandline" )
		{
			next_line( sstream ); 
		}
		else if ( command == "numJoints" )
		{
			sstream >> m_num_joints;
			m_joints.reserve( m_num_joints );
		}
		else if ( command == "numMeshes" )
		{
			sstream >> m_num_meshes;
			m_meshes.reserve( m_num_meshes );
		}
		else if ( command == "joints" )
		{
			discard( sstream, "{" );
			md5_joint joint;
			for ( size_t i = 0; i < m_num_joints; ++i )
			{
				sstream >> joint.name >> joint.parent_id;
				discard( sstream, "(" );
				sstream >> joint.pos.x >> joint.pos.y >> joint.pos.z;
				discard( sstream, ")" );
				discard( sstream, "(" );
				sstream >> joint.orient.x >> joint.orient.y >> joint.orient.z;
				remove_quotes( joint.name );
				unit_quat_w( joint.orient );
				m_joints.push_back( joint );
				next_line( sstream );
			}
			discard( sstream, "}" );
		}
		else if ( command == "mesh" )
		{
			md5_mesh_data* mesh = new md5_mesh_data();

			int num_verts, num_tris, num_weights;

			discard( sstream, "{" );
			sstream >> command;
			while ( command != "}" ) 
			{
				if ( command == "shader" )
				{
					sstream >> mesh->m_shader;
					remove_quotes( mesh->m_shader );
					// texturePath.replace_extension( ".tga" );
					next_line( sstream );
				}
				else if ( command == "numverts")
				{
					sstream >> num_verts; 

					{
						mesh_raw_channel* ch_pnt = mesh_raw_channel::create<md5_vtx_pnt>( num_verts );
						mesh_raw_channel* ch_t   = mesh_raw_channel::create<md5_vtx_t>( num_verts );
						mesh->m_pntdata          = (md5_vtx_pnt*)ch_pnt->data;
						mesh->m_tdata            = (md5_vtx_t*)ch_t->data;
						mesh->add_channel( ch_pnt );
						mesh->add_channel( ch_t );
					}
					mesh->m_vtx_data.resize( num_verts );

					next_line( sstream );
					std::string line;
					for ( int i = 0; i < num_verts; ++i )
					{
						md5_vtx_data& vdata = mesh->m_vtx_data[i];
						size_t weight_count;
						size_t start_weight;
						vec2 texcoord;

						std::getline( sstream, line );
						sscanf( line.c_str(), "%*s %*u ( %f %f ) %u %u", &(texcoord.x), &(texcoord.y), &(start_weight), &(weight_count) );
						vdata.start_weight = start_weight;
						vdata.weight_count = weight_count;
						mesh->m_tdata[i].texcoord = texcoord;
					}  
				}
				else if ( command == "numtris" )
				{
					sstream >> num_tris;

					mesh_raw_index_channel* ch_i = mesh_raw_index_channel::create<uint32>( num_tris * 3 );
					uint32* vtx_i                = (uint32*)ch_i->data;
					mesh->m_idata                = vtx_i;
					uint32 idx = 0;
					mesh->set_index_channel( ch_i );

					next_line( sstream );
					std::string line;
					for ( int i = 0; i < num_tris; ++i )
					{
						size_t ti0;
						size_t ti1;
						size_t ti2;

						std::getline( sstream, line );
						sscanf( line.c_str(), "%*s %*u %u %u %u )", &(ti0), &(ti1), &(ti2));

						vtx_i[idx++] = (uint32)ti0;
						vtx_i[idx++] = (uint32)ti1;
						vtx_i[idx++] = (uint32)ti2;
					}              
				}
				else if ( command == "numweights" )
				{
					sstream >> num_weights;
					mesh->m_weights.reserve( num_weights );
					next_line( sstream );
					std::string line;
					for ( int i = 0; i < num_weights; ++i )
					{
						md5_weight weight;

						std::getline( sstream, line );
						sscanf( line.c_str(), "%*s %*u %u %f ( %f %f %f )", &(weight.joint_id), &(weight.bias), &(weight.pos.x), &(weight.pos.y), &(weight.pos.z));
 						mesh->m_weights.push_back(weight);
					}
				}
				else
				{
					next_line( sstream );
				}

				sstream >> command;
			}

			prepare_mesh( mesh );

			m_meshes.push_back(mesh);
		}
		sstream >> command;
	}

	assert( m_joints.size() == m_num_joints );
	assert( m_meshes.size() == m_num_meshes );
	return true;
}

bool md5_loader::prepare_mesh( md5_mesh_data* mdata )
{
	uint32 vtx_count = mdata->m_vtx_data.size();
	md5_vtx_pnt* vtcs = mdata->m_pntdata;

	for ( uint32 i = 0; i < vtx_count; ++i )
	{
		md5_vtx_data& vdata = mdata->m_vtx_data[i];
		md5_vtx_pnt& vtc = vtcs[i];

		vtc.position = glm::vec3(0);
		vtc.normal   = glm::vec3(0);
		vtc.tangent  = glm::vec3(0);

		for ( size_t j = 0; j < vdata.weight_count; ++j )
		{
			md5_weight& weight = mdata->m_weights[vdata.start_weight + j];
			md5_joint&  joint  = m_joints[weight.joint_id];

			glm::vec3 rot_pos = joint.orient * weight.pos;

			vtc.position += ( joint.pos + rot_pos ) * weight.bias;
		}
	}

	// Prepare normals
	uint32 tri_count = mdata->get_count() / 3;
	for ( unsigned int i = 0; i < tri_count; ++i )
	{
		uint32 ti0 = mdata->m_idata[ i * 3 ];
		uint32 ti1 = mdata->m_idata[ i * 3 + 1 ];
		uint32 ti2 = mdata->m_idata[ i * 3 + 2 ];
 
		glm::vec3 v1 = vtcs[ ti0 ].position;
		glm::vec3 v2 = vtcs[ ti1 ].position;
		glm::vec3 v3 = vtcs[ ti2 ].position;
		glm::vec3 xyz1 = v3 - v1;
		glm::vec3 xyz2 = v2 - v1;

		glm::vec3 normal = glm::cross( xyz1, xyz2 );

		vtcs[ ti0 ].normal += normal;
		vtcs[ ti1 ].normal += normal;
		vtcs[ ti2 ].normal += normal;

		const vec2& w1 = mdata->m_tdata[ ti0 ].texcoord;
		const vec2& w2 = mdata->m_tdata[ ti1 ].texcoord;
		const vec2& w3 = mdata->m_tdata[ ti2 ].texcoord;

		vec2 st1 = w3 - w1;
		vec2 st2 = w2 - w1;

		float coef = 1.0f / (st1.x * st2.y - st2.x * st1.y);

		vec3 tangent = (( xyz1 * st2.y ) - ( xyz2 * st1.y )) * coef;

		vtcs[ ti0 ].tangent += tangent;
		vtcs[ ti1 ].tangent += tangent;
		vtcs[ ti2 ].tangent += tangent;
	}

	for ( size_t i = 0; i < vtx_count; ++i )
	{
		md5_vtx_data& vdata = mdata->m_vtx_data[i];

		glm::vec3 normal  = glm::normalize( vtcs[i].normal );
		glm::vec3 tangent = glm::normalize( vtcs[i].tangent );
		vtcs[i].normal   = normal;
		vtcs[i].tangent  = tangent;

 		vdata.normal  = glm::vec3(0);
 		vdata.tangent = glm::vec3(0);
 
 		for ( size_t j = 0; j < vdata.weight_count; ++j )
 		{
 			const md5_weight& weight = mdata->m_weights[vdata.start_weight + j];
 			const md5_joint&  joint  = m_joints[weight.joint_id];
 			vdata.normal  += ( normal  * joint.orient ) * weight.bias;
 			vdata.tangent += ( tangent * joint.orient ) * weight.bias;
 		}
	}

	return true;
}


md5_animation::md5_animation()
	: m_md5_version( 0 )
	, m_num_frames( 0 )
	, m_num_joints( 0 )
	, m_frame_rate( 0 )
	, m_num_animated_components( 0 )
	, m_anim_duration( 0 )
	, m_frame_duration( 0 )
	, m_anim_time( 0 )
{

}

md5_animation::~md5_animation()
{

}

bool md5_animation::load_animation( stream& source )
{
	m_joint_infos.clear();
	m_bounds.clear();
	m_base_frames.clear();
	m_frames.clear();
	m_animated_skeleton.joints.clear();
	m_num_frames = 0;

	std_stream sstream( &source );
	std::string command;

	sstream >> command;
	while ( !sstream.eof() )
	{
		if ( command == "MD5Version" )
		{
			sstream >> m_md5_version;
			assert( m_md5_version == 10 );
		}
		else if ( command == "commandline" )
		{
			next_line( sstream ); 
		}
		else if ( command == "numFrames" )
		{
			sstream >> m_num_frames;
			next_line( sstream ); 
		}
		else if ( command == "numJoints" )
		{
			sstream >> m_num_joints;
			next_line( sstream ); 
		}
		else if ( command == "frameRate" )
		{
			sstream >> m_frame_rate;
			next_line( sstream ); 
		}
		else if ( command == "numAnimatedComponents" )
		{
			sstream >> m_num_animated_components;
			next_line( sstream ); 
		}
		else if ( command == "hierarchy" )
		{
			discard( sstream, "{" );
			for ( size_t i = 0; i < m_num_joints; ++i )
			{
				md5_joint_info joint;
				sstream >> joint.name >> joint.parent_id >> joint.flags >> joint.start_index;
				remove_quotes( joint.name );
				m_joint_infos.push_back( joint );
				next_line( sstream ); 
			}
			discard( sstream, "}" );
		}
		else if ( command == "bounds" )
		{
			discard( sstream, "{" );
			next_line( sstream ); 
			for ( size_t i = 0; i < m_num_frames; ++i ) 
			{
				md5_bound bound;
				discard( sstream, "(" );
				sstream >> bound.min.x >> bound.min.y >> bound.min.z;
				discard( sstream, ")" );
				discard( sstream, "(" );
				sstream >> bound.max.x >> bound.max.y >> bound.max.z;

				m_bounds.push_back( bound );

				next_line( sstream ); 
			}

			discard( sstream, "}" );
			next_line( sstream ); 
		}
		else if ( command == "baseframe" )
		{
			discard( sstream, "{" );
			next_line( sstream ); 

			for ( size_t i = 0; i < m_num_joints; ++i )
			{
				md5_base_frame base_frame;
				discard( sstream, "(" );
				sstream >> base_frame.pos.x >> base_frame.pos.y >> base_frame.pos.z;
				discard( sstream, ")" );
				discard( sstream, "(" );
				sstream >> base_frame.orient.x >> base_frame.orient.y >> base_frame.orient.z;
				next_line( sstream ); 

				m_base_frames.push_back( base_frame );
			}
			discard( sstream, "}" );
			next_line( sstream ); 
		}
		else if ( command == "frame" )
		{
			md5_frame_data frame;
			sstream >> frame.frame_id;
			discard( sstream, "{" );
			next_line( sstream ); 

			frame.frame_data.reserve( m_num_animated_components );
			char buf[50];
			for ( size_t i = 0; i < m_num_animated_components; ++i )
			{
				sstream >> buf;
				frame.frame_data.push_back((float)atof(buf));
			}

			m_frames.push_back(frame);

			build_frame_skeleton( m_skeletons, m_joint_infos, m_base_frames, frame );

			discard( sstream, "}" );
			next_line( sstream ); 
		}

		sstream >> command;
	} 

	m_animated_skeleton.joints.assign( m_num_joints, md5_skeleton_joint() );

	m_frame_duration = 1.0f / (float)m_frame_rate;
	m_anim_duration = ( m_frame_duration * (float)m_num_frames );
	m_anim_time = 0.0f;

	assert( m_joint_infos.size() == m_num_joints );
	assert( m_bounds.size()      == m_num_frames );
	assert( m_base_frames.size() == m_num_joints );
	assert( m_frames.size()      == m_num_frames );
	assert( m_skeletons.size()   == m_num_frames );

	return true;
}

void md5_animation::update( float delta_time )
{
	if ( m_num_frames < 1 ) return;

	m_anim_time += delta_time;

	while ( m_anim_time > m_anim_duration ) m_anim_time -= m_anim_duration;
	while ( m_anim_time < 0.0f ) m_anim_time += m_anim_duration;

	float frame_num = m_anim_time * (float)m_frame_rate;
	size_t frame0 = (size_t)floorf( frame_num );
	size_t frame1 = (size_t)ceilf( frame_num );
	frame0 = frame0 % m_num_frames;
	frame1 = frame1 % m_num_frames;

	float interpolate = fmodf( m_anim_time, m_frame_duration ) / m_frame_duration;

	interpolate_skeletons( m_animated_skeleton, m_skeletons[frame0], m_skeletons[frame1], interpolate );
}

void md5_animation::build_frame_skeleton( md5_frame_skeleton_list& skeletons, const md5_joint_info_list& joint_infos, const md5_base_frame_list& base_frames, const md5_frame_data& frame_data )
{
	md5_frame_skeleton skeleton;

	for ( unsigned int i = 0; i < joint_infos.size(); ++i )
	{
		unsigned int j = 0;

		const md5_joint_info& jinfo = joint_infos[i];
		md5_skeleton_joint animated_joint = base_frames[i];

		animated_joint.parent = jinfo.parent_id;

		if ( jinfo.flags & 1 )  animated_joint.pos.x    = frame_data.frame_data[ jinfo.start_index + j++ ];
		if ( jinfo.flags & 2 )  animated_joint.pos.y    = frame_data.frame_data[ jinfo.start_index + j++ ];
		if ( jinfo.flags & 4 )  animated_joint.pos.z    = frame_data.frame_data[ jinfo.start_index + j++ ];
		if ( jinfo.flags & 8 )  animated_joint.orient.x = frame_data.frame_data[ jinfo.start_index + j++ ];
		if ( jinfo.flags & 16 ) animated_joint.orient.y = frame_data.frame_data[ jinfo.start_index + j++ ];
		if ( jinfo.flags & 32 )	animated_joint.orient.z = frame_data.frame_data[ jinfo.start_index + j++ ];

		unit_quat_w( animated_joint.orient );

		if ( animated_joint.parent >= 0 ) // Has a parent joint
		{
			md5_skeleton_joint& pjoint = skeleton.joints[static_cast< size_t >( animated_joint.parent ) ];
			glm::vec3 rot_pos = pjoint.orient * animated_joint.pos;

			animated_joint.pos    = pjoint.pos + rot_pos;
			animated_joint.orient = pjoint.orient * animated_joint.orient;

			animated_joint.orient = glm::normalize( animated_joint.orient );
		}

		skeleton.joints.push_back( animated_joint );
	}

	skeletons.push_back( skeleton );
}

void md5_animation::interpolate_skeletons( md5_frame_skeleton& final_skeleton, const md5_frame_skeleton& skeleton0, const md5_frame_skeleton& skeleton1, float interpolate )
{
	for ( size_t i = 0; i < m_num_joints; ++i )
	{
		md5_skeleton_joint& final_joint = final_skeleton.joints[i];
		const md5_skeleton_joint& joint0 = skeleton0.joints[i]; 
		const md5_skeleton_joint& joint1 = skeleton1.joints[i];

		final_joint.parent = joint0.parent;

		final_joint.orient = glm::slerp( joint0.orient, joint1.orient, interpolate );
		final_joint.pos    = glm::mix( joint0.pos, joint1.pos, interpolate );
	}
}

bool md5_loader::check_animation( const md5_animation& animation ) const
{
	if ( m_num_joints != animation.get_num_joints() )
	{
		return false;
	}

	for ( uint32 i = 0; i < m_joints.size(); ++i )
	{
		const md5_joint& mjoint = m_joints[i];
		const md5_animation::md5_joint_info& ajoint = animation.get_joint_info( i );

		if ( mjoint.name != ajoint.name || mjoint.parent_id != ajoint.parent_id )
		{
			return false;
		}
	}

	return true;
}

mesh_data* nv::md5_loader::release_mesh_data( uint32 mesh )
{
	mesh_data* result = m_meshes[ mesh ];
	m_meshes[ mesh ] = nullptr;
	return result;
}

void nv::md5_mesh_data::apply( const md5_animation& animation )
{
	const md5_animation::md5_frame_skeleton& skeleton = animation.get_skeleton();

	for ( unsigned int i = 0; i < m_vtx_data.size(); ++i )
	{
		const md5_vtx_data& vert = m_vtx_data[i];
		md5_vtx_pnt& result = m_pntdata[i];

		result.position = glm::vec3(0);
		result.normal   = glm::vec3(0);
		result.tangent  = glm::vec3(0);

		for ( size_t j = 0; j < vert.weight_count; ++j )
		{
			const md5_weight& weight = m_weights[vert.start_weight + j];
			const md5_animation::md5_skeleton_joint& joint = skeleton.joints[weight.joint_id];

			glm::vec3 rot_pos = joint.orient * weight.pos;
			result.position += ( joint.pos + rot_pos ) * weight.bias;

			result.normal  += ( joint.orient * vert.normal  ) * weight.bias;
			result.tangent += ( joint.orient * vert.tangent ) * weight.bias;
		}
	}
}

nv::md5_loader::~md5_loader()
{
	for ( auto m : m_meshes ) { if (m) delete m; }
}