// Shave and a Haircut
// (c) 2019 Epic Games
// US Patent 6720962

#include <maya/MColor.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MFloatVector.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MFnMesh.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MGlobal.h>
#include <maya/MItDependencyGraph.h>
#include <maya/MItMeshPolygon.h>
#include <maya/MItMeshVertex.h>
#include <maya/MPlug.h>
#include <maya/MPointArray.h>
#include <maya/MSelectionList.h>
#include <maya/MString.h>
#include <maya/MTypeId.h>

#include <math.h>

#include "shaveVertexShader.h"

// Static data
MTypeId shaveVertexShader::id(0x00102998);
const MString shaveVertexShader::nodeTypeName("shaveVertexShader");


// Attributes 
MObject shaveVertexShader::pointAttr;
MObject shaveVertexShader::surfaceIndexAttr;
MObject shaveVertexShader::objectIdAttr;
MObject shaveVertexShader::primitiveIdAttr;
MObject shaveVertexShader::outColorAttr;
MObject shaveVertexShader::outAlphaAttr;


void shaveVertexShader::postConstructor()
{
	setMPSafe(true);
}


shaveVertexShader::shaveVertexShader()
{
}


shaveVertexShader::~shaveVertexShader()
{
}


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


MStatus shaveVertexShader::initialize()
{
    MFnNumericAttribute nAttr;

    pointAttr  = nAttr.create("pointObj", "po", MFnNumericData::k3Float);
	nAttr.setStorable(false);
    nAttr.setHidden(true);

	surfaceIndexAttr = nAttr.create("surfaceIndex", "si", MFnNumericData::kFloat);
    nAttr.setHidden(true);

	objectIdAttr = nAttr.create("objectId", "oi", MFnNumericData::kLong);
    nAttr.setHidden(true);

	primitiveIdAttr = nAttr.create("primitiveId", "pi", MFnNumericData::kLong);
    nAttr.setHidden(true);

    outColorAttr = nAttr.create("outColor", "oc", MFnNumericData::k3Float);
    nAttr.setStorable(false);
    nAttr.setReadable(true);
    nAttr.setWritable(false);

	outAlphaAttr = nAttr.create( "outAlpha", "oa", MFnNumericData::kFloat);
	nAttr.setDisconnectBehavior(MFnAttribute::kReset);
    nAttr.setStorable(false);
    nAttr.setReadable(true);
    nAttr.setWritable(false);

    addAttribute(pointAttr);
    addAttribute(outColorAttr);
    addAttribute(outAlphaAttr);
    addAttribute(surfaceIndexAttr);
    addAttribute(objectIdAttr);
    addAttribute(primitiveIdAttr);

    attributeAffects(pointAttr,			outColorAttr);
    attributeAffects(surfaceIndexAttr,	outColorAttr);
    attributeAffects(objectIdAttr,		outColorAttr);
    attributeAffects(primitiveIdAttr,	outColorAttr);

    attributeAffects(pointAttr,			outAlphaAttr);
    attributeAffects(surfaceIndexAttr,	outAlphaAttr);
    attributeAffects(objectIdAttr,		outAlphaAttr);
    attributeAffects(primitiveIdAttr,	outAlphaAttr);

	return MS::kSuccess;
}


