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

#include <float.h>
#include <math.h>

#include <maya/MDagPath.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MFnDagNode.h>
#include <maya/MFnTransform.h>
#include <maya/MGlobal.h>
#include <maya/MItDag.h>
#include <maya/MMatrix.h>
#include <maya/MPoint.h>
#include <maya/MRenderUtil.h>
#include <maya/MRenderView.h>
#include <maya/MVector.h>

#include "shaveDebug.h"
#include "shaveGlobals.h"
#include "shaveIO.h"
#include "shaveMaya.h"
#include "shaveMayaRenderer.h"
#include "shaveRender.h"
#include "shaveRenderCallback.h"


extern void lightBufferFileCleanup();
extern void lightProjectionNodeCleanup();

//	The different order of bytes passed by Maya to renderCallback() is
//	OSX-specific and is not an artifact of the byte-ordering on different
//	processors because both PPC and Intel on OSX use the same ordering,
//	which is different from that used by Linux and Windows.
#ifdef OSMac_
#define ALPHA 0
#define BLUE 1
#define GREEN 2
#define RED 3
#else
#define BLUE 0
#define GREEN 1
#define RED 2
#define ALPHA 3
#endif

#ifndef M_PI
#define M_PI 3.14159926535
#endif


bool shaveRenderCallback::mDoingGeomRender = false;
bool shaveRenderCallback::mDoTiles = false;
const MRenderData* shaveRenderCallback::mMayaRenderData = 0;
const shaveRender::Pixel* shaveRenderCallback::mShavePixels = 0;


void shaveRenderCallback::disableCamRender()
{
	ENTER();
	mDoingCamRender = false;
	LEAVE();
}


void shaveRenderCallback::enableCamRender(const MDagPath& renderCam)
{
	ENTER();
	mDoingCamRender = true;
	mRenderCamera = renderCam;

	LEAVE();
}


shaveRenderCallback::shaveRenderCallback(shaveMayaRenderer* renderer)
:	mDoingCamRender(false)
,	mRenderer(renderer)
{}


#if !defined(max)
inline float max(float first, float second)
{
	return(first >= second ? first : second);	
}
#endif

#if !defined(min)
inline float min(float first, float second)
{
	return(first <= second ? first : second);
}
#endif


bool shaveRenderCallback::shadowCastCallback (const MRenderShadowData &data)
{
	return true;
}


