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

#include <maya/MFloatArray.h>
#include <maya/MItMeshFaceVertex.h>
#include <maya/MObjectArray.h>
#include <maya/MStatus.h>

#include "shaveAPI.h"
#include "shaveAPIimpl.h"
#include "shaveConstant.h"
#include "shaveGlobals.h"
#include "shaveIO.h"
#include "shaveItHairImpl.h"
#include "shaveHairShape.h"
#include "shaveRender.h"
#include "shaveRenderer.h"
#include "shaveSDK.h"
#include "shaveUtil.h"


bool			shaveItHairImpl::mDoingInstances = false;

shaveConstant::RenderMode
				shaveItHairImpl::mHairRenderMode = shaveConstant::kNoRender;

MFloatArray		shaveItHairImpl::mInstanceUs;
MFloatArray		shaveItHairImpl::mInstanceVs;
int				shaveItHairImpl::mLastNodeIndex = -1;
shaveRenderer*	shaveItHairImpl::mRenderer = NULL;


MStatus shaveItHairImpl::clear()
{
	shaveAPIimpl::clearHairStack();
	mInstanceUs.clear();
	mInstanceVs.clear();
	mLastNodeIndex = -1;

	return MS::kSuccess;
}


MStatus shaveItHairImpl::doInit(bool instances, MObjectArray& shaveHairShapes)
{
	//
	// Get the render mode.
	//
	mHairRenderMode = mRenderer->getRenderMode();

	//
	// If 'instances' is true, remove any non-instanced shaveHairShapes.
	// If it's false, remove any instanced ones.
	//
	unsigned int	i;

	for (i = 0; i < shaveHairShapes.length();)
	{
		MFnDependencyNode	nodeFn(shaveHairShapes[i]);
		shaveHairShape*			nodePtr = (shaveHairShape*)nodeFn.userNode();

		if (nodePtr->isInstanced() != instances)
			shaveHairShapes.remove(i);
		else
			i++;
	}

	if (shaveHairShapes.length() == 0) return MS::kNotFound;

	mDoingInstances = instances;

	return shaveAPIimpl::createHairStack(
				shaveHairShapes, shaveRender::getFrameGlobals()
			);
}


MStatus shaveItHairImpl::init(bool instances, bool renderableOnly)
{
	shaveRender::saveFrameGlobals();

	mRenderer = shaveRender::getRenderer();

	MObjectArray	shaveHairShapes;

	if (renderableOnly)
		mRenderer->getRenderableShaveNodes(shaveHairShapes);
	else
		shaveUtil::getShaveNodes(shaveHairShapes);

	return doInit(instances, shaveHairShapes);
}


MStatus shaveItHairImpl::init(bool instances, MObjectArray& shaveHairShapes)
{
	shaveRender::saveFrameGlobals();

	mRenderer = shaveRender::getRenderer();

	unsigned int	i;

	for (i = 0; i < shaveHairShapes.length(); i++)
	{
		MFnDependencyNode	nodeFn(shaveHairShapes[i]);

		if (nodeFn.typeId() != shaveHairShape::id) return MS::kInvalidParameter;
	}

	return doInit(instances, shaveHairShapes);
}


