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

#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/MPlug.h>
#include <maya/MRenderUtil.h>
#include <maya/MString.h>
#include <maya/MTime.h>
#include <maya/MTypeId.h>

#include "shaveGlobals.h"
#include "shaveError.h"
#include "shaveIO.h"
#include "shaveSDK.h"
#include "shaveBackgroundShader.h"
#include "shaveUtil.h"

MTypeId shaveBackgroundShader::id(0x001029A3);

const MString shaveBackgroundShader::nodeTypeName("shaveBackground");

MObject	shaveBackgroundShader::aBackgroundSampler;
MObject	shaveBackgroundShader::aIncludeBackfacing;

MObject	shaveBackgroundShader::aLightDataArray;
MObject		shaveBackgroundShader::aLightDirection;
MObject		shaveBackgroundShader::aLightIntensity;
MObject		shaveBackgroundShader::aLightAmbient;
MObject		shaveBackgroundShader::aLightDiffuse;
MObject		shaveBackgroundShader::aLightSpecular;
MObject		shaveBackgroundShader::aLightShadowFraction;
MObject		shaveBackgroundShader::aPreShadowIntensity;
MObject		shaveBackgroundShader::aLightBlindData;
MObject	shaveBackgroundShader::aMatrixEyeToWorld;
MObject	shaveBackgroundShader::aNormalCamera;
MObject	shaveBackgroundShader::aObjectId;
MObject	shaveBackgroundShader::aOldShader;
MObject	shaveBackgroundShader::aOutColor;
MObject	shaveBackgroundShader::aOutMatteOpacity;
MObject	shaveBackgroundShader::aOutTransparency;

#ifdef _DEBUG
MObject	shaveBackgroundShader::aPixelCenter;
MObject		shaveBackgroundShader::aPixelCenterX;
MObject		shaveBackgroundShader::aPixelCenterY;
#endif

MObject	shaveBackgroundShader::aPointWorld;
MObject	shaveBackgroundShader::aRayDepth;
MObject	shaveBackgroundShader::aRayDirection;
MObject	shaveBackgroundShader::aRayOrigin;


shaveBackgroundShader::shaveBackgroundShader()
{
}


shaveBackgroundShader::~shaveBackgroundShader()
{
}


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



MStatus shaveBackgroundShader::initialize()
{
	MStatus					st;
	MFnCompoundAttribute	cAttr;
	MFnMatrixAttribute		aMat;
	MFnMessageAttribute		aMsg;
	MFnNumericAttribute		nAttr; 

	//
	// Renderer-Supplied Input Attributes
	//
	aBackgroundSampler = nAttr.createAddr("backgroundSampler", "rtb");
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aBackgroundSampler);


	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.createAddr("lightBlindData", "lbd", 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 = aMat.create(
								"matrixEyeToWorld",
								"etw",
								MFnMatrixAttribute::kFloat
							);
	aMat.setHidden(true);
	aMat.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);


#ifdef _DEBUG
	aPixelCenterX = nAttr.create(
							"pixelCenterX", "pcx", MFnNumericData::kFloat
						);

	aPixelCenterY = nAttr.create(
							"pixelCenterY", "pcy", MFnNumericData::kFloat
						);
	aPixelCenter = nAttr.create(
							"pixelCenter",
							"pc",
							aPixelCenterX,
							aPixelCenterY
						);
	nAttr.setHidden(true);
	nAttr.setStorable(false);
	addAttribute(aPixelCenter);
