//-
// ==========================================================================
// 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.
// ==========================================================================
//+

//
// Lasso selection within a user defined context.
//

#include <maya/MIOStream.h>
#include <math.h>
#include <stdlib.h>

#include <maya/MFnPlugin.h>
#include <maya/MString.h>
#include <maya/MGlobal.h>
#include <maya/M3dView.h>
#include <maya/MDagPath.h>
 
#include <maya/MItSelectionList.h>

#include <maya/MPxContextCommand.h>
#include <maya/MPxContext.h>
#include <maya/MCursor.h>
#include <maya/MEvent.h>

#if defined(OSMac_MachO_)
#include <OpenGL/gl.h>
#include <OpenGL/glu.h>
#else
#include <GL/gl.h>
#include <GL/glu.h>
#endif

#include <maya/MPoint.h>
#include <maya/MPointArray.h>
#include <maya/MItCurveCV.h>
#include <maya/MItSurfaceCV.h>
#include <maya/MItMeshVertex.h>
#include <maya/MItMeshEdge.h>
#include <maya/MItMeshPolygon.h>

#ifdef _WIN32
LPCSTR lassoToolCursor = "lassoToolCursor.cur";
#else
#include <lassoToolCursor.h>
#include <lassoToolCursorMask.h>

#define lassoToolCursor_x_hot 1
#define lassoToolCursor_y_hot 16
#endif

class coord {
public:
	short h;
	short v;
};

extern "C" int xycompare( coord *p1, coord *p2 );
int xycompare( coord *p1, coord *p2 )
{
	if ( p1->v > p2->v )
		return 1;
	if ( p2->v > p1->v )
		return -1;
	if ( p1->h > p2->h )
		return 1;
	if ( p2->h > p1->h )
		return -1;

	return 0;
}

//////////////////////////////////////////////
// The user Context
//////////////////////////////////////////////

const int initialSize		= 1024;
const int increment			=  256;
const char helpString[]		= "drag mouse to select points by encircling";

class lassoTool : public MPxContext
{
public:
					lassoTool();
	virtual			~lassoTool();
	void*			creator();

	virtual void	toolOnSetup( MEvent & event );
	virtual MStatus	doPress( MEvent & event );
	virtual MStatus	doDrag( MEvent & event );
	virtual MStatus	doRelease( MEvent & event );

private:
	void					append_lasso( short x, short y );
	void					draw_lasso();
	bool					point_in_lasso( coord pt );

	bool					firstDraw;
	coord					min;
	coord					max;
	unsigned				maxSize;
	unsigned				num_points;
	coord*					lasso;
	MGlobal::ListAdjustment	listAdjustment;
	M3dView 				view;
	MCursor					lassoCursor;

};

lassoTool::lassoTool()
	: maxSize(0)
	, num_points(0)
	, lasso(NULL)
#ifdef _WIN32
	, lassoCursor(lassoToolCursor)
#else
	, lassoCursor(lassoToolCursor_width,
			  lassoToolCursor_height,
			  lassoToolCursor_x_hot,
			  lassoToolCursor_y_hot,
			  lassoToolCursor_bits,
			  lassoToolCursorMask_bits)
#endif
{
	setTitleString ( "Lasso Pick" );
	
	// set the initial state of the cursor
	setCursor(lassoCursor);

	// Tell the context which XPM to use so the tool can properly
	// be a candidate for the 6th position on the mini-bar.
	setImage("lassoTool.xpm", MPxContext::kImage1 );
}

lassoTool::~lassoTool() {}

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

void lassoTool::toolOnSetup ( MEvent & )
{
	setHelpString( helpString );
}

MStatus lassoTool::doPress( MEvent & event )
// Set up for overlay drawing, and remember our starting point
{
	// Figure out which modifier keys were pressed, and set up the
	// listAdjustment parameter to reflect what to do with the selected points.
	if (event.isModifierShift() || event.isModifierControl() ) {
		if ( event.isModifierShift() ) {
			if ( event.isModifierControl() ) {
				// both shift and control pressed, merge new selections
				listAdjustment = MGlobal::kAddToList;
			} else {
				// shift only, xor new selections with previous ones
				listAdjustment = MGlobal::kXORWithList;
			}
		} else if ( event.isModifierControl() ) {
			// control only, remove new selections from the previous list
			listAdjustment = MGlobal::kRemoveFromList; 
		}
	} else {
		listAdjustment = MGlobal::kReplaceList;
	}

	// Get the active 3D view.
	//
	view = M3dView::active3dView();

	// Create an array to hold the lasso points. Assume no mem failures
	maxSize = initialSize;
	lasso = (coord*) malloc (sizeof(coord) * maxSize);

	coord start;
	event.getPosition( start.h, start.v );
	num_points = 1;
	lasso[0] = min = max = start;

	firstDraw = true;

	return MS::kSuccess;
}

