source: trunk/src/formats/assimp_loader.cc @ 451

// Copyright (C) 2014-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/formats/assimp_loader.hh"

#include "nv/interface/data_channel_access.hh"
#include "nv/stl/hash_store.hh"
#include "nv/lib/assimp.hh"

using namespace nv;

const unsigned MAX_BONES = 64;

struct assimp_plain_vtx 
{
        vec3 position;
        vec3 normal;
        vec2 texcoord;
        vec4 tangent;

        assimp_plain_vtx() {}
        assimp_plain_vtx( const vec3& v, const vec2& t, const vec3& n, const vec4& g )
        {
                position = v;
                texcoord = t;
                normal   = n;
                tangent  = g;
        }
};

struct assimp_skinned_vtx 
{
        vec3  position;
        vec3  normal;
        vec2  texcoord;
        vec4  tangent;
        ivec4 boneindex;
        vec4  boneweight;

        assimp_skinned_vtx() {}
        assimp_skinned_vtx( const vec3& v, const vec2& t, const vec3& n, const vec4& g )
        {
                position = v;
                texcoord = t;
                normal   = n;
                tangent  = g;
        }
};

struct assimp_key_p  { float time; vec3 translation; };
struct assimp_key_r  { float time; quat rotation; };
struct assimp_key_s  { float time; vec3 scale; };
struct assimp_key_tr { transform tform; };


nv::assimp_loader::assimp_loader( string_table* strings, const string_view& a_ext, uint32 a_assimp_flags /*= 0 */ )
        : mesh_loader( strings ), m_scene( nullptr ), m_mesh_count(0)
{
        m_ext   = a_ext;
        m_assimp_flags = a_assimp_flags;
        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 )
{
        load_assimp_library();
        if ( m_scene != nullptr ) aiReleaseImport( reinterpret_cast<const aiScene*>( m_scene ) );
        m_scene = nullptr;
        m_mesh_count = 0;
        NV_LOG_NOTICE( "AssImp loading file..." );
        size_t size = source.size();
        char* data  = new char[ size ];
        source.read( data, size, 1 );
        const aiScene* scene = aiImportFileFromMemory( data, size, m_assimp_flags, m_ext.data() );

        if( !scene)
        {
                NV_LOG_ERROR( aiGetErrorString() );
                return false;
        }
        m_scene      = scene;
        m_mesh_count = scene->mNumMeshes;
        NV_LOG_NOTICE( "Loading successfull" );
        return true;
}

data_channel_set* nv::assimp_loader::release_mesh_data( size_t index /*= 0 */ )
{
        if ( index >= m_mesh_count ) return nullptr;
        data_channel_set* result = data_channel_set_creator::create_set( 2 );
        load_mesh_data( result, index );
        return result;
}
void nv::assimp_loader::load_mesh_data( data_channel_set* data, size_t index )
{
        const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
        const aiMesh*  mesh  = scene->mMeshes[ index ];

        bool skinned = mesh->mNumBones > 0;

        data_descriptor desc;
        if ( skinned )
                desc.initialize< assimp_skinned_vtx >();
        else
                desc.initialize< assimp_plain_vtx >();
        data_channel_set_creator maccess( data );
        const char* name = mesh->mName.data;
        maccess.set_name( make_name( name ) );
        uint8*  cdata   = maccess.add_channel( desc, mesh->mNumVertices ).raw_data();
        uint16* indices = reinterpret_cast<uint16*>( maccess.add_channel< index_u16 >( mesh->mNumFaces * 3 ).raw_data() );

        if ( mesh->mTangents && mesh->mBitangents )
                for ( unsigned int i = 0; i < mesh->mNumVertices; i++ )
                {
                        vec3 v = assimp_vec3_cast( mesh->mVertices[i] );
                        vec3 n = glm::normalize( assimp_vec3_cast( mesh->mNormals[i] ) );
                        vec3 t = glm::normalize( assimp_vec3_cast( mesh->mTangents[i] ) );
                        vec3 b = glm::normalize( assimp_vec3_cast( mesh->mBitangents[i] ) );
                        vec2 s = assimp_st_cast( mesh->mTextureCoords[0][i] );

                        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 )
                                reinterpret_cast<assimp_skinned_vtx*>( cdata )[i] = assimp_skinned_vtx( v, s, n, vt );
                        else
                                reinterpret_cast<assimp_plain_vtx*>( cdata )[i] = assimp_plain_vtx( v, s, n, vt );
                }

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

        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] );
                }
        }

}

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

