clusterWeightFunction.cpp

//-
// ==========================================================================
// Copyright 1995,2006,2008 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 <maya/MIOStream.h>
#include <math.h>
#include <maya/MSimple.h>
#include <maya/MGlobal.h>
#include <maya/MDagPath.h>
#include <maya/MSelectionList.h>
#include <maya/MObjectArray.h>
#include <maya/MFnWeightGeometryFilter.h>
#include <maya/MSyntax.h>
#include <maya/MArgDatabase.h>
#include <maya/MItGeometry.h>
#include <maya/MPoint.h>


#define kSineFlag                                                               "-s"
#define kSineFlagLong                                                   "-sine"
#define kSineDistanceFlag                                               "-sd"
#define kSineDistanceFlagLong                                   "-sineDistance"
#define kSineDistance2Flag                                              "-sd2"
#define kSineDistance2FlagLong                                  "-sineDistance2"
#define kDistanceSineDistanceFlag                               "-dsd"
#define kDistanceSineDistanceFlagLong                   "-distanceSineDistance"
#define kInverseDistanceSineDistanceFlag                "-ids"
#define kInverseDistanceSineDistanceFlagLong    "-inverseDistanceSineDistance"
#define kDistanceFlag                                                   "-d"
#define kDistanceFlagLong                                               "-distance"
#define kDistance2Flag                                                  "-d2"
#define kDistance2FlagLong                                              "-distance2"
#define kDistance3Flag                                                  "-d3"
#define kDistance3FlagLong                                              "-distance3"
#define kDistance4Flag                                                  "-d4"
#define kDistance4FlagLong                                              "-distance4"
#define kInverseDistanceFlag                                    "-id"
#define kInverseDistanceFlagLong                                "-inverseDistance"
#define kInverseDistance2Flag                                   "-id2"
#define kInverseDistance2FlagLong                               "-inverseDistance2"
#define kInverseDistance3Flag                                   "-id3"
#define kInverseDistance3FlagLong                               "-inverseDistance3"
#define kInverseDistance4Flag                                   "-id4"
#define kInverseDistance4FlagLong                               "-inverseDistance4"


class clusterWeightFunctionCmd : public MPxCommand
{
public:
        clusterWeightFunctionCmd();
        virtual ~clusterWeightFunctionCmd();
        virtual MStatus doIt(const MArgList &args);
        static void *creator();
        static MSyntax newSyntax();
        
        enum effectType {
                kSine,
                kSineDistance,
                kSineDistance2,
                kDistanceSineDistance,
                kInverseDistanceSineDistance,
                kDistance,
                kDistance2,
                kDistance3,
                kDistance4,
                kInverseDistance,
                kInverseDistance2,
                kInverseDistance3,
                kInverseDistance4
        };

private:
        MStatus parseArgs(const MArgList &);
        MStatus performWeighting(MFnWeightGeometryFilter &,
                                                         MDagPath &,
                                                         MObject &);

        effectType fEffectType;
};


clusterWeightFunctionCmd::clusterWeightFunctionCmd()
{
}


clusterWeightFunctionCmd::~clusterWeightFunctionCmd()
{
}


MStatus clusterWeightFunctionCmd::doIt(const MArgList &argList)
{
        MStatus status;

        MSelectionList list;
        MGlobal::getActiveSelectionList(list);
        
        // Get the cluster
        MFnWeightGeometryFilter cluster;
        MObject dgNode;
        list.getDependNode(0, dgNode);
        if (!cluster.setObject(dgNode))
                return MS::kFailure;

        // Get the object
        MDagPath dagPath;
        MObject component;
        list.getDagPath(1, dagPath, component);

        status = parseArgs(argList);
        if (MS::kSuccess == status)
                status = performWeighting(cluster, dagPath, component);
        return status;
}


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


MSyntax clusterWeightFunctionCmd::newSyntax()
{
        MSyntax syntax;

        syntax.addFlag(kSineFlag, kSineFlagLong);
        syntax.addFlag(kSineDistanceFlag, kSineDistanceFlagLong);
        syntax.addFlag(kSineDistance2Flag, kSineDistance2FlagLong);
        syntax.addFlag(kDistanceSineDistanceFlag, kDistanceSineDistanceFlagLong);
        syntax.addFlag(kInverseDistanceSineDistanceFlag, 
                                   kInverseDistanceSineDistanceFlagLong);
        syntax.addFlag(kDistanceFlag, kDistanceFlagLong);
        syntax.addFlag(kDistance2Flag, kDistance2FlagLong);
        syntax.addFlag(kDistance3Flag, kDistance3FlagLong);
        syntax.addFlag(kDistance4Flag, kDistance4FlagLong);
        syntax.addFlag(kInverseDistanceFlag, kInverseDistanceFlagLong);
        syntax.addFlag(kInverseDistance2Flag, kInverseDistance2FlagLong);
        syntax.addFlag(kInverseDistance3Flag, kInverseDistance3FlagLong);
        syntax.addFlag(kInverseDistance4Flag, kInverseDistance4FlagLong);

        return syntax;
}


