//-
//*****************************************************************************
// Copyright 2013 Autodesk, Inc. All rights reserved.
// 
// Use of this software is subject to the terms of the Autodesk license
// agreement provided at the time of installation or download, or which
// otherwise accompanies this software in either electronic or hard copy
// form.
//*****************************************************************************
//+

#include "AbcBulletStringTable.h"
#include "MayaTransformCollectionWriter.h"
#include "MayaUtility.h"
#include <maya/MDataHandle.h>
#include <maya/MFnPluginData.h>
#include <maya/MPxData.h>
#include <maya/MDagPath.h>
#include <maya/MSelectionList.h>
#include <maya/MItSelectionList.h>
#include <maya/MQuaternion.h>
#include <maya/MGlobal.h>
#include <maya/MStringResource.h>

class MotionStateStreamBuf : public std::streambuf
{
protected:
	enum { kInitialTranslations=8, kInitialOrientations=9, kMotionStates=10 } ;

	int mItemsRead;

	int mSection;
	int mNumItems;
	std::vector<float>	mMotionStateData;

public:
	// section header, numItems,  numItems x 3 x btTransform< btMatrix3x3<9 x float>, btVector3<3 x float> >
	enum {
		kSizeOfVector3 = 3,
		kSizeOfMatrix3 = 3*3,
		kSizeOfTransform = kSizeOfMatrix3 + kSizeOfVector3,
		kSizeOfMotionState = 3 * kSizeOfTransform,
	};

	MotionStateStreamBuf() : mItemsRead(0), mSection(kInitialTranslations) { ; }
	~MotionStateStreamBuf() { sync(); }

	int NumItems() { return mNumItems; };
	const std::vector<float>& data() const { return mMotionStateData; }

	void reset() { mItemsRead = 0; mMotionStateData.clear(); }

protected:
	virtual std::streamsize xsputn(const char * _Ptr, std::streamsize _Count) 
	{	// put _Count characters to stream
		std::streamsize _Copied(_Count);

		++mItemsRead;

		// section header
		if (1==mItemsRead)
		{
			::memcpy( &mSection, _Ptr, sizeof(int) );
		}
		// number of items
		else if (2==mItemsRead)
		{
			assert(_Count==sizeof(int));
			::memcpy( &mNumItems, _Ptr, sizeof(int) );
		}
		else
		{
			// TODO: we need to removed the initial transform information
			// off the motionstate stream to reduce the amount data read on each frame.
			switch (mSection)
			{
			case kInitialTranslations:
				// section header, numItems,  list of btVector3<float,float,float>
				if (mItemsRead == (mNumItems * 3) + 2)
				{
					mItemsRead=0;
				}
				break;
			case kInitialOrientations:
				// section header, numItems,  list of btQuaternion<float,float,float,float>
				if (mItemsRead == (mNumItems * 4) + 2)
				{
					mItemsRead=0;
				}
				break;
			case kMotionStates:
				// TODO: we need to only save the rotation & translation (4+3) instead of the entire transform (9+3)

				// section header, numItems,  numItems x 3 x btTransform< btMatrix3x3<9 x float>, btVector3<3 x float> >
				// motionstate = current transform, center of mass transform, start transform
				if (mItemsRead == (mNumItems * kSizeOfMotionState ) + 2)
				{
					mItemsRead=0;
				}
				else
				{
					float val(.0f);
					::memcpy( &val, _Ptr, sizeof(float) );

					mMotionStateData.push_back(val);
				}

				break;
			}
		}

		return (_Copied);
	}

};

// This is the output stream; its function is to format data (using mainly the <<
// operator) and send it to a streambuf to be stored and written to the output.
class MotionStateStream : public std::ostream
{
public:
	MotionStateStream() : ostream(new MotionStateStreamBuf()), ios(0) {}
	~MotionStateStream() { ; }

	void resetBuf() {  dynamic_cast<MotionStateStreamBuf*>(rdbuf())->reset(); }
};


