#include "stdafx.h"
#include "AssimpReader.h"

#include <LibQtGeoViewerCore/SceneMain.h>

#include "../FileUtil.h"

#include <C2/util/string_util.h>
#include <C2/lm/range1.h>
#include <LibGeo/Path.h>

#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>
#include "AssimpUtil.h"

#include <QFileInfo>
#include <sstream>




bool AssimpReader::Load(SceneMain& scene, const std::string& filename, geom::GeomFileFormat fmt)
{
	unsigned int ai_option =
		aiProcess_JoinIdenticalVertices;

	Assimp::Importer importer;
	const aiScene* ai_scene = importer.ReadFile(filename, ai_option);
	if (!ai_scene)
		return false;

	//PlotASNode(ai_scene->mRootNode);

	std::string dirpath = lib_geo::Path::GetParentDirPath(filename);

	GeomObject* geom = scene.CreateNewGeometry();
	for (unsigned int i = 0; i < ai_scene->mNumMeshes; ++i)
	{
		aiMesh* mesh = ai_scene->mMeshes[i];
		MeshBuf* mbuf = geom->CreateNewMeshBuf();
		scene.RefreshObjectIndex();

		lib_geo::BaseMesh& mesh_dst = mbuf->m_Mesh;
		mbuf->m_Name = mesh->mName.C_Str();

		geom->m_Name = mesh->mName.C_Str();
		if (geom->m_Name.empty())
		{
			std::ostringstream oss;
			oss << FileUtil::GetFileTitle(filename) << "-" << i;
			geom->m_Name = oss.str();
		}

		AssimpUtil::CopyAIVerts(mesh_dst, mesh);
		AssimpUtil::CopyAIFaces(mesh_dst, mesh);
		AssimpUtil::CopyMaterials(mbuf, ai_scene, scene, dirpath);

		ResetBone(mbuf, mesh);

		if (!mesh_dst.HasNormal())
		{
			mesh_dst.CreateNormalsEachVerts();
			mesh_dst.UpdateNormal();
		}

		mbuf->m_Bones.m_SrcVertPos = mesh_dst.m_Verts;
	}

	SetNodeTree(ai_scene, geom);

	geom::NodeMap nodemap;
	AssimpUtil::CreateBoneTree(nodemap, &geom->m_Nodes.m_RootNodes[0]);

	SetBoneNodeRef(nodemap, geom);

	LoadAnimation(nodemap, ai_scene, geom, scene.m_IOConfig.InterpolateAnimation);

	geom->SetFrame(scene.m_CurFrame);

	geom->InitializeBufferCommon();

	for (MeshBuf& mb : geom->m_MeshAry)
	{
		mb.UpdateBBox();
		mb.ResetIniBBox();
	}

	geom->m_FileFormat = fmt;
	geom->m_FilePath   = filename;

	//geom->m_Nodes.TraceNodeTree();

	scene.UpdateTransform();

	scene.ReportDoneEditGeometry();

	return true;
}

void AssimpReader::SetNodeTree(const aiScene* ai_scene, geom::GeomObject* geom)
{
	geom->m_Nodes.Clear();
	geom->m_Nodes.m_RootNodes.resize(1);
	SceneNode* n = &geom->m_Nodes.m_RootNodes[0];
	const aiNode* root = ai_scene->mRootNode;
	SetNodeTreeSub(ai_scene, root, geom, n);
}

void AssimpReader::SetNodeTreeSub(const aiScene* ai_scene, const aiNode* an, geom::GeomObject* geom, geom::SceneNode* gn)
{
	gn->m_Name = an->mName.C_Str();
	gn->m_MeshIDs.resize(an->mNumMeshes);
	for (unsigned int j = 0; j < an->mNumMeshes; ++j)
	{
		int mid = an->mMeshes[j];
		gn->m_MeshIDs[j] = mid;

		MeshBuf& mb = geom->m_MeshAry[mid];
		if (mb.m_Name.empty())
		{
			mb.m_Name = an->mName.C_Str();
		}
	}

	ae::ConvertMat(gn->m_Transform, an->mTransformation);

	for (unsigned int i = 0; i < an->mNumChildren; ++i)
	{
		aiNode* n = an->mChildren[i];
		SceneNode* gcn = new SceneNode();
		gn->m_Children.push_back(gcn);
		gcn->m_Parent = gn;
		SetNodeTreeSub(ai_scene, n, geom, gcn);
	}
}

void AssimpReader::ResetBone(MeshBuf* mbuf, aiMesh* mesh)
{
	mbuf->m_Bones.m_Bones.resize(mesh->mNumBones);
	for (unsigned int j = 0; j < mesh->mNumBones; ++j)
	{
		const aiBone* s_bone = mesh->mBones[j];
		Bone& d_bone = mbuf->m_Bones.m_Bones[j];
		d_bone.m_Name = s_bone->mName.C_Str();

		ConvertBoneWeight(s_bone, d_bone);

		lm::matrix4f offset;
		ae::ConvertMat(offset, s_bone->mOffsetMatrix);
		d_bone.SetOffset(offset);
	}
}