bool shaveRenderCallback::renderCallback (const MRenderData &data)
{
	ENTER();

	if (mDoingGeomRender)
	{
		mDoingGeomRender = false;
	}
	else
	{
		shaveRender::SceneInfo* shaveData = shaveRender::getSceneInfo();
		const shaveGlobals::Globals& globals = shaveRender::getFrameGlobals();

		if (mDoingCamRender && !shaveRenderCancelled)
		{
			initTileCallback(data, shaveData);

			if (shaveRender::renderCameraView(mRenderCamera) != 0)
				shaveRenderCancelled = true;

			cleanupTileCallback();
		}

		if (globals.doCompositing
		&&	!shaveRenderCancelled
		&&	!shaveGlobals::getDefaultNode().isNull())
		{
			float maxSZ = -SHAVE_FAR_CLIP;
			float minSZ = 0.0f;
			float maxMZ = -SHAVE_FAR_CLIP;
			float minMZ = 0.0f;
			float* depthData = data.depthArr;

			if (globals.composite2d)
			{
				unsigned char * imageData = data.rgbaArr;

				shaveMaya::getRenderGlobals();

				shaveRender::Pixel*		sPixel = NULL;
		
				if ((depthData != NULL) && (shaveData->shaveZBuffer != NULL))
				{
					float* MZ;
					float SZ;
					float nSZ;
					float tmp;
					MFloatPoint newPoint;
					for(int y = 0; y < data.ysize; y++)
					{
						for(int x = 0; x < data.xsize; x++)
						{
							MZ = &depthData[data.xsize*y+x];
							SZ = (float)shaveData->shaveZBuffer[
									data.resX*(y + data.bottom) + x + data.left];
							nSZ = shaveZtoMaya(SZ);
							tmp = (*MZ == 0 ? 0:1.0f/(*MZ));
							if(tmp != 0)
							{
								maxMZ = max(maxMZ, tmp);
								minMZ = min(minMZ, tmp);
							}
							if(nSZ != 0)
							{
								maxSZ = max(maxSZ, -SZ);
								minSZ = min(minSZ, -SZ);
							}
							if((nSZ < *MZ && nSZ != 0.0f) || *MZ == 0.0f)
								*MZ = (float)nSZ;
						}
					}
				}

				if (shaveData->shaveRenderPixels)
				{
					unsigned int	carryA;
					unsigned int	carryB;
					unsigned int	carryG;
					unsigned int	carryR;
					int				i;
					unsigned int	iALPHA = ALPHA * data.bytesPerChannel;
					unsigned int	iBLUE = BLUE * data.bytesPerChannel;
					unsigned int	iGREEN = GREEN * data.bytesPerChannel;
					unsigned int	iRED = RED * data.bytesPerChannel;
					unsigned char*	rPixel;
					unsigned int	temp;

					for (int y = 0; y < data.ysize; y++)
					{
						for (int x=0; x < data.xsize; x++)
						{
							rPixel = imageData + ((data.xsize*y+x)*4*data.bytesPerChannel);
							sPixel = &shaveData->shaveRenderPixels[data.resX
									* (y + data.bottom) + x + data.left];

							float sAlpha = (float)sPixel->a / 255.0f;
							float sTransp = 1.0f - sAlpha;

							//
							// Our own version of infinite precision
							// arithmetic.
							//
							carryA = carryB = carryG = carryR = 0;

#if defined(IRIX) || defined(__ppc__)
							for (i = data.bytesPerChannel-1; i >= 0; i--)
#else
							for (i = 0; i < data.bytesPerChannel; i++)
#endif
							{
								temp = (unsigned int)sPixel->r
									+ (unsigned int)(sTransp*(float)rPixel[iRED+i])
									+ carryR;

								carryR = temp >> 8;

								if ((data.bytesPerChannel == 1) && (temp > 255))
									rPixel[iRED+i] = 255;
								else
									rPixel[iRED+i] = (unsigned char)(temp & 0xff);

								temp = (unsigned int)sPixel->g
									+ (unsigned int)(sTransp*(float)rPixel[iGREEN+i])
									+ carryG;

								carryG = temp >> 8;

								if ((data.bytesPerChannel == 1) && (temp > 255))
									rPixel[iGREEN+i] = 255;
								else
									rPixel[iGREEN+i] = (unsigned char)(temp & 0xff);

								temp = (unsigned int)sPixel->b
									+ (unsigned int)(sTransp*(float)rPixel[iBLUE+i])
									+ carryB;

								carryB = temp >> 8;

								if ((data.bytesPerChannel == 1) && (temp > 255))
									rPixel[iBLUE+i] = 255;
								else
									rPixel[iBLUE+i] = (unsigned char)(temp & 0xff);

								temp = (unsigned int)sPixel->a
									+ (unsigned int)rPixel[iALPHA+i]
									+ carryA;

								carryA = temp >> 8;

								if ((data.bytesPerChannel == 1) && (temp > 255))
									rPixel[iALPHA+i] = 255;
								else
									rPixel[iALPHA+i] = (unsigned char)(temp & 0xff);
							}

							//
							// In theory, all of Shave's color values will
							// lie in the range 0 - alpha, so it should be
							// impossible for any of the values to
							// overflow.  However, let's be paranoid: if we
							// got a carry out the MSB of any channel, then
							// set that channel's MSB to its max value of
							// 255.  (We could set all of the bytes in the
							// channel, but MSB should be sufficient.)
							//
#if defined(IRIX) || defined(__ppc__)
							i++;
#else
							i--;
#endif
#if 0
							//
							// This *should* have had the same effect as
							// the more expensive 'temp > 255' checks in
							// the loop above, but Joe says that replacing
							// these with the checks got rid of 'pookies':
							// multicolored artifacts at the ends of
							// semi-transparent hairs.
							//
							if (carryR) rPixel[iRED+i] = 255;
							if (carryG) rPixel[iGREEN+i] = 255;
							if (carryB) rPixel[iBLUE+i] = 255;
							if (carryA) rPixel[iALPHA+i] = 255;
#endif
						}
					}
				}
			}
			//
			// Even if the hair itself is composited into the image within
			// a 3d volumetric shader, we still need to give Maya our depth
			// values so that post-processes such as depth of field will
			// work.
			//
			else
			{
				if ((depthData != NULL)
				&&	(shaveData->shaveZBuffer != NULL)
				&&	globals.doCompositing)
				{
					float* MZ;
					float SZ;
					float nSZ;
					float tmp;

					for(int y = 0; y < data.ysize; y++)
					{
						for(int x = 0; x < data.xsize; x++)
						{
							MZ = &depthData[data.xsize*y+x];
							SZ = (float)shaveData->shaveZBuffer[
									data.resX*(y + data.bottom) + x + data.left];
							nSZ = shaveZtoMaya(SZ);
							tmp = (*MZ == 0 ? 0:1.0f/(*MZ));
							if(tmp != 0)
							{
								maxMZ = max(maxMZ, tmp);
								minMZ = min(minMZ, tmp);
							}
							if(nSZ != 0)
							{
								maxSZ = max(maxSZ, -SZ);
								minSZ = min(minSZ, -SZ);
							}
							if((nSZ < *MZ && nSZ != 0.0f) || *MZ == 0.0f)
								*MZ = (float)nSZ;
						}
					}
				}
			}

			if (globals.verbose)
			{
#ifdef OSMac_
				//
				// Panther broke cerr somehow so that it crashes if you use
				// it to non-zero float or double values.  So let's try our
				// old pal stderr, instead.
				//
				if (stderr)
				{
					fprintf(
						stderr,
						"Z Bounds: Maya: %f, %f, Shave: %f, %f\n",
						-maxMZ, -minMZ, -maxSZ, -minSZ
					);
				}
#else
				cerr << "Z Bounds: Maya: " << -maxMZ << ", " << -minMZ
					 << ", Shave: " << -maxSZ << ", " << -minSZ << endl;
#endif
			}
		}
	}

	if (mRenderer) mRenderer->postCompositeCleanup();

	RETURN(true);
}


