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

#include <maya/MArrayDataBuilder.h>
#include <maya/MArrayDataHandle.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MFloatMatrix.h>
#include <maya/MFloatVector.h>
#include <maya/MFnCompoundAttribute.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MFnMatrixAttribute.h>
#include <maya/MFnMessageAttribute.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnNumericData.h>
#include <maya/MGlobal.h>
#include <maya/MPlug.h>
#include <maya/MRenderUtil.h>
#include <maya/MString.h>
#include <maya/MTypeId.h>

#include "shaveGlobals.h"
#include "shaveIO.h"
#include "shaveSDK.h"
#include "shaveShadowFilter.h"
#include "shaveUtil.h"

MTypeId shaveShadowFilter::id(0x001029BC);

const MString shaveShadowFilter::nodeTypeName("shaveShadowFilter");

MObject	shaveShadowFilter::aLightDataArray;
MObject		shaveShadowFilter::aLightDirection;
MObject		shaveShadowFilter::aLightIntensity;
MObject		shaveShadowFilter::aLightAmbient;
MObject		shaveShadowFilter::aLightDiffuse;
MObject		shaveShadowFilter::aLightSpecular;
MObject		shaveShadowFilter::aLightShadowFraction;
MObject		shaveShadowFilter::aPreShadowIntensity;
MObject		shaveShadowFilter::aLightBlindData;

MObject	shaveShadowFilter::aLightDataArrayOut;
MObject		shaveShadowFilter::aLightDirectionOut;
MObject		shaveShadowFilter::aLightIntensityOut;
MObject		shaveShadowFilter::aLightAmbientOut;
MObject		shaveShadowFilter::aLightDiffuseOut;
MObject		shaveShadowFilter::aLightSpecularOut;
MObject		shaveShadowFilter::aLightShadowFractionOut;
MObject		shaveShadowFilter::aPreShadowIntensityOut;
MObject		shaveShadowFilter::aLightBlindDataOut;

MObject	shaveShadowFilter::aMatrixEyeToWorld;
MObject	shaveShadowFilter::aNormalCamera;
MObject	shaveShadowFilter::aObjectId;

MObject	shaveShadowFilter::aPointWorld;
MObject	shaveShadowFilter::aRayDepth;
MObject	shaveShadowFilter::aRayDirection;
MObject	shaveShadowFilter::aRayOrigin;


shaveShadowFilter::shaveShadowFilter()
{
}


shaveShadowFilter::~shaveShadowFilter()
{
}


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