// copied from b3Matrix3x3.h
void getRotation( const float*  m3, double * result )
{
	static int _el0 = 0;
	static int _el1 = 3;
	static int _el2 = 6;
	static int _X = 0;
	static int _Y = 1;
	static int _Z = 2;

	float trace = m3[_el0+_X] + m3[_el1+_Y] + m3[_el2+_Z];

	if (trace > float(0.0)) 
	{
		float s = ::sqrt(trace + float(1.0));
		result[3]=(s * float(0.5));
		s = float(0.5) / s;

		result[0]=((m3[_el2+_Y] - m3[_el1+_Z]) * s);
		result[1]=((m3[_el0+_Z] - m3[_el2+_X]) * s);
		result[2]=((m3[_el1+_X] - m3[_el0+_Y]) * s);
	} 
	else 
	{
		int i = m3[_el0+_X] < m3[_el1+_Y] ? 
			(m3[_el1+_Y] < m3[_el2+_Z] ? 2 : 1) :
			(m3[_el0+_X] < m3[_el2+_Z] ? 2 : 0); 
		int j = (i + 1) % 3;  
		int k = (i + 2) % 3;

		float s = ::sqrt( m3[(i*3)+i] - m3[(j*3)+j] - m3[(k*3)+k] + float(1.0));
		result[i] = s * float(0.5);
		s = float(0.5) / s;

		result[3] = (m3[(k*3)+j] - m3[(j*3)+k]) * s;
		result[j] = (m3[(j*3)+i] + m3[(i*3)+j]) * s;
		result[k] = (m3[(k*3)+i] + m3[(i*3)+k]) * s;
	}
}

double MayaTransformCollectionItem::asDouble(Alembic::AbcGeom::XformOperationType channelOp, Alembic::Util::uint32_t channelNum)
{
	double result = 0.0;

	int offset = mItemID * MotionStateStreamBuf::kSizeOfMotionState;

	if (isTranslationChannel(channelOp,channelNum))
	{
		int trsOffset = offset + MotionStateStreamBuf::kSizeOfMatrix3;
		int idx = trsOffset + channelNum;
		result = mBuf->data()[idx];

#ifdef _DEBUG
		if (channelNum==0)
		{
			if (mVerbose)
			{
				MString str;
				MStringArray args(4,"");

				unsigned int i = 0;

				args[i++] = mName;
				args[i++].set(mBuf->data()[trsOffset+0]);
				args[i++].set(mBuf->data()[trsOffset+1]);
				args[i++].set(mBuf->data()[trsOffset+2]);

				str.format("^1s: translate ^2s, ^3s, ^4s", args);
				MGlobal::displayInfo( str );
			}
		}
#endif
	}
	else if (isRotationChannel(channelOp,channelNum)) 
	{
		int index = (channelOp==Alembic::AbcGeom::kRotateXOperation) ? 0 : (channelOp==Alembic::AbcGeom::kRotateYOperation) ? 1 : 2;

		if (channelOp==Alembic::AbcGeom::kRotateXOperation)
		{
			// TODO: we should do this on the read and cache the result.
			const float* m3 = &mBuf->data()[offset];

			double qData[4];

			getRotation(m3,qData);

			MQuaternion q(qData);

			mEulerRotation = q;

#ifdef _DEBUG
			if (mVerbose)
			{
				MString str;
				MStringArray args(10,"");

				unsigned int i = 0;

				args[i++].set(qData[0]);
				args[i++].set(qData[1]);
				args[i++].set(qData[2]);
				args[i++].set(qData[3]);

				args[i++].set(mEulerRotation[0]);
				args[i++].set(mEulerRotation[1]);
				args[i++].set(mEulerRotation[2]);

				args[i++].set(Alembic::AbcGeom::RadiansToDegrees(mEulerRotation[0]));
				args[i++].set(Alembic::AbcGeom::RadiansToDegrees(mEulerRotation[1]));
				args[i++].set(Alembic::AbcGeom::RadiansToDegrees(mEulerRotation[2]));

				str.format(": rotate quat[^1s, ^2s, ^3s, ^4s] euler radians[^5s, ^6s, ^7s] degrees[^8s, ^9s, ^10s]", args);
				str = mName + str;
				MGlobal::displayInfo( str );
			}
#endif
		}
		result = mEulerRotation[index];
	}

	return result;
}


