#include "StdAfx.h"
#include "mayacryboneutil.h"
#include "MayaUtils.h"

CMayaCryBoneUtil::CMayaCryBoneUtil(void)
{
}

CMayaCryBoneUtil::~CMayaCryBoneUtil(void)
{
}


// returns the number of bones in the hierarchy
int CMayaCryBoneUtil::numBones() 
{
	validateBoneHierarchy();
	return m_arrBones.size ();
}

// returns the ith bone object
const CMayaCryBoneUtil::BoneEntry& CMayaCryBoneUtil::getBone(int nBone)
{
	validateBoneHierarchy();
	return m_arrBones[nBone];
}


//////////////////////////////////////////////////////////////////////////
// returns the array of bones valid for the current state of the utility
void CMayaCryBoneUtil::getBones (BoneHierarchy& arrBones)
{
	validateBoneHierarchy();
	arrBones = m_arrBones;
}


//////////////////////////////////////////////////////////////////////////
// validates the linearized hierarchy (m_arrBones) - constructs it from the m_arrRootBones if necessary
void CMayaCryBoneUtil::validateBoneHierarchy()
{
	if (m_arrBones.size () == 0 && m_arrRootBones.length () > 0)
	{
		constructBoneHierarchy();
	}
}



//////////////////////////////////////////////////////////////////////////
// flags the bone hierarchy (m_arrBones) invalid so that the next time the validate function is called, it's recalculated
void CMayaCryBoneUtil::invalidateBoneHierarchy()
{
	m_arrBones.clear ();
	m_mapBoneIndex.clear();
}


//////////////////////////////////////////////////////////////////////////
// constructs the bone hierarchy out from the array of root bones.
// assumes that the root bones are not children of each other
void CMayaCryBoneUtil::constructBoneHierarchy()
{
	m_arrBones.clear ();
	m_mapBoneIndex.clear();

#ifdef _DEBUG
	{
		unsigned i, j;
		for (i = 0;i < m_arrRootBones.length (); ++i)
			for (j = i + 1; j < m_arrRootBones.length (); ++j)
			{
				MFnDagNode fnNode (m_arrRootBones[i]);
				assert (!fnNode.isChildOf (m_arrRootBones[j].node ()));
				assert (!fnNode.hasChild (m_arrRootBones[j].node ()));
			}
	}
#endif

	for (unsigned i = 0; i < m_arrRootBones.length () ; ++i)
		addBoneHierarchy (m_arrRootBones[i], BoneEntry::kNoBone); // the root bone has no parent
}



//////////////////////////////////////////////////////////////////////////
// adds the node and its subhierarchy to the array of bones and bone index map
void CMayaCryBoneUtil::addBoneHierarchy (MDagPath& pathBone, int nParentId)
{
	int nBoneId = m_arrBones.size (); // id of this bone
	m_mapBoneIndex.insert (BoneIndexMap::value_type (pathBone, nBoneId));

	// find the number of children for the bone;
	// add the bone to the end of the list;
	// add all the children of the bone recursively

	MFnDagNode fnBone (pathBone);
	MObjectArray arrChildren;
	unsigned i;
	for (i = 0; i < fnBone.childCount (); ++i)
	{
		MObject objChild = fnBone.child (i);
		if (objChild.hasFn (MFn::kTransform))
			arrChildren.append (objChild);
	}

	// add all the children of the bone recursively
	m_arrBones.push_back (BoneEntry(pathBone, arrChildren.length (), nParentId));
	for (i = 0; i < arrChildren.length (); ++i)
	{
		MObject objChild = arrChildren[i];
		pathBone.push (objChild);
		addBoneHierarchy (pathBone, nBoneId);
		pathBone.pop ();
	}
}



//////////////////////////////////////////////////////////////////////////
// Adds the given shape's influencing nodes to the list of bones and returns
// the number of influencing bones (if 0, then the shape is static)
// This is used to determine the whole skeleton parts of which may influence
// the shape and add the whole skeleton to the list of bones to export.
int CMayaCryBoneUtil::addShapeInfluences(const MObject& objShape)
{
	MDagPathArray arrBones;
	findShapeInfluences (objShape, arrBones);

	unsigned i;
	for (i = 0; i < arrBones.length (); ++i)
	{
		addRootBone (arrBones[i]);
	}

	return arrBones.length ();
}

//////////////////////////////////////////////////////////////////////////
// adds another bone to the array of root bones
void CMayaCryBoneUtil::addRootBone (const MDagPath& pathRootBone)
{
	// determine if the bone is a child of one of the already existing root bones; then do nothing
	// determine if the bone is a parent of one of the already existing root bones; then delete them and add it
	unsigned i;
	
	bool bAlreadyAdded = false;

	for (i = 0; i < m_arrRootBones.length ();)
	{
		MFnDagNode fnRootNode(m_arrRootBones[i]);
		if (fnRootNode.isChildOf (pathRootBone.node ()))
		{
			// the already-existing root bone is a child of this one;
			// for the first time, replace it with the new one; the next time, just erase the child
			if (bAlreadyAdded)
			{
				m_arrRootBones.remove (i);
			}
			else
			{
				m_arrRootBones[i] = pathRootBone;
				bAlreadyAdded = true;
				++i;
			}
		}
		else
		if (fnRootNode.isParentOf (pathRootBone.node()))
		{
			assert (!bAlreadyAdded);
			bAlreadyAdded = true;
			break;
		}
		else
			++i; // the bones are irrelated
	}

	if (!bAlreadyAdded)
		m_arrRootBones.append (pathRootBone);
}



