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

#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MFloatVector.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnNumericData.h>
#include <maya/MPlug.h>
#include <maya/MString.h>
#include <maya/MTypeId.h>

#include <maya/MAnimControl.h>
#include <maya/MTime.h>

#include "shaveGlobals.h"
#include "shaveIO.h"
#include "shaveRender.h"
#include "shaveVolumeShader.h"

MTypeId	shaveVolumeShader::id(0x001029A2);
const MString shaveVolumeShader::nodeTypeName("shaveVolumeShader");

MObject	shaveVolumeShader::aFarPointCamera;
MObject	shaveVolumeShader::aFarPointWorld;
MObject	shaveVolumeShader::aOutColor;
MObject	shaveVolumeShader::aOutMatteOpacity;
MObject	shaveVolumeShader::aOutTransparency;
MObject	shaveVolumeShader::aPixelCenter;
MObject	shaveVolumeShader::aPixelCenterX;
MObject	shaveVolumeShader::aPixelCenterY;
MObject	shaveVolumeShader::aPointCamera;
MObject	shaveVolumeShader::aPointWorld;
MObject	shaveVolumeShader::aRayDepth;
MObject	shaveVolumeShader::aRayDirection;
MObject	shaveVolumeShader::aRayOrigin;


shaveVolumeShader::shaveVolumeShader()
{
}


shaveVolumeShader::~shaveVolumeShader()
{
}


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


MStatus shaveVolumeShader::initialize()
{
	MFnNumericAttribute	nAttr;

	//
	// Render Input Attributes
	//
	aFarPointCamera = nAttr.createPoint("farPointCamera", "fc");
	addAttribute(aFarPointCamera);


	aFarPointWorld = nAttr.createPoint("farPointWorld", "fw");
	addAttribute(aFarPointWorld);


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

	aPixelCenterY = nAttr.create(
							"pixelCenterY", "pcy", MFnNumericData::kFloat
						);

	aPixelCenter = nAttr.create(
							"pixelCenter",
							"pc",
							aPixelCenterX,
							aPixelCenterY
						);

	addAttribute(aPixelCenter);


	aPointCamera = nAttr.createPoint("pointCamera", "p");
	addAttribute(aPointCamera);


	aPointWorld = nAttr.createPoint("pointWorld", "pw");
	addAttribute(aPointWorld);


	aRayDepth = nAttr.create("rayDepth", "rd", MFnNumericData::kShort);
	addAttribute(aRayDepth);


	aRayDirection = nAttr.createPoint("rayDirection", "rad");
	addAttribute(aRayDirection);


	aRayOrigin = nAttr.createPoint("rayOrigin", "ro");
	addAttribute(aRayOrigin);


	//
	// Render Output Attributes
	//
	aOutColor = nAttr.createColor("outColor", "oc");
	addAttribute(aOutColor);


	aOutMatteOpacity = nAttr.createColor("outMatteOpacity", "omo");
	addAttribute(aOutMatteOpacity);


	aOutTransparency = nAttr.createColor("outTransparency", "ot");
	addAttribute(aOutTransparency);


	//
	// Set up attribute dependencies.
	//
	attributeAffects(aFarPointCamera,	aOutColor);
	attributeAffects(aFarPointWorld,	aOutColor);
	attributeAffects(aPixelCenter,		aOutColor);
	attributeAffects(aPointCamera,		aOutColor);
	attributeAffects(aPointWorld,		aOutColor);
	attributeAffects(aRayDepth,			aOutColor);
	attributeAffects(aRayDirection,		aOutColor);
	attributeAffects(aRayOrigin,		aOutColor);

	attributeAffects(aFarPointCamera,	aOutMatteOpacity);
	attributeAffects(aFarPointWorld,	aOutMatteOpacity);
	attributeAffects(aPixelCenter,		aOutMatteOpacity);
	attributeAffects(aPointCamera,		aOutMatteOpacity);
	attributeAffects(aPointWorld,		aOutMatteOpacity);
	attributeAffects(aRayDepth,			aOutMatteOpacity);
	attributeAffects(aRayDirection,		aOutMatteOpacity);
	attributeAffects(aRayOrigin,		aOutMatteOpacity);

	attributeAffects(aFarPointCamera,	aOutTransparency);
	attributeAffects(aFarPointWorld,	aOutTransparency);
	attributeAffects(aPixelCenter,		aOutTransparency);
	attributeAffects(aPointCamera,		aOutTransparency);
	attributeAffects(aPointWorld,		aOutTransparency);
	attributeAffects(aRayDepth,			aOutTransparency);
	attributeAffects(aRayDirection,		aOutTransparency);
	attributeAffects(aRayOrigin,		aOutTransparency);

	return MS::kSuccess;
}


void shaveVolumeShader::postConstructor()
{
}