MStatus shaveItHairImpl::nextHair(shaveAPI::HairInfo* hairInfo)
{
	if (hairInfo == NULL) return MS::kInvalidParameter;

	int	nodeIndex;

	//
	// If we're doing instanced hair, or if the we're doing non-instanced
	// hair but the render mode is Geom, then get polys.  Otherwise get
	// curves.
	//
	if (mDoingInstances
	||	(mHairRenderMode == shaveConstant::kGeometryRender))
	{
		nodeIndex = SHAVEexport_poly_iterator((HAIRTYPE*)hairInfo);

		//
		// If we're doing instances then Shave won't provide the UVs.
		// Instead we'll have to get those from shaveHairShape's instance input
		// surface.
		//
		if (mDoingInstances)
		{
			shaveHairShape*	nodePtr = shaveAPIimpl::getNodeFromStack(nodeIndex);

			if (nodePtr)
			{
				//
				// If this is a new node then cache its UVs.
				//
				if (nodeIndex != mLastNodeIndex)
				{
					MPlug	plug(
								nodePtr->thisMObject(),
								shaveHairShape::instanceMesh
							);
					MObject	meshObj;

					plug.getValue(meshObj);

					//
					// We cannot just grab the UV array from the mesh and
					// stuff it into hairInfo->uvws because the mesh uses a
					// shared vertex list while hairInfo currently does not.
					//
					// So we have to iterate through the vertices,
					// hopefully in the same order as was used when the
					// instance mesh was passed to Shave, and assign the
					// UVs individually.
					//
					float2				uv;
					MItMeshFaceVertex	faceVertIter(meshObj);

					mInstanceUs.clear();
					mInstanceVs.clear();

					for (; !faceVertIter.isDone(); faceVertIter.next())
					{
						faceVertIter.getUV(uv);
						mInstanceUs.append(uv[0]);
						mInstanceVs.append(uv[1]);
					}
				}

				//
				// Copy uvs from the cache to each vertex of each strand.
				//
				// Note that for instanced hair 'numHairs' is neither the
				// number of hairs nor the the number of strands, but the
				// number of faces in the polymesh for all of the hair's
				// strands.  So if the hair has 3 strands of 10 polys each
				// then 'numHairs' will be 30.
				//
				// So we have to work backward to the strand count.
				//
				int	numVertsPerStrand = mInstanceUs.length();
				int	numStrands;

				if (numVertsPerStrand > 0)
					numStrands = hairInfo->numHairVertices / numVertsPerStrand;
				else
					numStrands = 0;

				int i = 0;
				int	strand;
				int	v;

				for (strand = 0; strand < numStrands; strand++)
				{
					for (v = 0; v < numVertsPerStrand; v++)
					{
						hairInfo->uvws[i].x = mInstanceUs[v];
						hairInfo->uvws[i].y = 1.0f-mInstanceVs[v];
						hairInfo->uvws[i].z = 0.0f;
						i++;
					}
				}
			}
		}
	}
	else
	{
		nodeIndex = SHAVEexport_iterator((HAIRTYPE*)hairInfo);
	}

	if (nodeIndex == -1) return MS::kNotFound;

	mLastNodeIndex = nodeIndex;

	return MS::kSuccess;
}


MStatus shaveItHairImpl::nextHairCounts(shaveAPI::HairInfo* hairInfo)
{
	if (hairInfo == NULL) return MS::kInvalidParameter;

	int	nodeIndex = SHAVEexport_iteratorROOT((HAIRTYPE*)hairInfo);

	//
	// A return of -1 indicates that we've run out of hairs.
	//
	if (nodeIndex == -1) return MS::kNotFound;

	//
	// The returned hairInfo will not contain correct counts for the number
	// of vertices and hair-vertices in this hair.  So we need to calculate
	// for ourselves, using the shaveHairShape's segment count.
	//
	shaveHairShape*	nodePtr;

	nodePtr = shaveAPIimpl::getNodeFromStack((unsigned)nodeIndex);

	hairInfo->numVertices = 0;
	hairInfo->numHairVertices = 0;

	if (nodePtr != NULL)
	{
		SHAVENODE*	hairNode = nodePtr->getHairNode();

		if (hairNode != NULL)
		{
			int	numSegs = hairNode->shavep.segs[nodePtr->getHairGroup()];

			hairInfo->numVertices = hairInfo->numHairs * numSegs;
			hairInfo->numHairVertices = hairInfo->numVertices;
		}
	}

	return MS::kSuccess;
}


MStatus shaveItHairImpl::reset()
{
	SHAVEreset_iterator();

	return MS::kSuccess;
}