GTASource/rage/base/tools/dcc/maya/rageSnapNodes/rageSnapNodesCmd.cpp

//
// Copyright (C) Rockstar San Diego
// 
// File: rageSnapNodesCmd.cpp
//
// MEL Command: rageSnapNodes
//
// Author: Steven Waller
//

#include "mayahelp.h"

#include "rageSnapNodesCmd.h"


#define REQUIRE_IOSTREAM
#pragma warning(push)
#pragma warning(disable : 4668)
#include <maya/MArgDatabase.h>
#include <maya/MDagPath.h>
#include <maya/MFloatPoint.h>
#include <maya/MFloatVector.h>
#include <maya/MFnMesh.h>
#include <maya/MFnTransform.h>
#include <maya/MGlobal.h>
#include <maya/MItSelectionList.h>
#include <maya/MSelectionList.h>
#include <maya/MString.h>
#include <maya/MStringArray.h>
#include <maya/MSyntax.h>
#include <maya/MVector.h>
#include <maya/MVectorArray.h>
#pragma warning(pop)


#define rageErrorCheckGeneral(stat, message)	\
	if ( MS::kSuccess != stat )				\
	{										\
		MGlobal::displayError(message);		\
		return MS::kFailure;				\
	}

#define rageErrorCheckDagPath(stat, dagPath, message)	\
	if ( MS::kSuccess != stat )							\
	{													\
		MStatus status;									\
		MString fullPathName, fullMessage;				\
		fullPathName = dagPath.fullPathName(&status);	\
		fullMessage = (fullPathName + ":" + message);	\
		MGlobal::displayError(fullMessage);				\
		return MS::kFailure;							\
	}


typedef unsigned int UINT;


const char	*axisFlag = "-a", *axisLongFlag = "-axis";
MString		X_AXIS = "x", Y_AXIS = "y", Z_AXIS = "z";


MSyntax rageSnapNodes::addFlags()
{
	MSyntax syntax;

	syntax.addFlag( axisFlag, axisLongFlag, MSyntax::kString );

	return syntax;
}


MStatus isAxisValid(MString *axisLetter)
{
	if ( *axisLetter == X_AXIS || *axisLetter == "X" )
	{
		*axisLetter = X_AXIS;
	}
	else if ( *axisLetter == Y_AXIS || *axisLetter == "Y" )
	{
		*axisLetter = Y_AXIS;
	}
	else if ( *axisLetter == Z_AXIS || *axisLetter == "Z" )
	{
		*axisLetter = Z_AXIS;
	}
	else
	{
		return MS::kFailure;
	}

	return MS::kSuccess;
}


MStatus getIntersection(MFloatPoint raySource, MFloatVector rayDirection, MDagPath dpMesh, MFloatPoint &hitPoint, bool &bGotHit)
{
	MStatus stat;
	
	MFnMesh fnDestMesh(dpMesh, &stat);
	rageErrorCheckDagPath(stat, dpMesh, "Unable to initialize the MFnMesh function set in getIntersection!");

	bGotHit = fnDestMesh.closestIntersection ( raySource, rayDirection, NULL, NULL, 0, MSpace::kWorld, 10000.0, 0, NULL, hitPoint, NULL, NULL, NULL, NULL, NULL, 0.0, &stat );
	rageErrorCheckGeneral(stat, "Unable to determine the closest intersection!");

	return stat;
}


UINT isTolerant(MFloatPoint raySource, MFloatPoint hitPoint, MString axisLetter, float fTolerance)
{
	UINT iIsTolerant = 0;
	
	//Handle only hits within our tolerance.
	float fRayToleranceX = ((int)(raySource.x * fTolerance) / fTolerance);
	float fHitToleranceX = ((int)(hitPoint.x * fTolerance) / fTolerance);
	float fRayToleranceY = ((int)(raySource.y * fTolerance) / fTolerance);
	float fHitToleranceY = ((int)(hitPoint.y * fTolerance) / fTolerance);
	float fRayToleranceZ = ((int)(raySource.z * fTolerance) / fTolerance);
	float fHitToleranceZ = ((int)(hitPoint.z * fTolerance) / fTolerance);
	
	if(axisLetter == X_AXIS)
	{
		if(fRayToleranceY == fHitToleranceY && fRayToleranceZ == fHitToleranceZ)
		{
			iIsTolerant = 1;
		}
	}
	else if(axisLetter == Y_AXIS)
	{
		if(fRayToleranceX == fHitToleranceX && fRayToleranceZ == fHitToleranceZ)
		{
			iIsTolerant = 1;
		}
	}
	else if(axisLetter == Z_AXIS)
	{
		if(fRayToleranceX == fHitToleranceX && fRayToleranceY == fHitToleranceY)
		{
			iIsTolerant = 1;
		}
	}

	return iIsTolerant;
}


