cvColorNode.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/MPxLocatorNode.h> 
#include <maya/MString.h> 
#include <maya/MTypeId.h> 
#include <maya/MPlug.h>
#include <maya/MMatrix.h>
#include <maya/MPoint.h>
#include <maya/MPointArray.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MColor.h>
#include <maya/M3dView.h>
#include <maya/MFnPlugin.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnTypedAttribute.h>
#include <maya/MFnNurbsSurface.h>
#include <maya/MFnNurbsSurfaceData.h>
#include <maya/MFnPointArrayData.h>
#include <maya/MFnMatrixData.h>
#include <maya/MIOStream.h>

class cvColor : public MPxLocatorNode
{
public:
                                                        cvColor();
        virtual                                 ~cvColor(); 

    virtual MStatus             compute( const MPlug& plug, MDataBlock& data );

        virtual void            draw( M3dView & view, const MDagPath & path, 
                                                                  M3dView::DisplayStyle style,
                                                                  M3dView::DisplayStatus status );

        virtual bool            isBounded() const;

        static  void *          creator();
        static  MStatus         initialize();

        static  MTypeId                 id;
        static  MObject                 drawingEnabled;
        static  MObject                 pointSize;
        static  MObject         inputSurface;
        static  MObject                 cvLocations;
};

MTypeId cvColor::id( 0x80006 );
MObject cvColor::inputSurface;
MObject cvColor::cvLocations;
MObject cvColor::pointSize;
MObject cvColor::drawingEnabled;

cvColor::cvColor() {}
cvColor::~cvColor() {}

MStatus cvColor::compute( const MPlug& plug, MDataBlock& data )
{ 
        MStatus stat;

        // cout << "cvColor::compute\n";

        if ( plug == cvLocations ) {
                MDataHandle inputData = data.inputValue ( inputSurface, &stat );
                if (!stat) {
                        stat.perror("cvColor::compute get inputSurface");
                        return stat;
                }

                MObject surf = inputData.asNurbsSurface();
                MFnNurbsSurface surfFn (surf, &stat);
                if (!stat) {
                        stat.perror("cvColor::compute surface creator");
                        return stat;
                }

                MDataHandle outputData = data.outputValue ( cvLocations, &stat );
                if (!stat) {
                        stat.perror("cvColor::compute get cvLocations");
                        return stat;
                }
                
                MObject cvs = outputData.data();
                MFnPointArrayData cvData(cvs, &stat);
                if (!stat) {
                        stat.perror("cvColor::compute point array data creator");
                        return stat;
                }

                MPointArray cvArray;
                stat = surfFn.getCVs( cvArray, MSpace::kObject);
                if (!stat) {
                        stat.perror("cvColor::compute getCVs");
                        return stat;
                }

                stat = cvData.set ( cvArray );
                if (!stat) {
                        stat.perror("cvColor::compute setCVs");
                        return stat;
                }

                outputData.set ( cvs );

                stat = data.setClean ( plug );
                if (!stat) {
                        stat.perror("cvColor::compute setClean");
                        return stat;
                }
        } else {
                return MS::kUnknownParameter;
        }

        return MS::kSuccess;
}

