// Copyright (C) 2012-2015 ChaosForge Ltd
// http://chaosforge.org/
//
// This file is part of Nova libraries. 
// For conditions of distribution and use, see copying.txt file in root folder.

#include "nv/gfx/mesh_creator.hh"

struct nv_key_transform { nv::transform tform; };

void nv::mesh_nodes_creator::pre_transform_keys()
{
	if ( m_data->m_flat ) return;
	merge_keys();
	uint32 max_frames = 0;
	for ( size_t i = 0; i < m_data->get_count(); ++i )
	{
		sint16 parent_id = m_data->m_nodes[i].parent_id;
		key_data* keys   = m_data->m_nodes[i].data;
		key_data* pkeys  = ( parent_id != -1 ? m_data->m_nodes[parent_id].data : nullptr );
		size_t count     = ( keys ? keys->get_channel(0)->element_count() : 0 );
		size_t pcount    = ( pkeys ? pkeys->get_channel(0)->element_count() : 0 );
		max_frames = nv::max<uint32>( count, max_frames );
		if ( pkeys && pkeys->get_channel_count() > 0 && keys && keys->get_channel_count() > 0 )
		{
			nv_key_transform*  channel = reinterpret_cast<nv_key_transform*>(keys->get_channel(0)->data);
			nv_key_transform* pchannel = reinterpret_cast<nv_key_transform*>(pkeys->get_channel(0)->data);
			for ( unsigned n = 0; n < count; ++n )
			{
				channel[n].tform = pchannel[ nv::min( n, pcount-1 ) ].tform * channel[n].tform;
			}
		}
	}

	// DAE pre_transform hack
	if ( m_data->m_frame_rate == 1 )
	{
		m_data->m_frame_rate = 32;
		m_data->m_duration   = static_cast<float>( max_frames );
	}

	m_data->m_flat = true;
}

// TODO: DELETE
struct assimp_key_p  { float time; nv::vec3 position; };
struct assimp_key_r  { float time; nv::quat rotation; };


void nv::mesh_nodes_creator::merge_keys()
{
	for ( size_t i = 0; i < m_data->get_count(); ++i )
	{
		key_data* old_keys = m_data->m_nodes[i].data;
		if ( old_keys && old_keys->get_channel_count() > 0 )
		{
			size_t chan_count = old_keys->get_channel_count();
			if ( chan_count == 1 
				&& old_keys->get_channel(0)->desc.slot_count() == 1 
				&& old_keys->get_channel(0)->desc[0].etype == TRANSFORM ) continue;

			size_t max_keys = 0;
			for ( size_t c = 0; c < chan_count; ++c )
			{
				max_keys = nv::max( max_keys, old_keys->get_channel(c)->count );
			}

			raw_data_channel* raw_channel = raw_data_channel::create<nv_key_transform>( max_keys );
			key_data* new_keys = new key_data;
			new_keys->add_channel( raw_channel );
			nv_key_transform* channel = reinterpret_cast<nv_key_transform*>(raw_channel->data);
			data_descriptor final_key = old_keys->get_final_key();

			for ( unsigned n = 0; n < max_keys; ++n )
			{
				float key[ 16 ];
				float* pkey = key;

				for ( uint16 c = 0; c < chan_count; ++c )
				{
					size_t idx = nv::min( old_keys->get_channel(c)->count - 1, n );
					pkey += old_keys->get_raw( old_keys->get_channel(c), idx, pkey );
				}
				channel[n].tform = extract_transform_raw( final_key, key );
			}

			delete old_keys;
			m_data->m_nodes[i].data = new_keys;
		}
	}
}

void nv::mesh_nodes_creator::transform( float scale, const mat3& r33 )
{
	mat3 ri33 = glm::inverse( r33 );
	mat4 pre_transform ( scale * r33 );
	mat4 post_transform( 1.f/scale * ri33 ); 

	for ( size_t i = 0; i < m_data->get_count(); ++i )
	{
		mesh_node_data& node = m_data->m_nodes[i];
		node.transform = pre_transform * node.transform * post_transform;
		if ( node.data )
		{
			key_data* kdata  = node.data;
			for ( size_t c = 0; c < kdata->get_channel_count(); ++c )
			{
				const raw_data_channel* channel = kdata->get_channel(c);
				size_t key_size = channel->desc.element_size();
				for ( size_t n = 0; n < channel->count; ++n )
				{
					transform_key_raw( channel->desc, channel->data + n * key_size, scale, r33, ri33 );
				}
			}
		}
	}
}