MStatus isApiType(MDagPath dpNode, MFn::Type apiType)
{
	if(dpNode.apiType() != apiType)
	{
		return MS::kFailure;
	}

	return MS::kSuccess;
}


MStatus rageSnapNodes::doIt( const MArgList &args )
//
//	Description:
//		implements the MEL rageSnapNodes command.
//
//	Arguments:
//		args - the argument list that was passes to the command from MEL
//
//	Return Value:
//		MS::kSuccess - command succeeded
//		MS::kFailure - command failed (returning this value will cause the 
//                     MEL script that is being run to terminate unless the
//                     error is caught using a "catch" statement.
//
{
	//float f_TOLERANCE = 10000.0; //Tolerance for the intersection operation.

	MDagPath		dpNode;
	MSelectionList	selectionList;
	MStatus			stat;

	//Get the axis for snapping nodes to the destination mesh.
	MArgDatabase argData( syntax(), args );
	MString axisLetter;

	if( argData.isFlagSet( axisFlag ) )
	{
		stat = argData.getFlagArgument( axisFlag, 0, axisLetter );
		rageErrorCheckGeneral(stat, "Invalid axis letter specified!");
	}
	else
	{
		MGlobal::displayError ("No -a (-axis) flag set!");
		return MS::kFailure;
	}

	stat = isAxisValid(&axisLetter);
	rageErrorCheckGeneral(stat, "Invalid -a (-axis) flag!  Please use x, y or z...");
	
	// Get a list of currently selected objects
	MGlobal::getActiveSelectionList( selectionList );

	//Make sure the User has the proper selection made for the command to function.
	UINT numSelectedItems = selectionList.length();
	
	if(numSelectedItems < 2)
	{
		MGlobal::displayError ("You must have at least two nodes selected - node(s) to snap and then destination mesh (last selected)!");
		return MS::kFailure;
	}

	//Get the destination mesh node info.
	UINT indexDestMesh = numSelectedItems - 1;

	stat = selectionList.getDagPath( indexDestMesh, dpNode );
	rageErrorCheckGeneral(stat, "Unable to get the dagPath for the destination mesh!");

	//Check to see that they have a transform node selected.
	stat = isApiType(dpNode, MFn::kTransform);
	rageErrorCheckDagPath(stat, dpNode, "The destination (last selected) polygonal mesh needs to have its transform node selected!");

	//Check to see how many shape nodes are below this transform.
	UINT numShapes;
	stat = dpNode.numberOfShapesDirectlyBelow(numShapes);
	rageErrorCheckGeneral(stat, "Unable to get the shape node below the destination mesh transform node!");

	if(numShapes != 1)
	{
		MStatus status;
		MString fullPathName;
		fullPathName = dpNode.fullPathName(&status);
		MString errorMessage ("There must be one and only one shape node (mesh) beneath the destination polygonal transform node!  `");
		errorMessage += fullPathName;
		errorMessage += "` is not valid...";
		MGlobal::displayError(errorMessage);
		return MS::kFailure;
	}

	//Attempt to get the shape node of the transform node.
	stat = dpNode.extendToShapeDirectlyBelow(0);
	rageErrorCheckDagPath(stat, dpNode, "Unable to get the shape node directly below the selected destination transform node!");

	//Check to see what type of object we have selected.
	stat = isApiType(dpNode, MFn::kMesh);
	rageErrorCheckDagPath(stat, dpNode, "The destination (last selected) node must have a polygonal mesh (shape) node!");

	//Set up the destination mesh object for intersection testing.
	dpDestMesh = dpNode;
	MFnMesh fnDestMesh(dpDestMesh, &stat);
	rageErrorCheckDagPath(stat, dpDestMesh, "Unable to initialize the MFnMesh function set!");

	//Remove the destination object from the list of items that need to be iterated.
	stat = selectionList.remove( indexDestMesh );
	rageErrorCheckGeneral(stat, "Unable to remove the destination mesh transform node from the selection list!");

	MItSelectionList iterSel( selectionList );

	//Check to see if the target nodes intersect the destination mesh and save their original and destination transforms if they do.
	UINT i = 0;//, iIsTolerant = 0;;
	for ( ; !iterSel.isDone(); iterSel.next() )
	{								
		MVector	vOrigPos, vNewPos;
		
		iterSel.getDagPath( dpNode );

		//Check to see that we have transform nodes selected.
		stat = isApiType(dpNode, MFn::kTransform);

		if (! stat)  //Was not a transform node.
		{
			MStatus status;
			MString fullPathName;
			fullPathName = dpNode.fullPathName(&status);
			MString warningMessage ("You must select the transform nodes of your target objects!  `");
			warningMessage += fullPathName;
			warningMessage += "` is not a transform node, so skipping it...";
			MGlobal::displayWarning(warningMessage);
		}
		else
		{
			//Get the position of the node in world space.
			MFnTransform fnTransform( dpNode, &stat );
			rageErrorCheckDagPath(stat, dpNode, "Unable to initialize the MFnTransform function set!");

			vOrigPos = fnTransform.translation( MSpace::kWorld, &stat );

			if ( MS::kFailure == stat )
			{
				MStringArray strArray;
				stat = iterSel.getStrings(strArray);
				MString errorMessage ("Could not get the translation on  `");
				errorMessage += strArray[0];
				errorMessage += "`!";
				MGlobal::displayError(errorMessage);
				return MS::kFailure;
			}

			float fRayPosX = (float)vOrigPos.x;
			float fRayPosY = (float)vOrigPos.y;
			float fRayPosZ = (float)vOrigPos.z;

			//Check to see if we get an intersection.
			MFloatPoint raySource(fRayPosX, fRayPosY, fRayPosZ, 1);
			MFloatPoint hitPoint;
			MFloatVector  rayDirection;

			rayDirection.x = 0.0;
			rayDirection.y = 0.0;
			rayDirection.z = 0.0;

			if(axisLetter == X_AXIS)
			{
				rayDirection.x = -1.0;
			}
			else if(axisLetter == Y_AXIS)
			{
				rayDirection.y = -1.0;
			}
			else if(axisLetter == Z_AXIS)
			{
				rayDirection.z = -1.0;
			}

            bool bGotHit = false;
			getIntersection(raySource, rayDirection, dpDestMesh, hitPoint, bGotHit);

			//iIsTolerant = isTolerant(raySource, hitPoint, axisLetter, f_TOLERANCE);

			//If we did not get a hit, shoot the ray in the opposite direction.
			if (! bGotHit)//if(! iIsTolerant)
			{
				if(axisLetter == X_AXIS)
				{
					rayDirection.x = -rayDirection.x;
				}
				else if(axisLetter == Y_AXIS)
				{
					rayDirection.y = -rayDirection.y;
				}
				else if(axisLetter == Z_AXIS)
				{
					rayDirection.z = -rayDirection.z;
				}

				getIntersection(raySource, rayDirection, dpDestMesh, hitPoint, bGotHit);

				//iIsTolerant = isTolerant(raySource, hitPoint, axisLetter, f_TOLERANCE);
			}

			//If we got a hit, store the destination position for the node and add it to an edit list.
			if(bGotHit) //if(iIsTolerant)
			{
				if(axisLetter == X_AXIS)
				{
					vNewPos.x = hitPoint.x;
					vNewPos.y = vOrigPos.y;
					vNewPos.z = vOrigPos.z;
				}
				else if(axisLetter == Y_AXIS)
				{
					vNewPos.x = vOrigPos.x;
					vNewPos.y = hitPoint.y;
					vNewPos.z = vOrigPos.z;
				}
				else if(axisLetter == Z_AXIS)
				{
					vNewPos.x = vOrigPos.x;
					vNewPos.y = vOrigPos.y;
					vNewPos.z = hitPoint.z;
				}

				stat = vaOrigPos.append(vOrigPos);
				rageErrorCheckGeneral(stat, "Unable to append the original position to the vector array!");

				stat = vaNewPos.append(vNewPos);
				rageErrorCheckGeneral(stat, "Unable to append the new position to the new position vector array!");

				stat = editList.add( dpNode );
				rageErrorCheckGeneral(stat, "Unable to remove the non-intersecting node from the selection list!");
			}
			else //There was  no hit.
			{
				/*TODO: Temporarily commented out for speed purposes.  Works just fine.
				MStringArray strArray;
				stat = iterSel.getStrings(strArray);
				MString message ("The node `");
				message += strArray[0];
				message += ("` does not intersect your destination mesh along your designated axis, so skipping it...");
				MGlobal::displayInfo(message);
				*/
			}
		}
		
		i++;
	}

	// Typically, the doIt() method only collects the infomation required
	// to do/undo the action and then stores it in class members.  The 
	// redo method is then called to do the actuall work.  This prevents
	// code duplication.
	//
	return redoIt();
}