bool shaveRenderCallback::postProcessCallback (const MRenderData &data)
{		
	return true;
}


float shaveZtoMaya(float shaveVal)
{
	if(shaveVal == SHAVE_FAR_CLIP)
		return 0.0f;
	else
		return -1.0f/shaveVal;
}


MStatus shaveRenderCallback::getRenderCamera(
		const MRenderData& renderData, MDagPath& renderCamPath
)
{
	//
	// Unfortunately, there's no way to simply ask Maya which camera is
	// currently being rendered.  So instead we have to walk through all
	// available cameras and find out which one's parameters most closely
	// match those which are available to use through the passed-in
	// MRenderData object.
	//
	MItDag		iter(MItDag::kDepthFirst, MFn::kCamera);
	MDagPath	testCamPath;
	float		aspectDiff;
	float		fovDiff;
	const float	kTolerance = 0.001f;
	
	for (; !iter.isDone(); iter.next())
	{
		iter.getPath(testCamPath);

		MFnCamera	cameraFn(testCamPath);

		//
		// Make sure that the perspective/ortho and other settings match.
		//
		aspectDiff = renderData.aspectRatio - (float)cameraFn.aspectRatio();

		fovDiff = renderData.fieldOfView
				- (float)cameraFn.horizontalFieldOfView();

		if ((cameraFn.isOrtho() != renderData.perspective)
		&&	(fabs(aspectDiff) < kTolerance)
		&&	(fabs(fovDiff) < kTolerance))
		{
			//
			// The preliminaries look good, now let's see if the camera's
			// eye point is in the right place.
			//
			MPoint		tempPoint = cameraFn.eyePoint(MSpace::kWorld);

			MFloatPoint	eyePoint(
							(float)tempPoint.x,
							(float)tempPoint.y,
							(float)tempPoint.z
						);

			//
			// For some inexplicable reason, MRenderData::worldToEyeMatrix
			// is an MFloatMatrix rather than a normal MMatrix.  That means
			// that we lose a lot of precision in the calculations below.
			// Enough so that even our relatively generous kTolerance value
			// may end up rejecting valid matches.
			//
			// So let's instead calculate a tolerance which is in keeping
			// with the overall scale of the vectors involved.
			//
			float	scaledTolerance = eyePoint.distanceTo(MFloatPoint::origin)
									/ 1000000.0f;

			//
			// Convert it to the render camera's local space.
			//
			eyePoint *= renderData.worldToEyeMatrix;
			if (eyePoint.isEquivalent(renderData.eyePoint, scaledTolerance))
			{
				//
				// The eye point is right, now let's check the viewing
				// direction.
				//
				MFloatVector	viewDir(cameraFn.viewDirection());

				if (viewDir.isEquivalent(
						renderData.viewDirection, scaledTolerance))
				{
					renderCamPath = testCamPath;

					return MS::kSuccess;
				}
			}
		}
	}

	MGlobal::displayError(
		"shaveRenderCallback::getRenderCamera - could not find the render camera."
	);

	return MS::kFailure;
}