void nv::mesh_data_creator::transform( float scale, const mat3& r33 )
{
	vec3 vertex_offset     = vec3(); 
	mat3 vertex_transform  = scale * r33;
	mat3 normal_transform  = r33;

	for ( uint32 c = 0; c < m_data->get_channel_count(); ++c )
	{
		const raw_data_channel* channel = m_data->get_channel(c);
		const data_descriptor&  desc    = channel->desc;
		uint8* raw_data = channel->data;
		uint32 vtx_size = desc.element_size();
		int p_offset = -1;
		int n_offset = -1;
		int t_offset = -1;
		for ( const auto& cslot : desc  )
			switch ( cslot.vslot )
			{
				case slot::POSITION : if ( cslot.etype == FLOAT_VECTOR_3 ) p_offset = int( cslot.offset ); break;
				case slot::NORMAL   : if ( cslot.etype == FLOAT_VECTOR_3 ) n_offset = int( cslot.offset ); break;
				case slot::TANGENT  : if ( cslot.etype == FLOAT_VECTOR_4 ) t_offset = int( cslot.offset ); break;
				default             : break;
			}

		if ( p_offset != -1 )
			for ( uint32 i = 0; i < channel->count; i++)
			{
				vec3& p = *reinterpret_cast<vec3*>( raw_data + vtx_size*i + p_offset );
				p = vertex_transform * p + vertex_offset;
			}

		if ( n_offset != -1 )
			for ( uint32 i = 0; i < channel->count; i++)
			{
				vec3& n = *reinterpret_cast<vec3*>( raw_data + vtx_size*i + n_offset );
				n = glm::normalize( normal_transform * n );
			}
		if ( t_offset != -1 )
			for ( uint32 i = 0; i < channel->count; i++)
			{
				vec4& t = *reinterpret_cast<vec4*>(raw_data + vtx_size*i + t_offset );
				t = vec4( glm::normalize( normal_transform * vec3(t) ), t[3] );
			}
	}
}

struct vertex_g
{
	nv::vec4 tangent;
};

void nv::mesh_data_creator::flip_normals()
{
	int ch_n  = m_data->get_channel_index( slot::NORMAL );
	size_t n_offset = 0;
	if ( ch_n == -1 ) return;
	raw_data_channel* channel = m_data->m_channels[ unsigned( ch_n ) ];
	for ( const auto& cslot : channel->desc )
		if ( cslot.vslot == slot::NORMAL )
		{
			n_offset  = cslot.offset;
		}

	for ( uint32 i = 0; i < channel->count; ++i )
	{
		vec3& normal = *reinterpret_cast<vec3*>( channel->data + channel->desc.element_size() * i + n_offset );
		normal = -normal;
	}
}