void cvColor::draw( M3dView & view, const MDagPath & path, 
                                                         M3dView::DisplayStyle style,
                                                         M3dView::DisplayStatus status )
{ 
        // cout << "cvColor::draw\n";

        MStatus         stat;
        MObject         thisNode = thisMObject();

        MPlug enPlug( thisNode, drawingEnabled );
        bool doDrawing; 
        stat = enPlug.getValue ( doDrawing );
        if (!stat) {
                stat.perror("cvColor::draw get drawingEnabled");
                return;
        }

        if (!doDrawing)
                return;

        MPlug szPlug( thisNode, pointSize );
        float ptSize; 
        stat = szPlug.getValue ( ptSize );
        if (!stat) {
                stat.perror("cvColor::draw get pointSize");
                ptSize = 4.0;
        }

        MPlug cvPlug( thisNode, cvLocations );
        MObject cvObject;
        stat = cvPlug.getValue(cvObject);
        if (!stat) {
                stat.perror("cvColor::draw get cvObject");
                return;
        }

        MFnPointArrayData cvData(cvObject, &stat);
        if (!stat) {
                stat.perror("cvColor::draw get point array data");
                return;
        }

        MPointArray cvs = cvData.array( &stat );
        if (!stat) {
                stat.perror("cvColor::draw get point array");
                return;
        }

        // Extract the 'worldMatrix' attribute that is inherited from 'dagNode'
        //
        MFnDependencyNode fnThisNode( thisNode );
        MObject worldSpaceAttribute = fnThisNode.attribute( "worldMatrix" );
        MPlug matrixPlug( thisNode, worldSpaceAttribute);

        // 'worldMatrix' is an array so we must specify which element the plug
        // refers to
        matrixPlug = matrixPlug.elementByLogicalIndex (0);

        // Get the value of the 'worldMatrix' attribute
        //
        MObject matObject;
        stat = matrixPlug.getValue(matObject);
        if (!stat) {
                stat.perror("cvColor::draw get matObject");
                return;
        }

        MFnMatrixData matrixData(matObject, &stat);
        if (!stat) {
                stat.perror("cvColor::draw get world matrix data");
                return;
        }

        MMatrix worldSpace = matrixData.matrix( &stat );
        if (!stat) {
                stat.perror("cvColor::draw get world matrix");
                return;
        }

        if ( view.isColorIndexMode() ) {
                cerr << "Can't update cv colors in color index mode\n";
                return;
        }

        view.beginGL(); 

        // Push the color settings
        // 
        glPushAttrib( GL_CURRENT_BIT | GL_POINT_BIT );
        glPointSize( ptSize );
        glDisable ( GL_POINT_SMOOTH );  // Draw square "points"

        glBegin( GL_POINTS );

                int numCVs = cvs.length();
                for (int i = 0; i < numCVs; i++) {
                        // cout << "cv[" << i << "]: " << cvs[i] << ": ";
                        MPoint          cv( cvs[i] );
                        MColor          cvColor;

                        cv *= worldSpace;

                        if (cv.x < 0 && cv.y < 0) {
                                // cout << "Red";
                                cvColor.r = 1.0;
                                cvColor.g = 0.0;
                                cvColor.b = 0.0;
                        } else if (cv.x < 0 && cv.y >= 0) {
                                // cout << "Cyan";
                                cvColor.r = 0.0;
                                cvColor.g = 1.0;
                                cvColor.b = 1.0;
                        } else if (cv.x >= 0 && cv.y < 0) {
                                // cout << "Blue";
                                cvColor.r = 0.0;
                                cvColor.g = 0.0;
                                cvColor.b = 1.0;
                        } else if (cv.x >= 0 && cv.y >= 0) {
                                // cout << "Yellow";
                                cvColor.r = 1.0;
                                cvColor.g = 1.0;
                                cvColor.b = 0.0;
                        }
                        // cout << endl;

                        view.setDrawColor ( cvColor );
                        glVertex3f( (float)cvs[i].x, (float)cvs[i].y, (float)cvs[i].z);
                } 
        glEnd();

        glPopAttrib();

        view.endGL();
}

bool cvColor::isBounded() const
{ 
        return false;
}

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

MStatus cvColor::initialize()
{ 
        MStatus                         stat;
        MFnNumericAttribute     numericAttr;
        MFnTypedAttribute       typedAttr;

        drawingEnabled = numericAttr.create( "drawingEnabled", "en",
                                                                        MFnNumericData::kBoolean, 1, &stat );
        if (!stat) {
                stat.perror("create drawingEnabled attribute");
                return stat;
        }

        pointSize = numericAttr.create( "pointSize", "ps",
                                                                        MFnNumericData::kFloat, 4.0, &stat );
        if (!stat) {
                stat.perror("create pointSize attribute");
                return stat;
        }
        
        inputSurface = typedAttr.create( "inputSurface", "is",
                                                         MFnNurbsSurfaceData::kNurbsSurface, &stat);
        if (!stat) {
                stat.perror("create inputSurface attribute");
                return stat;
        }

        cvLocations = typedAttr.create( "cvLocations", "cv",
                                                         MFnPointArrayData::kPointArray, &stat);
        if (!stat) {
                stat.perror("create cvLocations attribute");
                return stat;
        }

        MPointArray                     defaultPoints;
        MFnPointArrayData       defaultArray;
        MObject                         defaultAttr;

        defaultPoints.clear(); // Empty array
        defaultAttr = defaultArray.create (defaultPoints);
        stat = typedAttr.setDefault(defaultAttr);
        if (!stat) {
                stat.perror("could not create default output attribute");
                return stat;
        }

        stat = addAttribute (drawingEnabled);
                if (!stat) { stat.perror("addAttribute"); return stat;}
        stat = addAttribute (pointSize);
                if (!stat) { stat.perror("addAttribute"); return stat;}
        stat = addAttribute (inputSurface);
                if (!stat) { stat.perror("addAttribute"); return stat;}
        stat = addAttribute (cvLocations);
                if (!stat) { stat.perror("addAttribute"); return stat;}

        stat = attributeAffects( inputSurface, cvLocations );
                if (!stat) { stat.perror("attributeAffects"); return stat;}
        stat = attributeAffects( drawingEnabled, cvLocations );
                if (!stat) { stat.perror("attributeAffects"); return stat;}
        stat = attributeAffects( pointSize, cvLocations );
                if (!stat) { stat.perror("attributeAffects"); return stat;}
        
        return MS::kSuccess;
}

MStatus initializePlugin( MObject obj )
{ 
        MStatus   status;
        MFnPlugin plugin( obj, PLUGIN_COMPANY, "3.0", "Any");

        status = plugin.registerNode( "cvColor", cvColor::id, 
                                                 &cvColor::creator, &cvColor::initialize,
                                                 MPxNode::kLocatorNode );
        if (!status) {
                status.perror("registerNode");
                return status;
        }
        return status;
}

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

        status = plugin.deregisterNode( cvColor::id );
        if (!status) {
                status.perror("deregisterNode");
                return status;
        }

        return status;
}