#endif


	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
	//
	aIncludeBackfacing = nAttr.create(
		"includeBackfacing", "ibf", MFnNumericData::kBoolean, 0.0, &st
	);
	MChkErr(st, "cannot create 'includeBackfacing' attribute");
	nAttr.setDefault(true);
	st = addAttribute(aIncludeBackfacing);
	MChkErr(st, "cannot add 'includeBackfacing' attribute");


	//
	// Standard Renderer Output Attributes
	//
	aOutColor = nAttr.createColor("outColor", "oc");
	nAttr.setDefault(0.0f, 0.0f, 0.0f);
	nAttr.setHidden(true);
	nAttr.setWritable(false);
	nAttr.setUsedAsColor(true);
	addAttribute(aOutColor);


	aOutMatteOpacity = nAttr.createColor("outMatteOpacity", "omo");
	nAttr.setDefault(0.0f, 0.0f, 0.0f);
	nAttr.setHidden(true);
	nAttr.setWritable(false);
	nAttr.setUsedAsColor(true);
	addAttribute(aOutMatteOpacity);


	aOutTransparency = nAttr.createColor("outTransparency", "ot");
	nAttr.setDefault(0.0f, 0.0f, 0.0f);
	nAttr.setHidden(true);
	nAttr.setWritable(false);
	nAttr.setUsedAsColor(true);
	addAttribute(aOutTransparency);


	//
	// Custom Output Attributes
	//
	aOldShader = aMsg.create("oldShader", "oshd");
	aMsg.setHidden(true);
	aMsg.setWritable(false);
	aMsg.setHidden(true);
	addAttribute(aOldShader);


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

	attributeAffects(aIncludeBackfacing,	aOutMatteOpacity);
	attributeAffects(aBackgroundSampler,	aOutMatteOpacity);
	attributeAffects(aLightDataArray,	aOutMatteOpacity);
	attributeAffects(aMatrixEyeToWorld,	aOutMatteOpacity);
	attributeAffects(aNormalCamera,		aOutMatteOpacity);
	attributeAffects(aObjectId,			aOutMatteOpacity);
	attributeAffects(aPointWorld,		aOutMatteOpacity);
	attributeAffects(aRayDepth,			aOutMatteOpacity);
	attributeAffects(aRayDirection,		aOutMatteOpacity);
	attributeAffects(aRayOrigin,			aOutMatteOpacity);

	attributeAffects(aIncludeBackfacing,	aOutTransparency);
	attributeAffects(aBackgroundSampler,	aOutTransparency);
	attributeAffects(aLightDataArray,	aOutTransparency);
	attributeAffects(aMatrixEyeToWorld,	aOutTransparency);
	attributeAffects(aNormalCamera,		aOutTransparency);
	attributeAffects(aObjectId,			aOutTransparency);
	attributeAffects(aPointWorld,		aOutTransparency);
	attributeAffects(aRayDepth,			aOutTransparency);
	attributeAffects(aRayDirection,		aOutTransparency);
	attributeAffects(aRayOrigin,			aOutTransparency);

#ifdef _DEBUG
	attributeAffects(aPixelCenter,		aOutColor);
	attributeAffects(aPixelCenter,		aOutMatteOpacity);
	attributeAffects(aPixelCenter,		aOutTransparency);
#endif

    return MS::kSuccess;
}


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


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

		const MFloatVector&	normal = datablock
										.inputValue(aNormalCamera)
										.asFloatVector();
#ifdef _DEBUG
		float	pixelX = datablock.inputValue(aPixelCenterX).asFloat();
		float	pixelY = datablock.inputValue(aPixelCenterY).asFloat();
#endif

		bool	includeBackfacing = datablock.inputValue(aIncludeBackfacing).asBool();
		MArrayDataHandle	lightArray =
								datablock.inputArrayValue(aLightDataArray);

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

		unsigned int	i;
		float			lightShadow = 0.0f;
		float			netIllum = 0.0f;
		float			netShadow = 0.0f;
		float			preShadowIntensity;
		float			totalPreShadowIntensity = 0.0f;

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