MStatus clusterWeightFunctionCmd::parseArgs(const MArgList &argList)
{
        MArgDatabase argData(syntax(), argList);

        if (argData.isFlagSet(kSineFlag)) {
                fEffectType = kSine;
        }
        else if (argData.isFlagSet(kSineDistanceFlag)) {
                fEffectType = kSineDistance;
        }
        else if (argData.isFlagSet(kSineDistance2Flag)) {
                fEffectType = kSineDistance2;
        }
        else if (argData.isFlagSet(kDistanceSineDistanceFlag)) {
                fEffectType = kDistanceSineDistance;
        }
        else if (argData.isFlagSet(kInverseDistanceSineDistanceFlag)) {
                fEffectType = kInverseDistanceSineDistance;
        }
        else if (argData.isFlagSet(kDistanceFlag)) {
                fEffectType = kDistance;
        }
        else if (argData.isFlagSet(kDistance2Flag)) {
                fEffectType = kDistance2;
        }
        else if (argData.isFlagSet(kDistance3Flag)) {
                fEffectType = kDistance3;
        }
        else if (argData.isFlagSet(kDistance4Flag)) {
                fEffectType = kDistance4;
        }
        else if (argData.isFlagSet(kInverseDistanceFlag)) {
                fEffectType = kInverseDistance;
        }
        else if (argData.isFlagSet(kInverseDistance2Flag)) {
                fEffectType = kInverseDistance2;
        }
        else if (argData.isFlagSet(kInverseDistance3Flag)) {
                fEffectType = kInverseDistance3;
        }
        else if (argData.isFlagSet(kInverseDistance4Flag)) {
                fEffectType = kInverseDistance4;
        }

        return MS::kSuccess;
}


MStatus clusterWeightFunctionCmd::performWeighting(MFnWeightGeometryFilter &cluster,
                                                                                   MDagPath &dagPath,
                                                                                   MObject &component)
{
        MStatus status;
        MItGeometry geomIter(dagPath, component, &status);
        MObject comp;
        MObjectArray compArray;
        MDoubleArray weightArray;
        double weight = 0.0;

        if (MS::kSuccess == status) {
                for (; !geomIter.isDone(); geomIter.next()) {
                        comp = geomIter.component();

                        MPoint pnt = geomIter.position(MSpace::kWorld, &status);

                        if (MS::kSuccess != status) {
                                return MS::kFailure;
                        }

                        if (kSine == fEffectType) {
                                weight = sin(pnt.x)*sin(pnt.z);
                        }
                        else if (kSineDistance == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = sin(distance);
                        }
                        else if (kSineDistance2 == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = sin(distance*distance);
                        }
                        else if (kDistanceSineDistance == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = distance*sin(distance);
                        }
                        else if (kInverseDistanceSineDistance == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = sin(distance)/(distance+1);
                        }
                        else if (kDistance == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = distance;
                        }
                        else if (kDistance2 == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = distance*distance;
                        }
                        else if (kDistance3 == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = distance*distance*distance;
                        }
                        else if (kDistance4 == fEffectType) {
                                double distance = pnt.distanceTo(MPoint::origin);
                                weight = distance*distance*distance*distance;
                        }
                        else if (kInverseDistance == fEffectType) {
                                double dist = pnt.distanceTo(MPoint::origin) + 1;
                                weight = 1/(dist);
                        }
                        else if (kInverseDistance2 == fEffectType) {
                                double dist = pnt.distanceTo(MPoint::origin) + 1;
                                weight = 1/(dist*dist);
                        }
                        else if (kInverseDistance3 == fEffectType) {
                                double dist = pnt.distanceTo(MPoint::origin) + 1;
                                weight = 1/(dist*dist*dist);
                        }
                        else if (kInverseDistance4 == fEffectType) {
                                double dist = pnt.distanceTo(MPoint::origin) + 1;
                                weight = 1/(dist*dist*dist*dist);
                        }

                        compArray.append(comp);
                        weightArray.append(weight);
                }

                unsigned length = compArray.length();
                for (unsigned i = 0; i < length; i++) {
                        cluster.setWeight(dagPath, compArray[i], (float) weightArray[i]);
                }

                return MS::kSuccess;
        } else
                return MS::kFailure;
}


extern MStatus initializePlugin(MObject obj)
{
        MStatus status;
        MFnPlugin plugin(obj, PLUGIN_COMPANY, "3.0", "Any");
        
        status = plugin.registerCommand("clusterWeightFunction",
                                                                        clusterWeightFunctionCmd::creator,
                                                                        clusterWeightFunctionCmd::newSyntax);
        
        if (!status) {
                status.perror("registerCommand");
        }

        return status;
}


extern MStatus uninitializePlugin(MObject obj)
{
        MStatus status;
        MFnPlugin plugin(obj);
        
        status = plugin.deregisterCommand("clusterWeightFunction");

        if (!status) {
                status.perror("deregisterCommand");
        }       

        return status;
}