MStatus shaveVertexShader::compute( const MPlug& plug, MDataBlock& block ) 
{
	if ((plug != outColorAttr) && (plug.parent() != outColorAttr) && 
		(plug != outAlphaAttr))
	{
		return MS::kUnknownParameter;
	}

	MStatus status;
	MObject thisNode = thisMObject();
	MColor	resultColour(0.5, 0.5, 0.5, 1.0);

	long surfaceIndex = (long)(block.inputValue(surfaceIndexAttr).asFloat() + 0.5);
	long triangleID = block.inputValue(primitiveIdAttr).asLong();

	// Location of the point we are shading
	float3&			samplePtIn = block.inputValue(pointAttr).asFloat3();
	MFloatVector	samplePt(samplePtIn);

	// Find the Mesh object
	MPlug outColorPlug = MFnDependencyNode(thisNode).findPlug("outColor", 
															  &status);
	
	MItDependencyGraph depIt( outColorPlug, MFn::kShadingEngine,
							  MItDependencyGraph::kDownstream, 
							  MItDependencyGraph::kBreadthFirst,
							  MItDependencyGraph::kNodeLevel, 
							  &status);
	
	depIt.enablePruningOnFilter();

	MObject shadingEngineNode = depIt.thisNode();
	
	MPlug dagSetMembersPlug = MFnDependencyNode(shadingEngineNode).findPlug(
								"dagSetMembers", &status
							);

	if (surfaceIndex < (long)dagSetMembersPlug.numElements())
	{
		MPlug dagSetMembersElementPlug;

		dagSetMembersElementPlug = dagSetMembersPlug.elementByLogicalIndex( 
										surfaceIndex, &status
									);

		MPlugArray meshPlugArray;

		dagSetMembersElementPlug.connectedTo(
			meshPlugArray, true, false, &status
		);

		if (meshPlugArray.length() > 0)
		{
			MObject meshNode = meshPlugArray[0].node();

			int polygonID = 0;
			int prevIndex = 0;

			int i;
			MFnMesh meshFn(meshNode);
			MItMeshPolygon iter(meshNode);

			int		numPolygons = meshFn.numPolygons();

			//
			// If the mesh is empty, then there's nothing to do.
			//
			if (numPolygons > 0)
			{
				//	We used to cache the face and triangle IDs for the
				//	current object, under the assumption that we would get
				//	many consecutive hits on the same object, thereby
				//	speeding things up.  But the caching makes this node MP
				//	unsafe and disabling MP slows things down considerably.
				//	Also, it appears that the caching alone is slower than
				//	just handling each sample on the fly.  So we no longer
				//	cache
				for (i = 0; i < numPolygons; i++)
				{
					int nTri;
					iter.setIndex(i, prevIndex);
					iter.numTriangles(nTri);

					if (triangleID < nTri)
					{
						polygonID = i;
						break;
					}

					triangleID -= nTri;
				}

				iter.setIndex(polygonID, prevIndex);

				MPointArray v;
				MIntArray vertIndices;

				//
				// Get the triangle's vertices.
				//
				status = iter.getTriangle(
							triangleID, v, vertIndices, MSpace::kObject
						);

				MFloatVector p1((float)v[0].x, (float)v[0].y, (float)v[0].z);
				MFloatVector p2((float)v[1].x, (float)v[1].y, (float)v[1].z);
				MFloatVector p3((float)v[2].x, (float)v[2].y, (float)v[2].z);

				MColor colour1;
				MColor colour2;
				MColor colour3;

				MItMeshVertex vertIter(meshNode);

				prevIndex = 0;

				vertIter.setIndex(vertIndices[0], prevIndex);
				vertIter.getColor(colour1, polygonID);
		
				vertIter.setIndex(vertIndices[1], prevIndex);
				vertIter.getColor(colour2, polygonID);

				vertIter.setIndex(vertIndices[2], prevIndex);
				vertIter.getColor(colour3, polygonID);

				//
				// Calculate the point's barycentric coordinates.
				//
				samplePt = samplePt - p3;
			    p1 = p1 - p3;
			    p2 = p2 - p3;

			    MFloatVector norm = p1 ^ p2;
			    float lenSquared = norm * norm;

			    lenSquared = (norm * samplePt) / lenSquared;

				//
				// The point may not be exactly on the triangle, so move it
				// there.
				//
			    samplePt = samplePt - (lenSquared * norm);

			    float aa = p1 * p1;
			    float bb = p2 * p2;
			    float ab = p1 * p2;
			    float am = p1 * samplePt;
			    float bm = p2 * samplePt;
			    float det = aa*bb - ab*ab;

				MFloatVector	abc;
			    abc.x = (am*bb - bm*ab) / det;
			    abc.y = (bm*aa - am*ab) / det;
				abc.z = 1 - abc.x - abc.y;

				resultColour = (abc.x*colour1)
							 + (abc.y*colour2)
							 + (abc.z*colour3);
			}
		}
	}

	MDataHandle outColorHandle = block.outputValue(outColorAttr);
	MFloatVector& outColor = outColorHandle.asFloatVector();

	outColor.x = resultColour.r;
	outColor.y = resultColour.g;
	outColor.z = resultColour.b;
	outColorHandle.setClean();

	MDataHandle outAlphaHandle = block.outputValue(outAlphaAttr);
	float& outAlpha = outAlphaHandle.asFloat();

	outAlpha = resultColour.a;
	outAlphaHandle.setClean();

	return MS::kSuccess;
}