void AssimpReader::ConvertBoneWeight(const aiBone* s_bone, geom::Bone& d_bone)
{
	d_bone.m_Weights.resize(s_bone->mNumWeights);
	for (unsigned int k = 0; k < s_bone->mNumWeights; ++k)
	{
		BoneWeight& bw = d_bone.m_Weights[k];
		aiVertexWeight& vw = s_bone->mWeights[k];
		bw.m_Vid    = vw.mVertexId;
		bw.m_Weight = vw.mWeight;
	}
}

void AssimpReader::SetBoneNodeRef(geom::NodeMap& nodemap, geom::GeomObject* geom)
{
	for (MeshBuf& m : geom->m_MeshAry)
	{
		for (Bone& b : m.m_Bones.m_Bones)
		{
			b.m_RefNode = nodemap.Find(b.m_Name);
		}
	}
}

void AssimpReader::LoadAnimation(geom::NodeMap& nodemap, const aiScene* ai_scene, geom::GeomObject* geom, bool Interpolate)
{
	if (geom->m_Nodes.m_RootNodes.empty())
		return;

	if (ai_scene->mNumAnimations <= 0)
		return;

	const aiAnimation* anim = ai_scene->mAnimations[0];
	double FPS = 30.0 / anim->mDuration;
	double FrameStep = anim->mDuration / 30.0;

	for (unsigned int j = 0; j < anim->mNumChannels; ++j)
	{
		aiNodeAnim* ch = anim->mChannels[j];
		SceneNode* bn = nodemap.Bones[ch->mNodeName.C_Str()];
		if (bn == NULL)
			continue;

		for (size_t k = 0; k < geom->m_MeshAry.size(); ++k)
		{
			MeshBuf& m = geom->m_MeshAry[k];
			for (Bone& b : m.m_Bones.m_Bones)
			{
				if (b.m_Name == bn->m_Name)
				{
					bn->m_MeshIDs.push_back((int)k);
					bn->m_Bone = &b;
					break;
				}
			}
		}

		if (Interpolate)
			ReadKeys_Time(ch, bn, FrameStep);
		else
			ReadKeys_Raw(ch, bn);

		bn->m_KeyIdxOffset = GetNumMinusKeyFrames(ch);
	}
}

void AssimpReader::ReadKeys_Raw(aiNodeAnim* ch, SceneNode* bn)
{
	int num_key = AssimpUtil::GetNumKeyOfFrame(ch);

	// ファイルフォーマットの仕様としては, 属性ごとにキー要素数が異なる可能性があるが,
	// 制御が面倒なので一致しているという前提で処理を作っている.
	bn->m_Translate.resize(num_key);
	bn->m_Scale.resize(num_key);
	bn->m_Rotate.resize(num_key);

	for (unsigned int i = 0; i < ch->mNumPositionKeys; ++i)
	{
		ae::ConvertVec(bn->m_Translate[i], ch->mPositionKeys[i].mValue);
	}

	for (unsigned int i = 0; i < ch->mNumScalingKeys; ++i)
	{
		ae::ConvertVec(bn->m_Scale[i], ch->mScalingKeys[i].mValue);
	}

	std::vector<double> vt;
	for (unsigned int i = 0; i < ch->mNumRotationKeys; ++i)
	{
		vt.push_back(ch->mRotationKeys[i].mTime);
		ae::AiQuatToLmQuat(bn->m_Rotate[i], ch->mRotationKeys[i].mValue);
	}
}

int AssimpReader::RoundF(double f)
{
	int i = (int)f;
	double d = f - (double)i;
	if (d >= 0.9)
		return i + 1;
	else
		return i;
}

int AssimpReader::ToFrame(double time, double TimeMin, double FrameStep)
{
	return RoundF((time - TimeMin) / FrameStep);
}