static unsigned int count = 0;
if (++count == 7453) count = 0;

		for (i = 0; i < numLights; i++)
		{
			MDataHandle	light = lightArray.inputValue();

			//
			// If the preShadowIntensity is zero then the light doesn't
			// even point at us (or we're beyond it's range) in which case
			// we don't count it.
			//
			preShadowIntensity = light.child(aPreShadowIntensity).asFloat();

			if (preShadowIntensity <= 0.0f) continue;

			//
			// If the normal is facing away from the light then and we're
			// not including backfacing polys in the matte, then we don't want
			// to include that light in the calculation. Otherwise,
			// backfacing polys are considered to be totally shadowed from
			// the light.
			//
			float			hairIllum = 1.0f;
			MFloatVector	dir =
							light.child(aLightDirection).asFloatVector();

			if (dir * normal < 0.04f)
			{
				if (includeBackfacing)
					lightShadow = 1.0f;
				else
					continue;
			}
			else
				lightShadow = light.child(aLightShadowFraction).asFloat();

			numAffectingLights++;
			totalPreShadowIntensity += preShadowIntensity;

			//
			// If the point is already completely in shadow, then there's
			// no point in getting Shave's expensive opinion on the matter.
			//
			// Similarly, if Shave's shadow density is set to zero, then no
			// hair shadows will be displayed so once again we don't need
			// to ask Shave if the point is in shadow.
			//
			if ((lightShadow < 1.0f) && (shaveShadowDensityGlob > 0.0))
			{
				//
				// 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;
				void* blindData = light.child(aLightBlindData).asAddr();
				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();

					dir *= camToWorld;

					lightIndex = shaveUtil::getLightIndex(point, dir);
				}

				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);
				}
			}

			//
			// lightShadow gives the degree to which the point is shadowed
			// from that light, with 1.0 being fully shadowed and 0.0 being
			// fully lit.
			//
			// So (preShadowIntensity - lightShadow) gives the amount of
			// illumination from the light which hits the point, after
			// shadows but before taking hair into account.
			//
			float	lightIllum = preShadowIntensity - lightShadow;

			//
			// We multiply that by the amount of illumination that the hair
			// will let through to get the final illumination of the point.
			//
			float	finalIllum = lightIllum * hairIllum;

			//
			// The final amount of shadowing is (preShadowIntensity -
			// finalIllum).
			//
			lightShadow = preShadowIntensity - finalIllum;

			netShadow += lightShadow;

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

		if (totalPreShadowIntensity > 0.0f)
		{
			netShadow /= totalPreShadowIntensity;
			netIllum = 1.0f - netShadow;
		}
		else
			netShadow = 0.0f;

#define USE_BG_SAMPLER

#ifdef USE_BG_SAMPLER
		MFloatVector	traceColor(0.0f, 0.0f, 0.0f);

		if (netIllum > 0.0f)
		{
			int	rayDepth = datablock.inputValue(aRayDepth).asInt();
			void*&	objectId = datablock.inputValue(aObjectId).asAddr();
			void*&	backgroundSampler =
						datablock.inputValue(aBackgroundSampler).asAddr();

			const MFloatVector&	rayDirection =
						datablock.inputValue(aRayDirection).asFloatVector();

			const MFloatVector&	rayOrigin =
						datablock.inputValue(aRayOrigin).asFloatVector();

			MFloatVector	traceTransp;

			MRenderUtil::raytrace(
				rayOrigin,
				rayDirection,
				objectId,
				backgroundSampler,
				(short)rayDepth,
				traceColor,
				traceTransp,
				true
			);
		}
#endif

		MFloatVector&	outColor = datablock
										.outputValue(aOutColor)
										.asFloatVector();

		MFloatVector&	outMatteOpacity = datablock
											.outputValue(aOutMatteOpacity)
											.asFloatVector();

		MFloatVector&	outTransparency = datablock
											.outputValue(aOutTransparency)
											.asFloatVector();

#ifdef USE_BG_SAMPLER
		outColor = traceColor * netIllum;
#else
		outColor.x = 0.0;
		outColor.y = 0.0;
		outColor.z = 0.0;
#endif

		outMatteOpacity.x = netShadow;
		outMatteOpacity.y = netShadow;
		outMatteOpacity.z = netShadow;

#ifdef USE_BG_SAMPLER
		outTransparency.x = 0.0f;
		outTransparency.y = 0.0f;
		outTransparency.z = 0.0f;
#else
		outTransparency.x = netIllum;
		outTransparency.y = netIllum;
		outTransparency.z = netIllum;
#endif

		datablock.setClean(aOutColor);
		datablock.setClean(aOutMatteOpacity);
		datablock.setClean(aOutTransparency);
	}
	else
		return MS::kUnknownParameter;

	return MS::kSuccess;
}