MayaTransformCollectionWriter::MayaTransformCollectionWriter(Alembic::AbcGeom::OObject & iParent,
	MDagPath & iDag, Alembic::Util::uint32_t iTimeIndex, const JobArgs & iArgs)
{
	MStatus stat;

	// get initial state from iDag (SolvedState)
	MFnDependencyNode depSolvedState(iDag.node(&stat));

	// TODO: we want to access the current transform solution off the solved state
	// Ideally, we would access an array attribute which would contain the solved item's fullpath
	// and transform.
	MObject initialState;
	{
		MPlug isPlug = depSolvedState.findPlug("initialState", &stat);

		MPlugArray connections;
		isPlug.connectedTo(connections, true /*asDest*/, false/*asSource*/ );

		if (connections.length())
		{
			initialState = connections[0].node();
		}
	}

	MFnDependencyNode depInitState(initialState);
	
	// For backwards compatibility, use set membership if dagPathNames attribute is absent
	//
	MPlug dagPathNamesPlug = depInitState.findPlug("dagPathNames", &stat);
	MFnStringArrayData arrData(dagPathNamesPlug.asMObject());
	bool bUseSetMembership = ( (MS::kSuccess != stat) || (arrData.length()==0) );

	MObject dagPathNamesArrayData;
	MFnStringArrayData mfnStringArrayData;
	MSelectionList objects;
	
	if ( bUseSetMembership )
	{
		// access rigidset for node paths
		{
			MPlug setPlug = depInitState.findPlug("message", &stat);
			MPlugArray connections;
			setPlug.connectedTo(connections, false /*asDest*/, true/*asSource*/ );
			MObject rigidSet;

			for(unsigned j = 0, m = connections.length(); j < m; ++j) {
				MObject src = connections[j].node();

				if (src.hasFn(MFn::kSet)) {
					rigidSet = src;
					break;
				}
			}

			MFnDependencyNode rigidSetNode(rigidSet);
			MString error = MStringResource::getString( kDagPathOrderMissingErr, stat );
			error.format( error, depInitState.name(), rigidSetNode.name());
			MGlobal::displayError( error );

			MFnSet fnSet(rigidSet);
			fnSet.getMembers(objects, /*flatten*/true);
		}
	}
	else
	{
		mfnStringArrayData.setObject(dagPathNamesPlug.asMObject());
	}

	// access motionstate attribute
	mPlug = depInitState.findPlug("motionStates", &stat);

	// create data stream
	mDataStream = ostreamPtr(new MotionStateStream());
	MotionStateStreamBuf* buf = dynamic_cast<MotionStateStreamBuf*>(mDataStream->rdbuf());
	assert(buf);

	// write motionstate to memory stream
	readDataStream();

	// add motionstates to transform write list
	MString fullpath, name;

	MItSelectionList iterList( objects );

	MObject  component;
	MDagPath dagPath;

	unsigned int stringArraylength = mfnStringArrayData.length();

	for (unsigned i=0, m = buf->NumItems(); i < m ; i++)
	{
		if ( bUseSetMembership )
		{
			if(iterList.isDone()) break;

			iterList.getDagPath(dagPath, component);

			// add leaf node
			fullpath = dagPath.fullPathName( &stat );

			MFnDependencyNode depNode(dagPath.node());
			name =  depNode.name();
		}
		else
		{
			if(i >= stringArraylength) break;

			// add leaf node
			fullpath = mfnStringArrayData[i];
			name = fullpath.substringW(fullpath.rindex('|') + 1, fullpath.length()-1);

			MSelectionList sList;
			MStatus status = sList.add(fullpath);
			sList.getDagPath(0, dagPath);
		}

		MayaTransformCollectionItemPtr sampler = MayaTransformCollectionItemPtr(new MayaTransformCollectionItem(name, fullpath, i, buf, iArgs.verbose));
		mSamplerList.push_back(sampler);
		MayaTransformWriterPtr trans;
		
		// Copy the dag path because we'll be popping from it
		//
		MDagPath dag(dagPath);
	
		// precondition: iParent will already be at the root of the tree
		//	
		Alembic::Abc::OObject iRoot = iParent;
		Alembic::Abc::OObject iCurrent = iRoot;

		// Create the non-animation-based transforms (if any)
		//
		int j;
		int numPaths = dag.length();
		if (numPaths>1)
		{
			std::vector< MDagPath > dagList;

			for (j = numPaths - 1; j > -1; j--, dag.pop())
			{
				dagList.push_back(dag);
			}

			std::vector< MDagPath >::iterator iStart = dagList.begin();
			std::vector< MDagPath >::iterator iCur = dagList.end();
			iCur--;

			// now loop backwards over our DAG path list so we push ancestor nodes
			// first, all the way down to the current node
			//
			// This essentially reads the DAG paths left to right (or top down)
			// and checks to see if their components already exist in Alembic;
			// if a DAG path component is missing, it'll be added.
			//
			for (; iCur != iStart; iCur--)
			{
				MString currentDagPathName = (*iCur).fullPathName();
				MStringArray pathArray;
				currentDagPathName.split('|', pathArray);
				iCurrent = iRoot;
				for (unsigned int i = 0; i < pathArray.length(); i++) {
					// make sure you strip the namespace off the path before you
					// try to look up the alembic node.
					MString step = util::stripNamespaces(pathArray[i], iArgs.stripNamespace);
					Alembic::Abc::OObject iPriorCurrent = iCurrent;
					iCurrent = iCurrent.getChild(step.asChar());
					if (!iCurrent.valid()) {
						iCurrent = iPriorCurrent;
						trans = MayaTransformWriterPtr(new MayaTransformWriter(iCurrent, *iCur, iTimeIndex, iArgs));
						iCurrent = iCurrent.getChild(step.asChar());
					}
				}
			}
		}
		assert(iCurrent.valid());

		// Remember to continue the hierarchy - to include the animation-based transform
		// It is assumed that a leaf entry won't exist in Alembic due to the uniqueness of a full transform path
		//

		// Create the animation-based transform
		//
		trans = MayaTransformWriterPtr(new MayaTransformWriter(
			iCurrent, *sampler, iTimeIndex, iArgs));

		//TODO: Remove
		mTransList.push_back(trans);

		AttributesWriterPtr attrs = trans->getAttrs();
		if (attrs)
		{
			if (iTimeIndex != 0 && attrs->isAnimated())
				mTransAttrList.push_back(attrs);
		}

		if ( bUseSetMembership ) iterList.next();
	}
}