void AssimpReader::ReadKeys_Time(aiNodeAnim* ch, geom::SceneNode* bn, double FrameStep)
{
	lm::range1<double> time_range;
	if (ch->mNumPositionKeys > 0)
	{
		time_range.expand(ch->mPositionKeys[0].mTime);
		time_range.expand(ch->mPositionKeys[ch->mNumPositionKeys - 1].mTime);
	}
	if (ch->mNumScalingKeys > 0)
	{
		time_range.expand(ch->mScalingKeys[0].mTime);
		time_range.expand(ch->mScalingKeys[ch->mNumScalingKeys - 1].mTime);
	}
	if (ch->mNumRotationKeys > 0)
	{
		time_range.expand(ch->mRotationKeys[0].mTime);
		time_range.expand(ch->mRotationKeys[ch->mNumRotationKeys - 1].mTime);
	}

	int NumTimeKeys = RoundF(time_range.length() / FrameStep) + 1;
	double TimeMin = time_range.min_val();

	lm::vec3f default_t;
	if (ch->mNumPositionKeys > 0)
		ae::ConvertVec(default_t, ch->mPositionKeys[ch->mNumPositionKeys - 1].mValue);

	lm::vec3f default_s;
	if (ch->mNumScalingKeys > 0)
		ae::ConvertVec(default_s, ch->mScalingKeys[ch->mNumScalingKeys - 1].mValue);

	lm::quat4f default_r;
	if (ch->mNumRotationKeys > 0)
		ae::AiQuatToLmQuat(default_r, ch->mRotationKeys[ch->mNumRotationKeys - 1].mValue);

	bn->m_Translate.resize(NumTimeKeys, default_t);
	bn->m_Scale.resize(NumTimeKeys, default_s);
	bn->m_Rotate.resize(NumTimeKeys, default_r);

	for (unsigned int i = 0; i < ch->mNumPositionKeys; ++i)
	{
		lm::vec3f t;
		const aiVectorKey& pk = ch->mPositionKeys[i];
		ae::ConvertVec(t, pk.mValue);
		if (i == 0)
		{
			int ft = ToFrame(pk.mTime, TimeMin, FrameStep);
			for (int j = 0; j <= ft; ++j)
			{
				bn->m_Translate[j] = t;
			}
		}
		else
		{
			const aiVectorKey& pkl = ch->mPositionKeys[i - 1];
			int tmi0 = ToFrame(pkl.mTime, TimeMin, FrameStep);
			int tmi1 = ToFrame(pk.mTime, TimeMin, FrameStep);

			lm::vec3f tl = bn->m_Translate[tmi0];
			for (int j = tmi0 + 1; j <= tmi1; ++j)
			{
				double n = (double)(j - tmi0) / (double)(tmi1 - tmi0);
				bn->m_Translate[j] = t * (float)n + tl * (float)(1.0 - n);
			}
		}
	}

	for (unsigned int i = 0; i < ch->mNumScalingKeys; ++i)
	{
		lm::vec3f s;
		const aiVectorKey& sk = ch->mScalingKeys[i];
		ae::ConvertVec(s, sk.mValue);
		if (i == 0)
		{
			int ft = ToFrame(sk.mTime, TimeMin, FrameStep);
			for (int j = 0; j <= ft; ++j)
			{
				bn->m_Scale[j] = s;
			}
		}
		else
		{
			const aiVectorKey& skl = ch->mScalingKeys[i - 1];
			int tmi0 = ToFrame(skl.mTime, TimeMin, FrameStep);
			int tmi1 = ToFrame(sk.mTime, TimeMin, FrameStep);

			lm::vec3f sl = bn->m_Scale[tmi0];
			for (int j = tmi0 + 1; j <= tmi1; ++j)
			{
				double n = (double)(j - tmi0) / (double)(tmi1 - tmi0);
				float fn = (float)n;
				float frn = (float)(1.0 - n);
				bn->m_Scale[j].x = powf(s.x, fn) * powf(sl.x, frn);
				bn->m_Scale[j].y = powf(s.y, fn) * powf(sl.y, frn);
				bn->m_Scale[j].z = powf(s.z, fn) * powf(sl.z, frn);
			}
		}
	}

	for (unsigned int i = 0; i < ch->mNumRotationKeys; ++i)
	{
		lm::quat4f q;
		const aiQuatKey& rk = ch->mRotationKeys[i];
		ae::AiQuatToLmQuat(q, rk.mValue);
		if (i == 0)
		{
			int ft = ToFrame(rk.mTime, TimeMin, FrameStep);
			for (int j = 0; j <= ft; ++j)
			{
				bn->m_Rotate[j] = q;
			}
		}
		else
		{
			const aiQuatKey& rkl = ch->mRotationKeys[i - 1];
			size_t tmi0 = ToFrame(rkl.mTime, TimeMin, FrameStep);
			size_t tmi1 = ToFrame(rk.mTime, TimeMin, FrameStep);

			lm::quat4f ql;
			ae::AiQuatToLmQuat(ql, rkl.mValue);
			for (int j = tmi0 + 1; j <= tmi1; ++j)
			{
				double n = (double)(j - tmi0) / (double)(tmi1 - tmi0);
				bn->m_Rotate[j] = lm::quat4f::interpolate(ql, q, (float)n);
			}
		}
	}
}

void AssimpReader::PlotASNode(aiNode* n)
{
	static int layer = 0;

	{
		std::ostringstream ss;
		for (int i = 0; i < layer; ++i)
		{
			ss << "|";
		}

		ss << "+";
		ss << n->mName.C_Str();
		ss << "[" << n->mNumMeshes << "]";
		ss << std::endl;
		
		OutputDebugStringA(ss.str().c_str());
	}

	for (unsigned int i = 0; i < n->mNumChildren; ++i)
	{
		layer++;
		PlotASNode(n->mChildren[i]);
		layer--;
	}
}

// 0秒未満のキーフレーム数を数える
int AssimpReader::GetNumMinusKeyFrames(const aiNodeAnim* ch) const
{
	for (unsigned int i = 0; i < ch->mNumPositionKeys; ++i)
	{
		if(ch->mPositionKeys[i].mTime >= 0)
			return i;
	}

	// 全フレームが0秒未満
	assert(false);

	return ch->mNumPositionKeys;
}
