skinClusterWeights/skinClusterWeights.cpp

//
// skinClusterWeights.cpp
//
// see examples.mel in the same directory for examples
//
// skinClusterWeights -inf {"obj1", "obj2"} -sc {"skinCluster1", "skinCluster2"} {"myShape"};
// skinClusterWeights -edit -w {1, 2, 3, 4} "myShape.vtx[33]";
//
// The generic syntax is 
// skinClusterWeights -q/-edit -inf/influences $influenceArray -sc/skinClusters $skinClusterArray -w/weights $weightFloatArray $objectStringArray;
//

#include <maya/MPxCommand.h>
#include <maya/MArgList.h>
#include <maya/MArgParser.h>

#include <maya/MDagPath.h>
#include <maya/MItSelectionList.h>
#include <maya/MSelectionList.h>

#include <maya/MFnPlugin.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MPlugArray.h>
#include <maya/MPlug.h>
#include <maya/MItDependencyGraph.h>

#include <maya/MString.h>
#include <maya/MObjectArray.h>
#include <maya/MDoubleArray.h>
#include <maya/MDagPathArray.h>

#include <maya/MGlobal.h>

#include <maya/MFnSkinCluster.h>
#include <maya/MItMeshVertex.h>
#include <maya/MItSurfaceCV.h>
#include <maya/MItCurveCV.h>

#include <math.h>

#define kEditFlag                   "-e"
#define kEditFlagLong               "-edit"
#define kQueryFlag                  "-q"
#define kQueryFlagLong              "-query"
#define kInfluenceFlag              "-inf"
#define kInfluenceFlagLong          "-influences"
#define kSkinClusterFlag            "-sc"
#define kSkinClusterFlagLong        "-skinClusters"
#define kWeightFlag                 "-w"
#define kWeightFlagLong             "-weights"
#define kAssignAllToSingleFlag      "-as"
#define kAssignAllToSingleFlagLong  "-assignAllToSingle"

//
// Command class declaration
//
class skinClusterWeights : public MPxCommand
{
public:
                         skinClusterWeights();
        virtual     ~skinClusterWeights(); 

        virtual MStatus doIt(const MArgList& args);
        virtual MStatus undoIt();
        virtual MStatus redoIt();
        virtual bool    isUndoable() const {return true;}

    static void*    creator();

private:

    void        doItEdit();
    void        doItQuery();
    bool        editUsed;
    bool        queryUsed;

        // Return the skinCluster given a geometry
    MObject     findSkinCluster(MDagPath&);

        void            populateInfluenceIndexArray(MFnSkinCluster &skinClusterFn, MIntArray &influenceIndexArray);

        // true if 1) skin cluster is not specified in the command or
        //         2) MObject is one of the skin clusters specified in the command
    bool        isSkinClusterIncluded(MObject&);

    MStatus     parseArgs( const MArgList &args );
    MDagPathArray   influenceArray;
    MObjectArray    skinClusterArray;
    MDoubleArray    weightArray;
    MStringArray    geometryArray;
    bool        assignAllToSingle;

        // data structures for undo
        MDagPathArray   fDagPathArray;
        MObjectArray    fComponentArray;
    MIntArray      *fInfluenceIndexArrayPtrArray;
    MDoubleArray   *fWeightsPtrArray;
 };

skinClusterWeights::skinClusterWeights() : assignAllToSingle(false) 
{
    fInfluenceIndexArrayPtrArray = NULL;
    fWeightsPtrArray = NULL;
}

skinClusterWeights::~skinClusterWeights() 
{
    if (fInfluenceIndexArrayPtrArray != NULL) delete [] fInfluenceIndexArrayPtrArray;
    if (fWeightsPtrArray != NULL) delete [] fWeightsPtrArray;
}

void* skinClusterWeights::creator()
{
    return new skinClusterWeights;
}