MStatus lassoTool::doDrag( MEvent & event )
// Add to the growing lasso
{
	view.beginXorDrawing(true, true, 1.0f, M3dView::kStippleDashed);

	if (!firstDraw) {
		//	Redraw the old lasso to clear it.
		draw_lasso();
	} else {
		firstDraw = false;
	}

	coord currentPos;
	event.getPosition( currentPos.h, currentPos.v );
	append_lasso( currentPos.h, currentPos.v );

	//	Draw the new lasso.
	draw_lasso();

	view.endXorDrawing();

	return MS::kSuccess;
}

MStatus lassoTool::doRelease( MEvent & /*event*/ )
// Selects objects within the lasso
{
	MStatus							stat;
	MSelectionList					incomingList, boundingBoxList, newList;

	if (!firstDraw) {
		// Redraw the lasso to clear it.
		view.beginXorDrawing(true, true, 1.0f, M3dView::kStippleDashed);
		draw_lasso();
		view.endXorDrawing();
	}

	// We have a non-zero sized lasso.  Close the lasso, and sort
	// all the points on it.
	append_lasso(lasso[0].h, lasso[0].v);
	qsort( &(lasso[0]), num_points, sizeof( coord  ),
		(int (*)(const void *, const void *))xycompare);

	// Save the state of the current selections.  The "selectFromSceen"
	// below will alter the active list, and we have to be able to put
	// it back.
	MGlobal::getActiveSelectionList(incomingList);

	// As a first approximation to the lasso, select all components with
	// the bounding box that just contains the lasso.
	MGlobal::selectFromScreen( min.h, min.v, max.h, max.v,
							   MGlobal::kReplaceList );

	// Get the list of selected items from within the bounding box
	// and create a iterator for them.
	MGlobal::getActiveSelectionList(boundingBoxList);

	// Restore the active selection list to what it was before we
	// the "selectFromScreen"
	MGlobal::setActiveSelectionList(incomingList, MGlobal::kReplaceList);

	// Iterate over the objects within the bounding box, extract the
	// ones that are within the lasso, and add those to newList.
	MItSelectionList iter(boundingBoxList);
	newList.clear();

	bool	foundEntireObjects = false;
	bool	foundComponents = false;

	for ( ; !iter.isDone(); iter.next() ) {
		MDagPath	dagPath;
		MObject		component;
		MPoint		point;
		coord		pt;
		MObject     singleComponent;

		iter.getDagPath( dagPath, component );

		if (component.isNull()) {
			foundEntireObjects = true;
			continue; // not a component
		}

		foundComponents = true;

		switch (component.apiType()) {
		case MFn::kCurveCVComponent:
			{
				MItCurveCV curveCVIter( dagPath, component, &stat );
				for ( ; !curveCVIter.isDone(); curveCVIter.next() ) {
					point = curveCVIter.position(MSpace::kWorld, &stat );
					view.worldToView( point, pt.h, pt.v, &stat );
					if (!stat) {
						stat.perror("Could not get position");
						continue;
					}
					if ( point_in_lasso( pt ) ) {
						singleComponent = curveCVIter.cv();
						newList.add (dagPath, singleComponent);
					}
				}
				break;
			}

		case MFn::kSurfaceCVComponent:
			{
				MItSurfaceCV surfCVIter( dagPath, component, true, &stat );
				for ( ; !surfCVIter.isDone(); surfCVIter.next() ) {
					point = surfCVIter.position(MSpace::kWorld, &stat );
					view.worldToView( point, pt.h, pt.v, &stat );
					if (!stat) {
						stat.perror("Could not get position");
						continue;
					}
					if ( point_in_lasso( pt ) ) {
						singleComponent = surfCVIter.cv();
						newList.add (dagPath, singleComponent);
					}
				}
				break;
			}

		case MFn::kMeshVertComponent:
			{
				MItMeshVertex vertexIter( dagPath, component, &stat );
				for ( ; !vertexIter.isDone(); vertexIter.next() ) {
					point = vertexIter.position(MSpace::kWorld, &stat );
					view.worldToView( point, pt.h, pt.v, &stat );
					if (!stat) {
						stat.perror("Could not get position");
						continue;
					}
					if ( point_in_lasso( pt ) ) {
						singleComponent = vertexIter.vertex();
						newList.add (dagPath, singleComponent);
					}
				}
				break;
			}

		case MFn::kMeshEdgeComponent:
			{
				MItMeshEdge edgeIter( dagPath, component, &stat );
				for ( ; !edgeIter.isDone(); edgeIter.next() ) {
					point = edgeIter.center(MSpace::kWorld, &stat );
					view.worldToView( point, pt.h, pt.v, &stat );
					if (!stat) {
						stat.perror("Could not get position");
						continue;
					}
					if ( point_in_lasso( pt ) ) {
						singleComponent = edgeIter.edge();
						newList.add (dagPath, singleComponent);
					}
				}
				break;
			}

		case MFn::kMeshPolygonComponent:
			{
				MItMeshPolygon polygonIter( dagPath, component, &stat );
				for ( ; !polygonIter.isDone(); polygonIter.next() ) {
					point = polygonIter.center(MSpace::kWorld, &stat );
					view.worldToView( point, pt.h, pt.v, &stat );
					if (!stat) {
						stat.perror("Could not get position");
						continue;
					}
					if ( point_in_lasso( pt ) ) {
						singleComponent = polygonIter.polygon();
						newList.add (dagPath, singleComponent);
					}
				}
				break;
			}

		default:
#ifdef DEBUG
			cerr << "Selected unsupported type: (" << component.apiType()
				 << "): " << component.apiTypeStr() << endl;
#endif /* DEBUG */
			continue;
		}
	}

	// Warn user if zie is trying to select objects rather than components.
	if (foundEntireObjects && !foundComponents) {
		MGlobal::displayWarning("lassoTool can only select components, not entire objects.");
	}

	// Update the selection list as indicated by the modifier keys.
	MGlobal::selectCommand(newList, listAdjustment);

	// Free the memory that held our lasso points.
	free(lasso);
	lasso = (coord*) 0;
	maxSize = 0;
	num_points = 0;
 
	return MS::kSuccess;
}