bool nv::assimp_loader::load_bones( size_t index, array_ref< data_channel_set* > bones )
{
        if ( m_scene == nullptr ) return false;
        const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
        const aiMesh*  mesh  = scene->mMeshes[ index ];

        for (unsigned int m=0; m<mesh->mNumBones; m++)
        {
                aiBone* bone   = mesh->mBones[m];
                mat4    offset = assimp_mat4_cast( bone->mOffsetMatrix );
                bones[m] = data_channel_set_creator::create_set( 0 );
                data_channel_set_creator access( bones[m] );
                const char* name = bone->mName.data;
                access.set_name( make_name( name ) );
                access.set_transform( offset );
        }
        return true;
}

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

        NV_LOG_NOTICE( "------------------------" );
        NV_LOG_NOTICE( "Texture   count - ", scene->mNumTextures );
        NV_LOG_NOTICE( "Animation count - ", scene->mNumAnimations );
        NV_LOG_NOTICE( "Material  count - ", scene->mNumMaterials );
        NV_LOG_NOTICE( "Meshes    count - ", scene->mNumMeshes );
        NV_LOG_NOTICE( "------------------------" );

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

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

                        NV_LOG_NOTICE( "Mesh #", mc, "   - ", string_view( static_cast<char*>( mesh->mName.data ) ) );
                        NV_LOG_NOTICE( "  bones   - ", mesh->mNumBones );
                        NV_LOG_NOTICE( "  uvs     - ", mesh->mNumUVComponents[0] );
                        NV_LOG_NOTICE( "  verts   - ", mesh->mNumVertices );
                        NV_LOG_NOTICE( "  faces   - ", mesh->mNumFaces );

                        //                      NV_LOG_NOTICE(  "Bones:" );
                        //                      for (unsigned int m=0; m<mesh->mNumBones; m++)
                        //                      {
                        //                              aiBone* bone  = mesh->mBones[m];
                        //                              NV_LOG_NOTICE( bone->mName.C_Str() );
                        //                      }
                }
        }
        else
        {
                NV_LOG_NOTICE( "No meshes!" );
        }
        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_NOTICE( "  material - ", path.data );
        //                              texIndex++;
        //                      }
        //              }
        //      }
        //      else
        //      {
        //              NV_LOG_NOTICE( "No materials" );
        //      }
        //      NV_LOG_NOTICE( "------------------------" );

}

mesh_nodes_data* nv::assimp_loader::release_merged_bones( data_channel_set* meshes )
{
        const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
        mesh_nodes_data* result = new mesh_nodes_data( make_name( "bones" ) );
        hash_store< shash64, uint16 > names;
        for ( unsigned int m = 0; m < m_mesh_count; ++m )
        {
                uint16 translate[MAX_BONES];
                vector< data_channel_set* > bones;
                const aiMesh*  mesh  = scene->mMeshes[ m ];
                if ( mesh->mNumBones != 0 )
                {
                        bones.resize( mesh->mNumBones );
                        load_bones( m, bones );
                        for ( unsigned int b = 0; b < mesh->mNumBones; ++b )
                        {

                                data_channel_set* bone = bones[b];
                                auto iname = names.find( bone->get_name() );
                                if ( iname == names.end() )
                                {
                                        NV_ASSERT( result->size() < MAX_BONES, "Too many bones to merge!" );
                                        uint16 index = uint16( result->size() );
                                        result->push_back( bone );
                                        names[ bone->get_name() ] = index;
                                        translate[b] = index;
                                }
                                else
                                {
                                        translate[b] = iname->second;
                                }
                        }
                        if ( m > 0 && bones.size() > 0 )
                        {
                                data_channel_access< assimp_skinned_vtx > channel( const_cast< raw_data_channel* >( meshes[m].get_channel(0) ) );
                                for ( unsigned v = 0; v < channel.size(); ++v )
                                {
                                        assimp_skinned_vtx& vertex = channel.data()[v];

                                        for ( int i = 0 ; i < 4; ++i)
                                        {
                                                if ( vertex.boneweight[i] > 0.0f ) 
                                                {
                                                        vertex.boneindex[i] = int( translate[vertex.boneindex[i]] );
                                                }
                                        }
                                }
                        }
                }       
        }

        return result;
}

mesh_nodes_data* nv::assimp_loader::release_mesh_nodes_data( size_t index /*= 0*/ )
{
        if ( m_scene == nullptr ) return nullptr;
        const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
        if ( scene->mRootNode == nullptr || scene->mAnimations == nullptr || scene->mAnimations[index] == nullptr) return nullptr;

        const aiNode*      root = scene->mRootNode;
        const aiAnimation* anim = scene->mAnimations[index];

        uint32 count = count_nodes( scene->mRootNode );

        uint16 frame_rate     = static_cast<uint16>( anim->mTicksPerSecond );
        float  duration       = static_cast<float>( anim->mDuration );
        bool   flat           = false;

        data_channel_set** temp = new data_channel_set*[ count ];
        array_ref< data_channel_set* > temp_ref( temp, count );
        load_node( index, temp_ref, root, 0, -1 );

        mesh_nodes_data* result = new mesh_nodes_data( make_name( static_cast<const char*>( anim->mName.data ) ), frame_rate, duration, flat );
        for ( auto set : temp_ref )
        {
                result->push_back( set );
        }
        delete temp;
        return result;
}

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