// Custom parsing to support array argument. Not using MSyntax.
MStatus skinClusterWeights::parseArgs( const MArgList& args )
{
    MStatus     status = MS::kSuccess;
    MObject         node;
    MDagPath        dagPath;
    MSelectionList  selList;
    
    editUsed  = true;
    queryUsed = false;
    unsigned int i, nth = 0;
    unsigned int numArgs = args.length();   
    while(status == MS::kSuccess && nth < numArgs-1) {
    MString inputString = args.asString(nth, &status);

    if (status != MS::kSuccess) {
        MGlobal::displayError("skinClusterWeights syntax error");
        return status;
    }
      
    if (inputString == kEditFlag || inputString == kEditFlagLong) {
        editUsed  = true;
        queryUsed = false;
        nth++;
        continue;
    }
      
    if (inputString == kQueryFlag || inputString == kQueryFlagLong) {
        queryUsed = true;
        editUsed  = false;
        nth++;
        continue;
    }
      
    if (inputString == kInfluenceFlag || inputString == kInfluenceFlagLong) {       
        nth++;
        MStringArray stringArray = args.asStringArray(nth, &status);
        selList.clear();
        for (i = 0; i < stringArray.length(); i++) {
          selList.add(stringArray[i]);
        }
        for (i = 0; i < selList.length(); i++) {
          status = selList.getDagPath(i, dagPath);
          if (status == MS::kSuccess && dagPath.hasFn(MFn::kTransform)) {
        influenceArray.append(dagPath);
          } else {
        MGlobal::displayError(inputString + " is not a valid influence object.\n");
        return status;
          }
        }
        nth++;
        continue;
    }

    if (inputString == kSkinClusterFlag || inputString == kSkinClusterFlagLong) {
        nth++;
        MStringArray stringArray = args.asStringArray(nth, &status);
        selList.clear();
        for (i = 0; i < stringArray.length(); i++) {
          selList.add(stringArray[i]);
        }
        for (i = 0; i < selList.length(); i++) {
          status = selList.getDependNode(i, node);
          if (status == MS::kSuccess && node.hasFn(MFn::kSkinClusterFilter)) {
        skinClusterArray.append(node);
          } else {
        MGlobal::displayError(inputString + " is not a valid skinCluster.\n");
        return status;
          }
        }
        nth++;
        continue;
    }

    if (inputString == kWeightFlag || inputString == kWeightFlagLong) {
        nth++;
        weightArray = args.asDoubleArray(nth, &status);
        if (status != MS::kSuccess) {
        MGlobal::displayError("error while parsing weight array");
        }
        nth++;
        continue;
    }

    if (inputString == kAssignAllToSingleFlag || inputString == kAssignAllToSingleFlagLong) {
        assignAllToSingle = true;
        nth++;
        continue;
    }

    MGlobal::displayError("invalid command syntax at " + inputString);
    return MS::kFailure;
    }

    // parse command objects
    // nth should equals to numArgs-1 at this point
    geometryArray = args.asStringArray(nth, &status);
    if (status != MS::kSuccess) {
    MGlobal::displayError("Command object invalid");
    return status;
    }

    if (queryUsed) {
    if (assignAllToSingle) {
        MGlobal::displayWarning("-as/-assignAllToSingle is ignored with query flag");
    }
    if (weightArray.length() > 0) {
        MGlobal::displayWarning("-w/-weights is ignored with query flag");
    }
    }
    
    return status;
}

MStatus skinClusterWeights::doIt( const MArgList& args )
//
// Description
//
{
    MStatus status = parseArgs(args);
    if (status != MS::kSuccess) return status;

    if (editUsed) {
    redoIt();
    } else if (queryUsed) {
       doItQuery();
    }
    return MS::kSuccess;
}