MayaTransformCollectionWriter::~MayaTransformCollectionWriter()
{
}

void MayaTransformCollectionWriter::readDataStream()
{
	MStatus stat;

	MObject mData;
	stat = mPlug.getValue(mData);

	MFnPluginData pdFn(mData); 

	// get data handle to write to ostream
	MPxData* mPxData = pdFn.data(&stat);
	assert(mPxData);

	assert(mPxData);
	assert(mDataStream);

	if (mDataStream)
	{
		dynamic_cast<MotionStateStreamBuf*>(mDataStream->rdbuf())->reset();
		if (mPxData)
			mPxData->writeBinary(*mDataStream);
	}
}

void MayaTransformCollectionWriter::write()
{
	// read data from motionstates
	readDataStream();

	// write out transforms -- TODO: Remove
	std::vector< MayaTransformWriterPtr >::iterator tcur =
		mTransList.begin();

	std::vector< MayaTransformWriterPtr >::iterator tend =
		mTransList.end();

	for (; tcur != tend; tcur++)
	{
		(*tcur)->write();
	}

	std::vector< AttributesWriterPtr >::iterator tattrCur =
		mTransAttrList.begin();

	std::vector< AttributesWriterPtr >::iterator tattrEnd =
		mTransAttrList.end();

	for(; tattrCur != tattrEnd; tattrCur++)
	{
		(*tattrCur)->write();
	}
}

bool MayaTransformCollectionWriter::isAnimated() const
{
	return true;
}