void nv::mesh_data_creator::generate_tangents()
{
	int p_offset = -1;
	int n_offset = -1;
	int t_offset = -1;
	datatype i_type = NONE;
	uint32 n_channel_index = 0;

	const raw_data_channel* p_channel = nullptr;
	      raw_data_channel* n_channel = nullptr;
	const raw_data_channel* t_channel = nullptr;
	const raw_data_channel* i_channel = nullptr;

	for ( uint32 c = 0; c < m_data->get_channel_count(); ++c )
	{
		const raw_data_channel* channel = m_data->get_channel(c);

		for ( const auto& cslot : channel->desc )
		switch ( cslot.vslot )
		{
			case slot::POSITION : 
				if ( cslot.etype == FLOAT_VECTOR_3 )
				{
					p_offset  = int( cslot.offset );
					p_channel = channel;
				}
				break;
			case slot::NORMAL   : 
				if ( cslot.etype == FLOAT_VECTOR_3 )
				{
					n_offset  = int( cslot.offset );
					n_channel = m_data->m_channels[ c ];
					n_channel_index = c;
				}
				break;
			case slot::TEXCOORD : 
				if ( cslot.etype == FLOAT_VECTOR_2 )
				{
					t_offset  = int( cslot.offset );
					t_channel = channel;
				}
				break;
			case slot::INDEX    : 
				{
					i_type    = cslot.etype;
					i_channel = channel;
				}
				break;
			case slot::TANGENT  : return;
			default             : break;
		}
	}
	if ( !p_channel || !n_channel || !t_channel ) return;

	if ( p_channel->count != n_channel->count || p_channel->count % t_channel->count != 0 || ( i_type != UINT && i_type != USHORT && i_type != NONE ) )
	{
		return;
	}

	raw_data_channel* g_channel = raw_data_channel::create<vertex_g>( p_channel->count );
	vec4* tangents              = reinterpret_cast<vec4*>( g_channel->data );
	vec3* tangents2             = new vec3[ p_channel->count ];
	uint32 tri_count = i_channel ? i_channel->count / 3 : t_channel->count / 3;
	uint32 vtx_count = p_channel->count;
	uint32 sets      = p_channel->count / t_channel->count;

	for ( unsigned int i = 0; i < tri_count; ++i )
	{
		uint32 ti0 = 0;
		uint32 ti1 = 0;
		uint32 ti2 = 0;
		if ( i_type == UINT )
		{
			const uint32* idata = reinterpret_cast<const uint32*>( i_channel->data );
			ti0 = idata[ i * 3 ];
			ti1 = idata[ i * 3 + 1 ];
			ti2 = idata[ i * 3 + 2 ];
		}
		else if ( i_type == USHORT )
		{
			const uint16* idata = reinterpret_cast<const uint16*>( i_channel->data );
			ti0 = idata[ i * 3 ];
			ti1 = idata[ i * 3 + 1 ];
			ti2 = idata[ i * 3 + 2 ];
		}
		else // if ( i_type == NONE )
		{
			ti0 = i * 3;
			ti1 = i * 3 + 1;
			ti2 = i * 3 + 2;
		}

		const vec2& w1 = *reinterpret_cast<vec2*>(t_channel->data + t_channel->desc.element_size()*ti0 + t_offset );
		const vec2& w2 = *reinterpret_cast<vec2*>(t_channel->data + t_channel->desc.element_size()*ti1 + t_offset );
		const vec2& w3 = *reinterpret_cast<vec2*>(t_channel->data + t_channel->desc.element_size()*ti2 + t_offset );
		vec2 st1 = w3 - w1;
		vec2 st2 = w2 - w1;
		float stst = (st1.x * st2.y - st2.x * st1.y);
		float coef = ( stst != 0.0f ? 1.0f / stst : 0.0f );

		for ( uint32 set = 0; set < sets; ++set )
		{
			uint32 nti0 = t_channel->count * set + ti0;
			uint32 nti1 = t_channel->count * set + ti1;
			uint32 nti2 = t_channel->count * set + ti2;
			vec3 v1 = *reinterpret_cast<vec3*>(p_channel->data + p_channel->desc.element_size()*nti0 + p_offset );
			vec3 v2 = *reinterpret_cast<vec3*>(p_channel->data + p_channel->desc.element_size()*nti1 + p_offset );
			vec3 v3 = *reinterpret_cast<vec3*>(p_channel->data + p_channel->desc.element_size()*nti2 + p_offset );
			vec3 xyz1 = v3 - v1;
			vec3 xyz2 = v2 - v1;

			//vec3 normal = glm::cross( xyz1, xyz2 );
			//
			//vtcs[ ti0 ].normal += normal;
			//vtcs[ ti1 ].normal += normal;
			//vtcs[ ti2 ].normal += normal;
			vec3 tangent  = (( xyz1 * st2.y ) - ( xyz2 * st1.y )) * coef;
			vec3 tangent2 = (( xyz2 * st1.x ) - ( xyz1 * st2.x )) * coef;

			tangents[nti0] = vec4( vec3( tangents[nti0] ) + tangent, 0 );
			tangents[nti1] = vec4( vec3( tangents[nti1] ) + tangent, 0 );
			tangents[nti2] = vec4( vec3( tangents[nti2] ) + tangent, 0 );

			tangents2[nti0] += tangent2;
			tangents2[nti1] += tangent2;
			tangents2[nti2] += tangent2;
		}
	}

	for ( unsigned int i = 0; i < vtx_count; ++i )
	{
		const vec3 n = *reinterpret_cast<vec3*>( n_channel->data + n_channel->desc.element_size()*i + n_offset );
		const vec3 t = vec3(tangents[i]);
		if ( ! ( t.x == 0.0f && t.y == 0.0f && t.z == 0.0f ) )
		{
			tangents[i]    = vec4( glm::normalize(t - n * glm::dot( n, t )), 0.0f ); 
			tangents[i][3] = (glm::dot(glm::cross(n, t), tangents2[i]) < 0.0f) ? -1.0f : 1.0f;
		}
	}
	delete tangents2;

	m_data->m_channels[ n_channel_index ] = merge_channels( n_channel, g_channel );
	delete n_channel;
	delete g_channel;
}

nv::raw_data_channel* nv::mesh_data_creator::merge_channels( raw_data_channel* a, raw_data_channel* b )
{
	NV_ASSERT( a->count == b->count, "merge_channel - bad channels!" );
	data_descriptor adesc = a->desc;
	data_descriptor bdesc = b->desc;
	uint32          count = a->count;

	data_descriptor desc  = a->desc;
	for ( auto bslot : bdesc )
	{
		desc.push_slot( bslot.etype, bslot.vslot );
	}
	uint8* data = new uint8[ count * desc.element_size() ];
	for ( uint32 i = 0; i < count; ++i )
	{
		raw_copy_n( a->data + i * adesc.element_size(), adesc.element_size(), data + i*desc.element_size() );
		raw_copy_n( b->data + i * bdesc.element_size(), bdesc.element_size(), data + i*desc.element_size() + adesc.element_size() );
	}
	raw_data_channel* result = new raw_data_channel;
	result->count = count;
	result->desc  = desc;
	result->data  = data;
	return result;
}