MStatus rageSnapNodes::redoIt()
//
//	Description:
//		implements redo for the MEL rageSnapNodes command. 
//
//		This method is called when the user has undone a command of this type
//		and then redoes it.  No arguments are passed in as all of the necessary
//		information is cached by the doIt method.
//
//	Return Value:
//		MS::kSuccess - command succeeded
//		MS::kFailure - redoIt failed.  this is a serious problem that will
//                     likely cause the undo queue to be purged
//
{
	// Since this class is derived off of MPxCommand, you can use the 
	// inherited methods to return values and set error messages
	//
	
	MDagPath	dpDoItNode;
	MStatus		stat;
	
	MItSelectionList iterDoIt( editList );

	//UINT numSelectedItems = editList.length();
	
	UINT i = 0;
	for ( ; !iterDoIt.isDone(); iterDoIt.next() )
	{
		iterDoIt.getDagPath( dpDoItNode );

		MFnTransform fnTransform( dpDoItNode, &stat );
		rageErrorCheckDagPath(stat, dpDoItNode, "Unable to initialize the MFnTransform function set in doIt!");
			
		stat = fnTransform.setTranslation( vaNewPos[i], MSpace::kWorld );

		if ( MS::kSuccess != stat )
		{
			MStringArray strArray;
			stat = iterDoIt.getStrings(strArray);
			MString errorMessage ("Could not translate the node `");
			errorMessage += strArray[0];
			errorMessage += "`!";
			MGlobal::displayError(errorMessage);
			return MS::kFailure;
		}
		
		i++;
	}
	
	setResult( "rageSnapNodes command completed!  See Script Editor for details...\n" );

	return MS::kSuccess;
}