nv::sint16 nv::assimp_loader::load_node( uint32 anim_id, array_ref< data_channel_set* > nodes, const void* vnode, sint16 this_id, sint16 parent_id )
{
        const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
        const aiNode*  node  = reinterpret_cast<const aiNode*>( vnode );
        string_view name( static_cast< const char* >( node->mName.data ) );
        const aiAnimation* anim  = scene->mAnimations[anim_id];
        const aiNodeAnim*  anode = nullptr;

        for ( unsigned i = 0 ; i < anim->mNumChannels ; i++ )
        {
                anode = anim->mChannels[i];
                if ( string_view( static_cast< const char* >( anode->mNodeName.data ) ) == name ) break;
                anode = nullptr;
        }

        nodes[ this_id ] = anode ? create_keys( anode ) : data_channel_set_creator::create_set( 0 );

        data_channel_set_creator access( nodes[this_id] );
        access.set_name( make_name( name ) );
        access.set_parent_id( parent_id );
        // This value is ignored by the create_transformed_keys, but needed by create_direct_keys!
        // TODO: find a common solution!
        //       This is bad because create_transformed_keys never uses node-transformations for
        //       node's without keys
        access.set_transform( nv::assimp_mat4_cast( node->mTransformation ) );
        if ( this_id == 0 ) access.set_transform( mat4() );

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

        return next;
}

data_channel_set* nv::assimp_loader::create_keys( const void* vnode )
{
        const aiNodeAnim* node = reinterpret_cast< const aiNodeAnim* >( vnode );
        if ( node->mNumPositionKeys == 0 && node->mNumRotationKeys == 0 && node->mNumScalingKeys == 0 )
        {
                return data_channel_set_creator::create_set( 0 );
        }
        
        data_channel_set* set = data_channel_set_creator::create_set( 2 );
        data_channel_set_creator key_set( set );

        assimp_key_p* pchannel = key_set.add_channel< assimp_key_p >( node->mNumPositionKeys ).data();
        assimp_key_r* rchannel = key_set.add_channel< assimp_key_r >( node->mNumRotationKeys ).data();
//      assimp_key_s* schannel = key_set.add_channel< assimp_key_s >( node->mNumScalingKeys ).data();

        for ( unsigned np = 0; np < node->mNumPositionKeys; ++np )
        {
                pchannel[np].time        = static_cast<float>( node->mPositionKeys[np].mTime );
                pchannel[np].translation = assimp_vec3_cast(node->mPositionKeys[np].mValue);
        }
        for ( unsigned np = 0; np < node->mNumRotationKeys; ++np )
        {
                rchannel[np].time     = static_cast<float>( node->mRotationKeys[np].mTime );
                rchannel[np].rotation = assimp_quat_cast(node->mRotationKeys[np].mValue );
        }
//      if ( node->mNumScalingKeys > 0 )
//      {
//              vec3 scale_vec0 = assimp_vec3_cast( node->mScalingKeys[0].mValue );
//              float scale_value   = glm::length( math::abs( scale_vec0 - 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 )
//                      {
//                              schannel[np].time  = (float)node->mScalingKeys[np].mTime;
//                              schannel[np].scale = assimp_vec3_cast(node->mScalingKeys[np].mValue);
//                      }
//              }
//              else 
//              {
//                      schannel[0].time  = (float)node->mScalingKeys[0].mTime;
//                      schannel[0].scale = assimp_vec3_cast(node->mScalingKeys[0].mValue);
//              }
//      }
        return set;
}

mesh_data_pack* nv::assimp_loader::release_mesh_data_pack()
{
        if ( m_scene == nullptr || m_mesh_count == 0 ) return nullptr;
        const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
        bool has_bones = false;
        data_channel_set* meshes = data_channel_set_creator::create_set_array( m_mesh_count, 2 );
        for ( size_t m = 0; m < m_mesh_count; ++m )
        {
                const aiMesh* mesh = scene->mMeshes[ m ];
                data_channel_set_creator( &meshes[m] ).set_name( make_name( static_cast<const char*>( mesh->mName.data ) ) );
                if ( mesh->mNumBones > 0 ) has_bones = true;
                load_mesh_data(&meshes[m],m);
        }

        mesh_nodes_data* nodes = ( has_bones ? release_merged_bones( meshes ) : release_mesh_nodes_data(0) );
        return new mesh_data_pack( m_mesh_count, meshes, nodes );
}

nv::size_t nv::assimp_loader::get_nodes_data_count() const
{
        if ( m_scene == nullptr ) return 0;
        const aiScene* scene = reinterpret_cast<const aiScene*>( m_scene );
        return scene->mNumAnimations;   
}