MStatus shaveShadowFilter::initialize()
{
	MStatus					st;
	MFnCompoundAttribute	cAttr;
	MFnMatrixAttribute		matAttr;
	MFnMessageAttribute		msgAttr;
	MFnNumericAttribute		nAttr; 

	//
	// Renderer-Supplied Input Attributes
	//
	aLightDirection = nAttr.createPoint("lightDirection", "ld");
	nAttr.setDefault(1.0f, 1.0f, 1.0f);


	aLightIntensity = nAttr.createColor("lightIntensity", "li");
	nAttr.setDefault(1.0f, 1.0f, 1.0f);


	aLightAmbient = nAttr.create(
							"lightAmbient",
							"la",
							MFnNumericData::kBoolean,
							true
						);

	aLightDiffuse = nAttr.create(
							"lightDiffuse",
							"ldf",
							MFnNumericData::kBoolean,
							true
						);

	aLightSpecular = nAttr.create(
							"lightSpecular",
							"ls",
							MFnNumericData::kBoolean,
							false
						);

	aLightShadowFraction = nAttr.create(
							"lightShadowFraction",
							"lsf",
							MFnNumericData::kFloat,
							0.5
						);

	aPreShadowIntensity = nAttr.create(
							"preShadowIntensity",
							"psi",
							MFnNumericData::kFloat,
							1.0
						);

	aLightBlindData = nAttr.create(
							"lightBlindData", "lbd", MFnNumericData::kLong, 0
						);

	aLightDataArray = cAttr.create("lightDataArray", "ltd");
	cAttr.addChild(aLightDirection);
	cAttr.addChild(aLightIntensity);
	cAttr.addChild(aLightAmbient);
	cAttr.addChild(aLightDiffuse);
	cAttr.addChild(aLightSpecular);
	cAttr.addChild(aLightShadowFraction);
	cAttr.addChild(aPreShadowIntensity);
	cAttr.addChild(aLightBlindData);
	cAttr.setArray(true);
	cAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aLightDataArray);


	aMatrixEyeToWorld = matAttr.create(
								"matrixEyeToWorld",
								"etw",
								MFnMatrixAttribute::kFloat
							);
	matAttr.setHidden(true);
	matAttr.setStorable(false);
	addAttribute(aMatrixEyeToWorld);


	aNormalCamera = nAttr.createPoint("normalCamera", "n");
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aNormalCamera);


	aObjectId = nAttr.createAddr("objectId", "oi");
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aObjectId);


	aPointWorld = nAttr.createPoint("pointWorld", "pw");
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aPointWorld);


	aRayDepth = nAttr.create("rayDepth", "rd", MFnNumericData::kInt);
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aRayDepth);


	aRayDirection = nAttr.createPoint("rayDirection", "rad");
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aRayDirection);


	aRayOrigin = nAttr.createPoint("rayOrigin", "ro");
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aRayOrigin);


	//
	// Custom Input Attributes
	//


	//
	// Standard Renderer Output Attributes
	//

	//
	// Custom Output Attributes
	//
	aLightDirectionOut = nAttr.createPoint("lightDirectionOut", "ldo", &st);
	nAttr.setDefault(1.0f, 1.0f, 1.0f);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightDirectionOut attr: "
			+ st.errorString()
		);
		return st;
	}


	aLightIntensityOut = nAttr.createColor("lightIntensityOut", "lio", &st);
	nAttr.setDefault(1.0f, 1.0f, 1.0f);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightIntensityOut attr: "
			+ st.errorString()
		);
		return st;
	}


	aLightAmbientOut = nAttr.create(
							"lightAmbientOut",
							"lao",
							MFnNumericData::kBoolean,
							true,
							&st
						);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightAmbientOut attr: "
			+ st.errorString()
		);
		return st;
	}

	aLightDiffuseOut = nAttr.create(
							"lightDiffuseOut",
							"ldfo",
							MFnNumericData::kBoolean,
							true,
							&st
						);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightDiffuseOut attr: "
			+ st.errorString()
		);
		return st;
	}

	aLightSpecularOut = nAttr.create(
							"lightSpecularOut",
							"lso",
							MFnNumericData::kBoolean,
							false,
							&st
						);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightSpecularOut attr: "
			+ st.errorString()
		);
		return st;
	}

	aLightShadowFractionOut = nAttr.create(
							"lightShadowFractionOut",
							"lsfo",
							MFnNumericData::kFloat,
							0.5,
							&st
						);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightShadowFractionOut attr: "
			+ st.errorString()
		);
		return st;
	}

	aPreShadowIntensityOut = nAttr.create(
							"preShadowIntensityOut",
							"psio",
							MFnNumericData::kFloat,
							1.0,
							&st
						);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create preShadowIntensityOut attr: "
			+ st.errorString()
		);
		return st;
	}

	aLightBlindDataOut = nAttr.create(
							"lightBlindDataOut",
							"lbdo",
							MFnNumericData::kLong,
							0,
							&st
						);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightBlindDataOut attr: "
			+ st.errorString()
		);
		return st;
	}

	aLightDataArrayOut = cAttr.create("lightDataArrayOut", "ltdo", &st);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot create lightDataArrayOut attr: "
			+ st.errorString()
		);
		return st;
	}
	cAttr.addChild(aLightDirectionOut);
	cAttr.addChild(aLightIntensityOut);
	cAttr.addChild(aLightAmbientOut);
	cAttr.addChild(aLightDiffuseOut);
	cAttr.addChild(aLightSpecularOut);
	cAttr.addChild(aLightShadowFractionOut);
	cAttr.addChild(aPreShadowIntensityOut);
	cAttr.addChild(aLightBlindDataOut);
	cAttr.setArray(true);
	cAttr.setUsesArrayDataBuilder(true);
	cAttr.setHidden(true);
	cAttr.setStorable(false);
	st = addAttribute(aLightDataArrayOut);
	if (!st)
	{
		MGlobal::displayError(
			nodeTypeName + ": cannot add lightDataArrayOut attr: "
			+ st.errorString()
		);
		return st;
	}


	//
	// Set up attribute dependencies.
	//
	attributeAffects(aLightDataArray,	aLightDataArrayOut);
	attributeAffects(aMatrixEyeToWorld,	aLightDataArrayOut);
	attributeAffects(aNormalCamera,		aLightDataArrayOut);
	attributeAffects(aObjectId,			aLightDataArrayOut);
	attributeAffects(aPointWorld,		aLightDataArrayOut);
	attributeAffects(aRayDepth,			aLightDataArrayOut);
	attributeAffects(aRayDirection,		aLightDataArrayOut);
	attributeAffects(aRayOrigin,		aLightDataArrayOut);

    return MS::kSuccess;
}


void shaveShadowFilter::postConstructor( )
{
//	%%% Can we safely do this, given the use of globals by the Shave
//	    engine?
//
//	setMPSafe(true);
}