void lassoTool::append_lasso( short x, short y )
{
	int		cy, iy, ix, ydif, yinc, i;
	float	fx, xinc;

	iy = (int)lasso[num_points-1].v;
	ix = (int)lasso[num_points-1].h;
	ydif = abs( y - iy ); 
	if ( ydif == 0 )
		return;

	// Keep track of smallest rectangular area of the screen that
	// completely contains the lasso.
	if ( min.h > x )
		min.h = x;
	if ( max.h < x )
		max.h = x;
	if ( min.v > y )
		min.v = y;
	if ( max.v < y )
		max.v = y;

	if ( ((int)y - iy ) < 0 )
		yinc = -1;
	else
		yinc = 1;

	xinc = (float)((int)x - ix)/(float)ydif;
	fx = (float)ix + xinc;
	cy = iy + yinc;
	for ( i = 0; i < ydif; i++ ) {

		if ( num_points >= maxSize ) {
			// Make the array of lasso points bigger
			maxSize += increment;

			// If realloc() fails, it returns NULL but keeps the old block
			// of memory around, so let's not overwrite the contents of
			// 'lasso' until we know that realloc() worked.
			coord* newLasso = (coord*) realloc (lasso, sizeof(coord) * maxSize);
			if (newLasso == NULL) return;

			lasso = newLasso;
		}

		lasso[num_points].h = (short) fx;
		lasso[num_points].v = (short) cy;
		fx += xinc;
		cy += yinc;
		num_points++;
	}

	return;
}

void lassoTool::draw_lasso()
{
	glBegin( GL_LINE_LOOP );
	for ( unsigned i = 0; i < num_points ; i++ ){
		glVertex2i( lasso[i].h, lasso[i].v );
	}
	glEnd();
}

bool lassoTool::point_in_lasso( coord pt )
{
	unsigned i, sides;

	for ( i = 0; i < num_points; i++ ) {
		if ( lasso[i].v == pt.v ) {
			while ( (lasso[i].v == pt.v ) && (lasso[i].h < pt.h) )
				i++;
			if ( lasso[i].v != pt.v )
				return( false );
			sides = 0;
			i++;
			while ( lasso[i].v == pt.v  ) {
				i++;
				sides++;
			}
			if ( sides % 2 )
				return( false );
			else
				return( true );
		}
	}
	return( false );
}

//////////////////////////////////////////////
// Command to create contexts
//////////////////////////////////////////////

class lassoContextCmd : public MPxContextCommand
{
public: 
						lassoContextCmd() {};
	virtual MPxContext* makeObj();
	static void*		creator();
};

MPxContext* lassoContextCmd::makeObj()
{
	return new lassoTool;
}

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

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

	status = plugin.registerContextCommand( "lassoToolContext",
										    lassoContextCmd::creator );

	if (!status) {
		status.perror("registerContextCommand");
		return status;
	}

	// set the mel scripts to be run when the plugin is loaded / unloaded
	status = plugin.registerUI("lassoToolCreateUI", "lassoToolDeleteUI");
	if (!status) {
		status.perror("registerUIScripts");
		return status;
	}

	return status;
}

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

	status = plugin.deregisterContextCommand( "lassoToolContext" );

	return status;
}