nv::raw_data_channel* nv::mesh_data_creator::append_channels( raw_data_channel* a, raw_data_channel* b, uint32 frame_count )
{
	if ( a->desc != b->desc ) return nullptr;
	if ( a->count % frame_count != 0 ) return nullptr;
	if ( b->count % frame_count != 0 ) return nullptr;
	size_t vtx_size = a->desc.element_size();

	uint8* data = new uint8[ ( a->count + b->count ) * vtx_size ];
	

	if ( frame_count == 1 )
	{
		size_t a_size = vtx_size * a->count;
		raw_copy_n( a->data, a_size, data );
		raw_copy_n( b->data, vtx_size * b->count, data + a_size );
	}
	else
	{
		size_t frame_size_a = ( a->count / frame_count ) * vtx_size;
		size_t frame_size_b = ( b->count / frame_count ) * vtx_size;
		size_t pos_a = 0;
		size_t pos_b = 0;
		size_t pos   = 0;
		for ( size_t i = 0; i < frame_count; ++i )
		{
			raw_copy_n( a->data + pos_a, frame_size_a, data + pos );
			raw_copy_n( b->data + pos_b, frame_size_b, data + pos + frame_size_a );				pos_a += frame_size_a;
			pos_b += frame_size_b; 
			pos   += frame_size_a + frame_size_b;
		}
	}

	raw_data_channel* result = new raw_data_channel;
	result->count = a->element_count() + b->element_count();
	result->desc  = a->descriptor();
	result->data  = data;
	return result;
}



bool nv::mesh_data_creator::is_same_format( mesh_data* other )
{
	if ( m_data->get_channel_count() != other->get_channel_count() ) return false;
	for ( uint32 c = 0; c < m_data->get_channel_count(); ++c )
	{
		if ( m_data->get_channel(c)->descriptor() != other->get_channel(c)->descriptor() )
			return false;
	}
	return true;
}

void nv::mesh_data_creator::merge( mesh_data* other )
{
	if ( !is_same_format( other ) ) return;
	int ch_pi  = m_data->get_channel_index( slot::POSITION );
	int ch_ti  = m_data->get_channel_index( slot::TEXCOORD );
	int och_pi = other->get_channel_index( slot::POSITION );
	int och_ti = other->get_channel_index( slot::TEXCOORD );
	if ( ch_pi == -1 || ch_ti == -1 ) return;
	size_t size   = m_data->m_channels[ unsigned(ch_ti) ]->element_count();
	size_t osize  =  other->m_channels[ unsigned(och_ti) ]->element_count();
	size_t count  = m_data->m_channels[ unsigned(ch_pi) ]->element_count();
	size_t ocount =  other->m_channels[ unsigned(och_pi) ]->element_count();
	if ( count % size != 0 || ocount % osize != 0 ) return;
	if ( count / size != ocount / osize ) return;
	
	for ( uint32 c = 0; c < m_data->get_channel_count(); ++c )
	{
		raw_data_channel* old = m_data->m_channels[c];
		size_t frame_count = ( old->get_buffer_type() == INDEX_BUFFER ? 1 : old->element_count() / size );
		m_data->m_channels[c] = append_channels( old, other->m_channels[c], frame_count );
		NV_ASSERT( m_data->m_channels[c], "Merge problem!" );
		if ( old->get_buffer_type() == INDEX_BUFFER )
		{
			switch ( old->desc[0].etype )
			{
			case USHORT : 
				{
					NV_ASSERT( size + osize < uint16(-1), "Index out of range!" );
					uint16* indexes = reinterpret_cast<uint16*>( m_data->m_channels[c]->data );
					for ( uint16 i = uint16( old->element_count() ); i < m_data->m_channels[c]->element_count(); ++i )
						indexes[i] += uint16( size );

				}
				break;
			case UINT   : 
				{
					uint32* indexes = reinterpret_cast<uint32*>( m_data->m_channels[c]->data );
					for ( uint32 i = old->element_count(); i < m_data->m_channels[c]->element_count(); ++i )
						indexes[i] += size;
				}
				break;
			default : NV_ASSERT( false, "Unsupported index type!" ); break;
			}
			m_data->m_index_channel = m_data->m_channels[c];
		}
		delete old;
	}
}