// Copyright (C) 2011-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/skeletal_mesh.hh"

#include "nv/interface/context.hh"
#include "nv/interface/device.hh"
#include "nv/stl/unordered_map.hh"

void nv::skeletal_animation_entry::initialize()
{
	m_root      = uint32(-1);
	m_prepared  = false;
	m_children  = nullptr;
	m_offsets   = nullptr;
	m_bone_ids  = new sint16[m_node_data->size()];

	NV_ASSERT( m_node_data, "node data empty!" );

	if ( !m_node_data->is_flat() )
	{
		m_children = new vector< uint32 >[m_node_data->size()];
		for ( uint32 n = 0; n < m_node_data->size(); ++n )
		{
			const data_channel_set* node = (*m_node_data)[n];
			if ( node->get_parent_id() != -1 )
			{
				m_children[ node->get_parent_id()].push_back( n );
			}
			else
			{
				if ( m_root >= 0 )
				{
					m_root = uint32( -2 );
				}
				else
					m_root = n;
			}
		}
		NV_ASSERT( m_root != uint32( -1 ), "Animation without root!" );
	}
}

void nv::skeletal_animation_entry::update_skeleton( mat4* data, uint32 a_ms_time ) const
{
	float  fframe   = ( a_ms_time * 0.001f ) * m_fps;
	uint32 frame    = math::floor( fframe );
	float  reminder = fframe - static_cast<float>( frame );
	uint32 duration = get_frame_count();
	if ( duration == 0 )
	{
		frame  = get_start_frame();
		fframe = static_cast<float>( frame );
	}
	else if ( frame >= duration )
	{
		if ( is_looping() )
		{
			frame  = frame % duration;
			fframe = static_cast<float>( frame ) + reminder;
		}
		else
		{
			frame  = get_end_frame();
			fframe = static_cast<float>( frame );
		}
	}

	if ( !m_node_data->is_flat() )
	{
		if ( m_root == uint32( -2 ) ) // multi-root
		{
			for ( uint32 n = 0; n < m_node_data->size(); ++n )
				if ( ( *m_node_data )[n]->get_parent_id() == -1 )
					animate_rec( data, fframe, n, mat4() );
		}
		else
			animate_rec( data, fframe, m_root, mat4() );
		return;
	}

	for ( uint32 n = 0; n < m_node_data->size(); ++n )
		if ( m_bone_ids[n] >= 0 )
		{
			const data_channel_set* node = (*m_node_data)[n];
			nv::mat4 node_mat( node->get_transform() );

			if ( node->size() > 0 )
			{
				raw_channel_interpolator interpolator( node, m_interpolation_key );
				node_mat = interpolator.get< mat4 >( fframe );
			}

			sint16 bone_id = m_bone_ids[n];
			data[ bone_id ] = node_mat * m_offsets[ bone_id ];
		}
}

void nv::skeletal_animation_entry::prepare( const mesh_nodes_data* bones )
{
	if ( m_prepared ) return;
	nv::hash_store< shash64, uint16 > bone_names;
	m_offsets = new mat4[ bones->size() ];
	for ( nv::uint16 bi = 0; bi < bones->size(); ++bi )
	{
		const data_channel_set* bone = (*bones)[ bi ];
		bone_names[ bone->get_name() ] = bi;
		m_offsets[bi] = bone->get_transform();
	}

	for ( uint32 n = 0; n < m_node_data->size(); ++n )
	{
		const data_channel_set* node = (*m_node_data)[ n ];
		sint16 bone_id = -1;

		auto bi = bone_names.find( node->get_name() );
		if ( bi != bone_names.end() )
		{
			bone_id = sint16( bi->second );
		}
		m_bone_ids[n] = bone_id;

		if ( m_interpolation_key.size() == 0 && node->size() > 0 )
			m_interpolation_key = node->get_interpolation_key();

	}
	m_prepared = true;
}

void nv::skeletal_animation_entry::animate_rec( mat4* data, float time, uint32 node_id, const mat4& parent_mat ) const
{
	// TODO: fix transforms, which are now embedded,
	//       see note in assimp_loader.cc:load_node
	const data_channel_set* node = ( *m_node_data )[ node_id ];
	mat4 node_mat( node->get_transform() );

	if ( node->size() > 0 )
	{
		raw_channel_interpolator interpolator( node, m_interpolation_key );
		node_mat = interpolator.get< mat4 >( time );
	}

	mat4 global_mat = parent_mat * node_mat;

	sint16 bone_id = m_bone_ids[ node_id ];
	if ( bone_id >= 0 )
	{
		data[ bone_id ] = global_mat * m_offsets[ bone_id ];
	}

	for ( auto child : m_children[ node_id ] )
	{
		animate_rec( data, time, child, global_mat );
	}
}

nv::skeletal_animation_entry::~skeletal_animation_entry()
{
	delete[] m_offsets;
	delete[] m_children;
	delete[] m_bone_ids;
}

nv::skeletal_mesh::skeletal_mesh( context* a_context, const data_channel_set* a_mesh, const mesh_nodes_data* a_bone_data )
	: m_context( a_context ), m_bone_data( a_bone_data ), m_index_count( 0 ), m_transform( nullptr ), m_parent_id(-1)
{
	if ( a_mesh )
	{
		m_va = a_context->create_vertex_array( a_mesh, nv::STATIC_DRAW );
		m_index_count = a_mesh->get_channel_size( slot::INDEX );
		m_parent_id = a_mesh->get_parent_id();
	}
	if ( m_bone_data )
	{
		m_transform = new mat4[ m_bone_data->size() ];
	}
}

void nv::skeletal_mesh::update_animation( animation_entry* a_anim, uint32 a_anim_time )
{
	if ( m_bone_data && a_anim )
	{
		skeletal_animation_entry * anim = static_cast<skeletal_animation_entry*>( a_anim );
		anim->prepare( m_bone_data );
		anim->update_skeleton( m_transform, a_anim_time );
	}
}

void nv::skeletal_mesh::update( program a_program )
{
	if ( m_bone_data )
		m_context->get_device()->set_opt_uniform_array( a_program, "nv_m_bones", m_transform, m_bone_data->size() );
}

nv::transform nv::skeletal_mesh::get_node_transform( uint32 node_id ) const
{
	if ( node_id == 0 ) return transform();
	if ( node_id == uint32(-1) ) return transform( m_transform[0] );
	return transform( m_transform[ node_id ] );
}

nv::mat4 nv::skeletal_mesh::get_node_matrix( uint32 node_id ) const
{
	return m_transform[ node_id ];
}