MStatus rageSnapNodes::undoIt()
//
//	Description:
//		implements undo for the MEL rageSnapNodes command.  
//
//		This method is called to undo a previous command of this type.  The 
//		system should be returned to the exact state that it was it previous 
//		to this command being executed.  That includes the selection state.
//
//	Return Value:
//		MS::kSuccess - command succeeded
//		MS::kFailure - redoIt failed.  this is a serious problem that will
//                     likely cause the undo queue to be purged
//
{
	MDagPath	dpUndoItNode;
	MStatus		stat;

	MItSelectionList iterUndoIt( editList );

	//UINT numSelectedItems = editList.length();

	UINT i = 0;
	for ( ; !iterUndoIt.isDone(); iterUndoIt.next() )
	{
		iterUndoIt.getDagPath( dpUndoItNode );

		MFnTransform fnTransform( dpUndoItNode, &stat );
		rageErrorCheckDagPath(stat, dpUndoItNode, "Unable to initialize the MFnTransform function set in undoIt!");

		stat = fnTransform.setTranslation( vaOrigPos[i], MSpace::kWorld );

		if ( MS::kSuccess != stat )
		{
			MStringArray strArray;
			stat = iterUndoIt.getStrings(strArray);
			MString errorMessage ("Could not undo translating the node `");
			errorMessage += strArray[0];
			errorMessage += "`!";
			MGlobal::displayError(errorMessage);
			return MS::kFailure;
		}

		i++;
	}

    MGlobal::displayInfo( "rageSnapNodes command undone!\n" );

	return MS::kSuccess;
}

void* rageSnapNodes::creator()
//
//	Description:
//		this method exists to give Maya a way to create new objects
//      of this type. 
//
//	Return Value:
//		a new object of this type
//
{
	return new rageSnapNodes();
}

rageSnapNodes::rageSnapNodes()
//
//	Description:
//		rageSnapNodes constructor
//
{}

rageSnapNodes::~rageSnapNodes()
//
//	Description:
//		rageSnapNodes destructor
//
{
}

bool rageSnapNodes::isUndoable() const
//
//	Description:
//		this method tells Maya this command is undoable.  It is added to the 
//		undo queue if it is.
//
//	Return Value:
//		true if this command is undoable.
//
{
	return true;
}