MStatus skinClusterWeights::redoIt()
{
    MStatus status;
    unsigned int ptr = 0;
    
    MSelectionList selList;
    int geomLen = geometryArray.length();
    fDagPathArray.setLength(geomLen);
    fComponentArray.setLength(geomLen);
    fInfluenceIndexArrayPtrArray = new MIntArray[geomLen];
    fWeightsPtrArray = new MDoubleArray[geomLen];

    for (int i = 0; i < geomLen; i++) {
    MDagPath dagPath;
    MObject  component;
    
    selList.clear();
    selList.add(geometryArray[i]);
    MStatus status = selList.getDagPath(0, dagPath, component);
    if (status != MS::kSuccess) {
        continue;
    }
    if (component.isNull()) dagPath.extendToShape();

    MObject skinCluster = findSkinCluster(dagPath);
    if (!isSkinClusterIncluded(skinCluster)) {
        continue;
    }

    MFnSkinCluster skinClusterFn(skinCluster, &status);
    if (status != MS::kSuccess) {
        continue; 
    }

    MIntArray influenceIndexArray;
    populateInfluenceIndexArray(skinClusterFn, influenceIndexArray);

    unsigned numInf = influenceIndexArray.length();
    if (numInf == 0) continue;
      
    unsigned numCV = 0;
    if (dagPath.node().hasFn(MFn::kMesh)) {
        MItMeshVertex polyIter(dagPath, component, &status);
        if (status == MS::kSuccess) {
        numCV = polyIter.count();
        } 
    } else if (dagPath.node().hasFn(MFn::kNurbsSurface)) {
        MItSurfaceCV nurbsIter(dagPath, component, true, &status);
        if (status == MS::kSuccess) {
        while (!nurbsIter.isDone()) {
            numCV++;
            nurbsIter.next();
        }
        }
    } else if (dagPath.node().hasFn(MFn::kNurbsCurve)) {
        MItCurveCV curveIter(dagPath, component, &status);
        if (status == MS::kSuccess) {
        while (!curveIter.isDone()) {
            numCV++;
            curveIter.next();
        }
        }
    }
    unsigned numEntry = numCV * numInf;

    if (numEntry > 0) {
        MDoubleArray weights(numEntry);
        unsigned int numWeights = weightArray.length();
        if (assignAllToSingle) {
        if (numInf <= numWeights) {
            for (unsigned j = 0; j < numEntry; j++) {
            weights[j] = weightArray[j % numInf];
            }
        } else {
            MGlobal::displayError("Not enough weights specified\n");
            return MS::kFailure;
        }
        } else {
        for (unsigned j = 0; j < numEntry; j++, ptr++) {
            if (ptr < numWeights) {
            weights[j] = weightArray[ptr];
            } else {
            MGlobal::displayError("Not enough weights specified\n");
            return MS::kFailure;
            }
        }
        }

        // support for undo
        fDagPathArray[i] = dagPath;
        fComponentArray[i] = component;
        fInfluenceIndexArrayPtrArray[i] = influenceIndexArray;
        MDoubleArray oldWeights;
        skinClusterFn.getWeights(dagPath, component, influenceIndexArray, oldWeights);
        fWeightsPtrArray[i] = oldWeights;
        
        skinClusterFn.setWeights(dagPath, component, influenceIndexArray, weights);
    }
    }
    return MS::kSuccess;
}

MStatus skinClusterWeights::undoIt()
{
    MStatus status;
    for (unsigned i = 0; i < fDagPathArray.length(); i++) {
        MDagPath &dagPath = fDagPathArray[i];
    MObject skinCluster = findSkinCluster(dagPath);
    if (!isSkinClusterIncluded(skinCluster)) {
        continue;
    }

    MFnSkinCluster skinClusterFn(skinCluster, &status);
    if (status != MS::kSuccess) {
        continue; 
    }

    MObject &component = fComponentArray[i];
    if (dagPath.isValid() && 
        fInfluenceIndexArrayPtrArray[i].length() > 0 && fWeightsPtrArray[i].length() > 0) {

        skinClusterFn.setWeights(dagPath, component, fInfluenceIndexArrayPtrArray[i], fWeightsPtrArray[i]);
    }
    }
    fDagPathArray.clear();
    fComponentArray.clear();
    delete [] fInfluenceIndexArrayPtrArray;
    delete [] fWeightsPtrArray;
    fInfluenceIndexArrayPtrArray = NULL;
    fWeightsPtrArray = NULL;
    return MS::kSuccess;
}

