#include "externals.hh"
#include "uniforms.hh"


/*= Uniform types ============================================================*/

E_UniformCategory GetUniformCategory(
	__rd__ const E_UniformType type )
{
	switch ( type ) {
		case E_UniformType::F1:
		case E_UniformType::F2:
		case E_UniformType::F3:
		case E_UniformType::F4:
			return E_UniformCategory::FLOAT;
		case E_UniformType::I1:
		case E_UniformType::I2:
		case E_UniformType::I3:
		case E_UniformType::I4:
			return E_UniformCategory::INT;
		case E_UniformType::SAMPLER2D:
			return E_UniformCategory::SAMPLER2D;
	}
	std::abort( );
}


uint32_t GetElementCount(
	__rd__ const E_UniformType type )
{
	switch ( type ) {
		case E_UniformType::F1:
		case E_UniformType::I1:
			return 1;
		case E_UniformType::F2:
		case E_UniformType::I2:
			return 2;
		case E_UniformType::F3:
		case E_UniformType::I3:
			return 3;
		case E_UniformType::F4:
		case E_UniformType::I4:
			return 4;

		case E_UniformType::SAMPLER2D:
			return 0;
	}
	std::abort( );
}

bool IsVectorCategory(
	__rd__ const E_UniformCategory category )
{
	assert( category != E_UniformCategory::__COUNT__ );
	return category == E_UniformCategory::FLOAT
		|| category == E_UniformCategory::INT;
}


/*= PackUniforms =============================================================*/

namespace {

T_PackedUniforms PackVectorUniforms_(
		__rd__ T_UniformDeclarations const& declarations ,
		__rd__ const E_UniformCategory category )
{
	std::vector< std::string > packOrder;
	packOrder.reserve( declarations.size( ) );
	for ( auto entry : declarations ) {
		auto const& u( entry.second );
		if ( GetUniformCategory( u.type ) == category ) {
			packOrder.push_back( entry.first );
		}
	}
	std::sort( packOrder.begin( ) , packOrder.end( ) ,
		[&declarations]( auto const& a , auto const& b )
		{
			auto const& ua( declarations.find( a )->second );
			auto const& ub( declarations.find( b )->second );
			return GetElementCount( ub.type ) < GetElementCount( ua.type );
		} );

	const auto n( packOrder.size( ) );
	auto index = 0u;
	T_PackedUniforms pu;

	// Place vec4's
	while ( index < n && GetElementCount( declarations.find( packOrder[ index ] )->second.type ) == 4 ) {
		pu.positions[ packOrder[ index ] ] = pu.elements;
		pu.elements += 4;
		index ++;
	}

	// Place vec3's
	uint32_t s3left( 0 );
	while ( index < n && GetElementCount( declarations.find( packOrder[ index ] )->second.type ) == 3 ) {
		pu.positions[ packOrder[ index ] ] = pu.elements;
		pu.elements += 4;
		s3left ++;
		index ++;
	}

	// Place vec2's
	const uint32_t s2s( pu.elements );
	bool s2spare( false );
	while ( index < n && GetElementCount( declarations.find( packOrder[ index ] )->second.type ) == 2 ) {
		if ( !s2spare ) {
			pu.elements += 4;
		}
		const uint32_t pos( pu.elements - ( s2spare ? 2 : 4 ) );
		pu.positions[ packOrder[ index ] ] = pos;
		s2spare = !s2spare;
		index ++;
	}

	// Place single elements
	uint32_t s1pos( s2spare ? ( pu.elements - 2 ) : pu.elements );
	while ( index < n ) {
		uint32_t pos;
		if ( s3left ) {
			pos = s2s - s3left * 4 + 3;
			s3left --;
		} else {
			if ( s1pos == pu.elements ) {
				pu.elements += 4;
			}
			pos = s1pos;
			s1pos ++;
		}
		pu.positions[ packOrder[ index ] ] = pos;
		index ++;
	}

	return pu;
}

T_PackedUniforms PackOpaqueUniforms_(
		__rd__ T_UniformDeclarations const& declarations ,
		__rd__ const E_UniformCategory category )
{
	T_PackedUniforms pu;
	for ( auto entry : declarations ) {
		auto const& u( entry.second );
		if ( GetUniformCategory( u.type ) == category ) {
			pu.positions[ entry.first ] = pu.elements;
			pu.elements ++;
		}
	}
	return pu;
}

} // namespace

/*----------------------------------------------------------------------------*/

T_PackedUniforms PackUniforms(
		__rd__ T_UniformDeclarations const& declarations ,
		__rd__ const E_UniformCategory category )
{
	assert( category != E_UniformCategory::__COUNT__ );
	if ( IsVectorCategory( category ) ) {
		return PackVectorUniforms_( declarations , category );
	}
	return PackOpaqueUniforms_( declarations , category );
}


/*= T_UniformSet =============================================================*/

constexpr uint32_t T_UniformSet::Invalid;

T_UniformSet::T_UniformSet( )
	: vectorElements_( 0 ) , opaqueElements_( 0 )
{ }

T_UniformSet::T_UniformSet(
		__rw__ T_UniformDeclarations declarations )
	: declarations_( std::move( declarations ) ) ,
		vectorElements_( 0 ) , opaqueElements_( 0 )
{
	for ( auto i = 0 ; i < int( E_UniformCategory::__COUNT__ ) ; i ++ ) {
		const auto cat{ E_UniformCategory( i ) };
		packed_[ i ] = PackUniforms( declarations_ , cat );
		if ( IsVectorCategory( cat ) ) {
			vectorElements_ += packed_[ i ].elements;
		} else {
			opaqueElements_ += packed_[ i ].elements;
		}
	}
}

uint32_t T_UniformSet::position(
	__rd__ std::string const& uniform ) const
{
	const auto pos( declarations_.find( uniform ) );
	if ( pos == declarations_.end( ) ) {
		return Invalid;
	}

	const int ci( int( GetUniformCategory( pos->second.type ) ) );
	const auto ePos( packed_[ ci ].positions.find( uniform ) );
	assert( ePos != packed_[ ci ].positions.end( ) );
	return ePos->second;
}