MStatus shaveVolumeShader::compute(const MPlug& plug, MDataBlock& block)
{
	bool			hit = false;
	MFloatVector*	outValuePtr;
	float			defaultValue = 0.0f;

	int		pixelX = (int)block.inputValue(aPixelCenterX).asFloat();
	int		pixelY = (int)block.inputValue(aPixelCenterY).asFloat();

	shaveRender::SceneInfo*	sceneInfo = shaveRender::getSceneInfo();

	float			r = 0.0f;
	float			g = 0.0f;
	float			b = 0.0f;
	float			a = 0.0f;

	MFloatVector&	rayOrigin = block.inputValue(aRayOrigin)
											.asFloatVector();

	//
	// Eye rays are handled by the buffer composite.
	//
	if (rayOrigin.isEquivalent(MFloatVector::zero))
	{
		//
		// If we're compositing and have an image from Shave, then check to
		// see if this pixel contains any hair.
		//
		if (doShaveCompsGlob && !doShaveComp2dGlob && sceneInfo->bufferValid)
		{
			//
			// Get shave's Z-value for this pixel.
			//
			float	shaveZ = sceneInfo->shaveZBuffer[
								sceneInfo->width * pixelY + pixelX
							];

			//
			// If Shave's Z-value equals its far clip constant, then
			// there's no hair in the pixel, so we can ignore it.
			//
			if (shaveZ < (SHAVE_FAR_CLIP - 1.0f))
			{
				//
				// Shave's Z-value is really the distance from the camera.
				//
				// The near and far points handed to us by the renderer are
				// in the camera's *transformational* space, but not in
				// perspective space.  So their Z values do not equate to
				// distance from the camera.
				//
				// So, we must compute the distances to the near and far
				// points.
				//
				MFloatVector&	farPoint = block.inputValue(aFarPointCamera)
												.asFloatVector();
				MFloatVector&	nearPoint = block.inputValue(aPointCamera)
												.asFloatVector();

				float	nearDist = nearPoint.length();
				float	farDist = farPoint.length();

				//
				// If it lies between the segment's endpoints, set the
				// colour for the pixel.
				//
				if ((shaveZ >= nearDist) && (shaveZ <= farDist))
				{
					hit = true;

					shaveRender::Pixel*	shavePixel;

					shavePixel = &sceneInfo->shaveRenderPixels[
									sceneInfo->width*pixelY + pixelX
								];
					r = (float)shavePixel->r / 255.0f;
					g = (float)shavePixel->g / 255.0f;
					b = (float)shavePixel->b / 255.0f;
					a = (float)shavePixel->a / 255.0f;
				}
			}
		}
	}
	//
	// Non-eye rays (e.g. reflections & refractions) are handled by a call
	// to SHAVEtrace().
	//
	else
	{
		if (visibleInReflectionsGlob)
		{
			MFloatVector&	farPoint = block.inputValue(aFarPointWorld)
											.asFloatVector();
			MFloatVector&	nearPoint = block.inputValue(aPointWorld)
											.asFloatVector();
			MFloatVector	ray = farPoint - nearPoint;

			VOXSAMP	sample;
			VERT	rbase;
			VERT	rdir;
			float	nearClip = 0.0f;
			float	farClip = ray.length();

			rbase.x = nearPoint.x;
			rbase.y = nearPoint.y;
			rbase.z = -nearPoint.z;

			rdir.x = ray.x;
			rdir.y = ray.y;
			rdir.z = -ray.z;

			if (!sceneInfo->shaveTraceInitialized)
			{
				SHAVEtrace_init();
				sceneInfo->shaveTraceInitialized = true;
			}

			if (SHAVEtrace(nearClip, farClip, rbase, rdir, 1, &sample))
			{
				hit = true;

				r = sample.color.x;
				g = sample.color.y;
				b = sample.color.z;
				a = sample.opacity;
			}
		}
	}

	//
	// In theory, the renderer could ask for either the parent plug's
	// value (e.g. outColor) or each of its children individually
	// (e.g. outColor.r).
	//
	// In practice, it only ever asks for the parent, and checking the
	// children adds a tiny but measurable amount of overhead to the
	// render.  So we'll only check the parent plugs and if that causes a
	// problem in some future version of the renderer, we'll make the
	// change then.
	//
	if (plug == aOutColor)
	{
		outValuePtr = &block.outputValue(aOutColor).asFloatVector();

		if (hit)
		{
			(*outValuePtr).x = r;
			(*outValuePtr).y = g;
			(*outValuePtr).z = b;
		}
		else
		{
			//
			// If we didn't get a hit, then make the pixel black otherwise
			// Maya's renderer will still use the colour, even if it's
			// transparent.
			//
			defaultValue = 0.0f;
		}
	}
	else if (plug == aOutMatteOpacity)
	{
		outValuePtr = &block.outputValue(aOutMatteOpacity).asFloatVector();

		if (hit)
		{
			(*outValuePtr).x = a;
			(*outValuePtr).y = a;
			(*outValuePtr).z = a;
		}
		else
			defaultValue = 0.0f;
	}
	else if (plug == aOutTransparency)
	{
		outValuePtr = &block.outputValue(aOutTransparency).asFloatVector();

		if (hit)
		{
			float	transparency = 1.0f - a;

			(*outValuePtr).x = transparency;
			(*outValuePtr).y = transparency;
			(*outValuePtr).z = transparency;
		}
		else
			defaultValue = 1.0f;
	}
	else
		return MS::kUnknownParameter;

	//
	// If we didn't get a hit on a hair, then just output the default
	// value, which is transparent black.
	//
	if (!hit)
	{
		(*outValuePtr).x = defaultValue;
		(*outValuePtr).y = defaultValue;
		(*outValuePtr).z = defaultValue;
	}

	block.setClean(plug);

	return MS::kSuccess;
}