//////////////////////////////////////////////////////////////////////////
// returns the index of the given bone in the linearized hierarchy, or -1 if not found
int CMayaCryBoneUtil::getBoneIndex (const MDagPath& pathBone)
{
	validateBoneHierarchy();

	BoneIndexMap::iterator it = m_mapBoneIndex.find (pathBone);
	if (it == m_mapBoneIndex.end())
		return -1; // didn't find the bone
	else
		return it->second;
}


//////////////////////////////////////////////////////////////////////////
// fills in the index map for the given array of dag paths:
// for i-th DagPath, it's index is put into the i-th element of the passed array
void CMayaCryBoneUtil::getBoneIndexMap (const MDagPathArray& arrBones, std::vector<int>& arrMap)
{
	arrMap.resize (arrBones.length());
	for (unsigned i  =0; i < arrBones.length(); ++i)
	{
		arrMap[i] = getBoneIndex (arrBones[i]);
	}
}

// returns the number of actual roots (hierarchies) in the list of bones
unsigned CMayaCryBoneUtil::numRoots()const
{
	return m_arrRootBones.length ();
}

// returns the ith root bone
const MDagPath& CMayaCryBoneUtil::getRoot (unsigned i)const
{
	assert (i>=0 && i < m_arrRootBones.length());
	return m_arrRootBones[i];
}


// cleans up the object
void CMayaCryBoneUtil::clear()
{
	m_arrBones.clear ();
	invalidateBoneHierarchy();
}


// puts all bones into the rest position
void CMayaCryBoneUtil::resetFromRestPosition ()
{
	MStatus status;
	BoneHierarchy::iterator it;
	MFnTransform fnTM;
	for (it = m_arrBones.begin (); it != m_arrBones.end(); ++it)
	{
		status = fnTM.setObject (it->pathBone);
		if (!status)
		{
			Log ("*ERROR* Bone %s is not a TM: %s", it->pathBone.fullPathName ().asChar (), status.errorString ().asChar ());
			continue;
		}
		
		status = fnTM.resetFromRestPosition ();
		if (!status)
		{
			Log ("*ERROR* Cannot set the rest TM for the Bone %s: %s", it->pathBone.fullPathName ().asChar (), status.errorString ().asChar ());
			continue;
		}
		
		/*
		MTransformationMatrix tm = fnTM.restPosition (&status);
		if (!status)
		{
			Log ("*ERROR* Bone %s has no rest position: %s", it->pathBone.fullPathName ().asChar (), status.errorString ().asChar ());
			continue;
		}

		status = fnTM.set (tm);
		if (!status)
		{
			Log ("*ERROR* Cannot set the rest TM for the Bone %s: %s", it->pathBone.fullPathName ().asChar (), status.errorString ().asChar ());
			continue;
		}
		*/
	}
}


//////////////////////////////////////////////////////////////////////////
// saves all bones' transform position
void CMayaCryBoneUtil::pushPosition()
{
	validateBoneHierarchy ();

	m_stkPositions.push (BonePositions());
	BonePositions& arrSaveTMs = m_stkPositions.top();

	MFnTransform fnTM;

	arrSaveTMs.resize (m_arrBones.size());
	for (unsigned i = 0; i < m_arrBones.size(); ++i)
	{
		if (!fnTM.setObject (m_arrBones[i].pathBone))
		{
			Log ("*ERROR* Bone %s is not a transform", m_arrBones[i].pathBone.fullPathName ().asChar ());
			continue;
		}
		arrSaveTMs[i] = fnTM.transformation ();
	}
}


//////////////////////////////////////////////////////////////////////////
// restores all bones' position
void CMayaCryBoneUtil::popPosition()
{
	validateBoneHierarchy();

	if (m_stkPositions.empty ())
		throw MStatus::kFailure;

	BonePositions& arrSavedTMs = m_stkPositions.top();
	if (arrSavedTMs.size() != m_arrBones.size())
		throw MStatus::kFailure;

	MFnTransform fnTM;

	for (unsigned i = 0; i < m_arrBones.size(); ++i)
	{
		if (!fnTM.setObject (m_arrBones[i].pathBone))
		{
			Log ("*ERROR* Cannot restore position for bone %s because it is not a TM", m_arrBones[i].pathBone.fullPathName().asChar ());
			continue;
		}

		fnTM.set (arrSavedTMs[i]);
	}

	m_stkPositions.pop();
}

// returns the comma-separated list of root bones
MString CMayaCryBoneUtil::getRootBoneDump ()const
{
	MString strBoneDump;
	unsigned i;
	for (i = 0; i < m_arrRootBones.length(); ++i)
	{
		if (i > 0)
			strBoneDump += ", ";
		strBoneDump += m_arrRootBones[i].partialPathName();
	}
	return strBoneDump;
}