void skinClusterWeights::doItQuery()
{
    MStatus status;
    unsigned int i, j;

    MDoubleArray weights;
    // To allow "skinClusterWeights -q" to return empty double array
    setResult(weights);
    MSelectionList selList;
    for (i = 0; i < geometryArray.length(); i++) {
    MDagPath dagPath;
    MObject  component;
      
    selList.clear();
    selList.add(geometryArray[i]);
    MStatus status = selList.getDagPath(0, dagPath, component);
    if (status != MS::kSuccess) {
        continue;
    }
    if (component.isNull()) dagPath.extendToShape();
      
    MObject skinCluster = findSkinCluster(dagPath);
    if (!isSkinClusterIncluded(skinCluster)) {
        continue;
    }  

    MFnSkinCluster skinClusterFn(skinCluster, &status);
    if (status != MS::kSuccess) {
        continue; 
    }

    MIntArray influenceIndexArray;
    populateInfluenceIndexArray(skinClusterFn, influenceIndexArray);
      
    weights.clear();
    skinClusterFn.getWeights(dagPath, component, influenceIndexArray, weights);

    if (weights.length() > 0) {
        for (j = 0; j < weights.length(); j++) {
        appendToResult(weights[j]);
        }
    }
    }
}

MObject skinClusterWeights::findSkinCluster(MDagPath& dagPath)
{
    MStatus status;
    MObject skinCluster;
    MObject geomNode = dagPath.node();
    MItDependencyGraph dgIt(geomNode, MFn::kSkinClusterFilter, MItDependencyGraph::kUpstream);
    if (!dgIt.isDone()) {
    skinCluster = dgIt.currentItem();
    }
    return skinCluster;
}

bool skinClusterWeights::isSkinClusterIncluded(MObject &node)
{
    MStatus   status;
    unsigned int i;

    if (skinClusterArray.length() == 0) return true;

    for (i = 0; i < skinClusterArray.length(); i++) {
    if (skinClusterArray[i] == node) return true;
    }

    return false;
}

void skinClusterWeights::populateInfluenceIndexArray(MFnSkinCluster &skinClusterFn, MIntArray &influenceIndexArray)
{
   MStatus status;

   MIntArray  allIndexArray;
   MDagPathArray pathArray;
   skinClusterFn.influenceObjects(pathArray, &status);
   for (unsigned j = 0; j < pathArray.length(); j++) {
    allIndexArray.append(skinClusterFn.indexForInfluenceObject(pathArray[j]));
   }

   if (influenceArray.length() > 0) {
        // Add the influence indices for the influence objects specified in the cmd
        for (unsigned j = 0; j < influenceArray.length(); j++) {
        unsigned int index = skinClusterFn.indexForInfluenceObject(influenceArray[j], &status);
        for (unsigned k = 0; k < allIndexArray.length(); k++) {
            if ((int)index == allIndexArray[k]) {
            influenceIndexArray.append(k);
        }
        }
    }
    } else {
        // Add the influence indices for all the influence objects of the skinCluster
    for (unsigned j = 0; j < allIndexArray.length(); j++) {
        influenceIndexArray.append(j);
    }
    }
}

//
// The following routines are used to register/unregister
// the command we are creating within Maya
//
MStatus initializePlugin( MObject obj )
{
    MStatus   status;
    MFnPlugin plugin(obj);

    status = plugin.registerCommand( "skinClusterWeights",
                     skinClusterWeights::creator);
    if (!status) {
    status.perror("registerCommand");
    return status;
    }

    return status;
}

MStatus uninitializePlugin( MObject obj)
{
    MStatus   status;
    MFnPlugin plugin( obj );

    status = plugin.deregisterCommand( "skinClusterWeights" );
    if (!status) {
    status.perror("deregisterCommand");
    return status;
    }
    
    return status;
}