void shaveRenderCallback::initTileCallback(
	const MRenderData& mayaRenderData, shaveRender::SceneInfo* shaveRenderData
)
{
	mMayaRenderData = &mayaRenderData;
	mShavePixels = shaveRenderData->shaveRenderPixels;
	mDoTiles = true;
}


void shaveRenderCallback::cleanupTileCallback()
{
	mMayaRenderData = 0;
	mShavePixels = 0;
	mDoTiles = false;
}


//	This gets called every time Shave has finished rendering a tile.
void shaveRenderCallback::tileRendered(
	unsigned int left, unsigned int right, unsigned int bottom, unsigned int top
)
{
	if (mDoTiles && MRenderView::doesRenderEditorExist())
	{
		//	Find the intersection between Shave's tile and Maya's render
		//	region.
		if (left < mMayaRenderData->left)
			left = mMayaRenderData->left;
		if (right >= (unsigned int)(mMayaRenderData->left + mMayaRenderData->xsize))
			right = mMayaRenderData->left + mMayaRenderData->xsize - 1;
		if (bottom < mMayaRenderData->bottom)
			bottom = mMayaRenderData->bottom;
		if (top >= (unsigned int)(mMayaRenderData->bottom + mMayaRenderData->ysize))
			top = mMayaRenderData->bottom + mMayaRenderData->ysize - 1;

		if ((right < left) || (top < bottom)) return;

		const unsigned xsize = right - left + 1;
		const unsigned ysize = top - bottom + 1;

		//	Create an image buffer for the render view.
		RV_PIXEL*	viewPixels =  new RV_PIXEL[xsize * ysize];

		//---------------------------------------------------------------
		//	Compose the shave pixels onto Maya's pixels and store in the
		//	render view buffer.
		//---------------------------------------------------------------
		const unsigned int	iALPHA = ALPHA * mMayaRenderData->bytesPerChannel;
		const unsigned int	iBLUE = BLUE * mMayaRenderData->bytesPerChannel;
		const unsigned int	iGREEN = GREEN * mMayaRenderData->bytesPerChannel;
		const unsigned int	iRED = RED * mMayaRenderData->bytesPerChannel;
		const unsigned int	mayaBytesPerPixel = 4 * mMayaRenderData->bytesPerChannel;

		//	If there are multiple bytes per channel, we only care about the
		//	most significant byte.
#if defined(__ppc__)
		unsigned int	mayaMSBOffset = 0;
#else
		unsigned int	mayaMSBOffset = mMayaRenderData->bytesPerChannel - 1;
#endif
		unsigned int	mayaByteIdx = 0;
		unsigned char*	mayaBytes = mMayaRenderData->rgbaArr;
		unsigned int	shavePixelIdx = 0;
		float			temp;
		unsigned int	viewPixelIdx = 0;

		for (unsigned int y = bottom; y <= top; ++y)
		{
			shavePixelIdx = mMayaRenderData->resX * y + left;
			mayaByteIdx = (mMayaRenderData->xsize * (y - mMayaRenderData->bottom)
						+ left - mMayaRenderData->left) * mayaBytesPerPixel
						+ mayaMSBOffset;

			for (unsigned int x = left; x <= right; ++x)
			{
				float shaveAlpha = (float)mShavePixels[shavePixelIdx].a / 255.0f;
				float shaveTransp = 1.0f - shaveAlpha;

				temp = (float)mShavePixels[shavePixelIdx].r
						+ shaveTransp * (float)mayaBytes[mayaByteIdx + iRED];
				if (temp > 255.0f) temp = 255.0f;
				viewPixels[viewPixelIdx].r = temp;

				temp = (float)mShavePixels[shavePixelIdx].g
						+ shaveTransp * (float)mayaBytes[mayaByteIdx + iGREEN];
				if (temp > 255.0f) temp = 255.0f;
				viewPixels[viewPixelIdx].g = temp;

				temp = (float)mShavePixels[shavePixelIdx].b
						+ shaveTransp * (float)mayaBytes[mayaByteIdx + iBLUE];
				if (temp > 255.0f) temp = 255.0f;
				viewPixels[viewPixelIdx].b = temp;

				temp = shaveAlpha + (float)mayaBytes[mayaByteIdx + iALPHA];
				if (temp > 255.0f) temp = 255.0f;
				viewPixels[viewPixelIdx].a = temp;

				mayaByteIdx += mayaBytesPerPixel;
				++shavePixelIdx;
				++viewPixelIdx;
			}
		}

		//	Update the render view with the new buffer.
		MRenderView::updatePixels(left, right, bottom, top, viewPixels);
		MRenderView::refresh(left, right, bottom, top);

		delete [] viewPixels;
	}
}