MStatus shaveShadowFilter::compute(const MPlug& plug, MDataBlock& datablock)
{
	if ((plug == aLightDataArrayOut)
	||	(plug.parent() == aLightDataArrayOut))
	{
		const MFloatVector&	point = datablock
										.inputValue(aPointWorld)
										.asFloatVector();

		const MFloatVector&	normal = datablock
										.inputValue(aNormalCamera)
										.asFloatVector();

		MArrayDataHandle	lightArrayIn =
								datablock.inputArrayValue(aLightDataArray);

		unsigned int	numAffectingLights = 0;
		unsigned int	numLights = lightArrayIn.elementCount();

		MArrayDataBuilder	lightArrayOutBuilder(
								&datablock, aLightDataArrayOut, numLights
							);

		int				blindData;
		unsigned int	i;
		MFloatVector	intensity;
		MFloatVector	lightDir;
		float			lightShadow = 0.0f;
		float			netIllum = 0.0f;
		float			netShadow = 0.0f;
		float			preShadowIntensity;
		float			shadowFraction;
		float			totalPreShadowIntensity = 0.0f;

		VERT wpoint;
		wpoint.x = point.x;
		wpoint.y = point.y;
		wpoint.z = point.z;

		for (i = 0; i < numLights; i++)
		{
			MDataHandle	lightIn = lightArrayIn.inputValue();
			MDataHandle	lightOut = lightArrayOutBuilder.addElement(
										lightArrayIn.elementIndex()
									);

			//
			// Get the values for this light.
			//
			blindData = lightIn.child(aLightBlindData).asAddr();
			intensity = lightIn.child(aLightIntensity).asFloatVector();
			lightDir = lightIn.child(aLightDirection).asFloatVector();
			preShadowIntensity = lightIn.child(aPreShadowIntensity).asFloat();
			shadowFraction = lightIn.child(aLightShadowFraction).asFloat();

			//
			// Checking for hair shadows is expensive, so let's avoid it if
			// at all possible.
			//
			if ((preShadowIntensity > 0.0f)
			&&	(shadowFraction < 1.0f)
			&&	(shaveShadowDensityGlob > 0.0f)
			&&	(intensity.x + intensity.y + intensity.z > 0.0f))
			{
				//
				// If the surface normal is facing away from the light then
				// there won't be any hair shadows on it.
				//
				float	hairIllum = 1.0f;

				if (lightDir * normal > 0.04f)
				{
					//
					// We need to get this light's index in the
					// globalLightList.
					//
					// If the light has a non-zero 'lightBlindData'
					// attribute value, then that will be its node pointer,
					// which is unique, so we can use it to look up the
					// light's index.
					//
					int	lightIndex = -1;

					if (blindData != 0)
						lightIndex = shaveUtil::getLightIndex(blindData);

					//
					// If we didn't find the light through its blind data,
					// then try finding it based on its direction.
					//
					if (lightIndex < 0)
					{
						const MFloatMatrix&	camToWorld =
								datablock.inputValue(aMatrixEyeToWorld).asFloatMatrix();

						lightDir *= camToWorld;
						lightIndex = shaveUtil::getLightIndex(point, lightDir);
					}

					if (lightIndex >= 0)
					{
						//
						// How much of the light's illumination gets
						// through the hair?
						//
						// Note that internally Shave creates multiple
						// sub-lights for some types of lights.  They're
						// not supposed to overlap but sometimes they do,
						// so we only want the brightest one.
						//
						int		id;
						float	temp;

						hairIllum = 10000.0f;

						for (id = shaveUtil::globalLightList[lightIndex].minId;
							id <= shaveUtil::globalLightList[lightIndex].maxId;
							id++)
						{
							temp = SHAVEilluminate_point(wpoint, id);
							if (temp < hairIllum) hairIllum = temp;
						}

						hairIllum = hairIllum * shaveShadowDensityGlob
									+ (1.0f - shaveShadowDensityGlob);

						//
						// Apply the hair's dimming effect to the light's
						// intensity.
						//
						intensity *= hairIllum;

						//
						// Add the dimming effect to the light's shadow
						// fraction.
						//
						shadowFraction += 1.0f - hairIllum;

						if (shadowFraction > 1.0f) shadowFraction = 1.0f;
					}
				}
			}

			//	Fill in the output values.
			lightOut.child(aLightDirectionOut).set(lightDir);
			lightOut.child(aLightIntensityOut).set(intensity);
			lightOut.child(aLightShadowFractionOut).set(shadowFraction);
			lightOut.child(aLightBlindDataOut).set(blindData);
			lightOut.child(aPreShadowIntensityOut).set(preShadowIntensity);

			lightOut.child(aLightAmbientOut).set(
				lightIn.child(aLightAmbient).asBool()
			);

			lightOut.child(aLightDiffuseOut).set(
				lightIn.child(aLightDiffuse).asBool()
			);

			lightOut.child(aLightSpecularOut).set(
				lightIn.child(aLightSpecular).asBool()
			);

			if (!lightArrayIn.next()) break;
		}

		datablock.outputArrayValue(aLightDataArrayOut).set(
			lightArrayOutBuilder
		);
	}
	else
		return MS::kUnknownParameter;

	return MS::kSuccess;
}