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

#include <maya/MColor.h>
#include <maya/MColorArray.h>
#include <maya/MIntArray.h>
#include <maya/MString.h>
#include <maya/MStringArray.h>

#include <math.h>
#include <string.h>

#include "shaveIO.h"
#include "shaveXPM.h"

shaveXPM::Node*	shaveXPM::Node::head = NULL;
unsigned long	shaveXPM::Node::newMinError;
unsigned long	shaveXPM::Node::nextIndex = 0;
unsigned long	shaveXPM::Node::numLeafNodes = 0;


shaveXPM::Node::Node(
		unsigned char irmin, unsigned char igmin, unsigned char ibmin,
		unsigned char irmax, unsigned char igmax, unsigned char ibmax
)
:	numPixels(0)
,	numLeafPixels(0)
,	sumLeafR(0)
,	sumLeafG(0)
,	sumLeafB(0)
,	quantError(0)
,	rmax(irmax)
,	rmid((irmax + irmin)/2)
,	rmin(irmin)
,	gmax(igmax)
,	gmid((igmax + igmin)/2)
,	gmin(igmin)
,	bmax(ibmax)
,	bmid((ibmax + ibmin)/2)
,	bmin(ibmin)
{
	int	i;

	for (i = 0; i < 8; i++)
		child[i] = NULL;
}


void shaveXPM::Node::init()
{
	numLeafNodes = 0;
	head = new Node(0, 0, 0, 255, 255, 255);
}


void shaveXPM::Node::cleanup()
{
	if (head)
	{
		head->doCleanup();

		delete head;
		head = NULL;
	}
}


void shaveXPM::Node::doCleanup()
{
	unsigned int i;

	for (i = 0; i < 8; i++)
	{
		if (child[i])
		{
			child[i]->doCleanup();
			delete child[i];
			child[i] = NULL;
		}
	}
}


void shaveXPM::Node::addPixel(unsigned char r, unsigned char g, unsigned char b)
{
	head->doAddPixel(r, g, b, 0);
}


void shaveXPM::Node::doAddPixel(
	unsigned char r, unsigned char g, unsigned char b, unsigned short level
)
{
	numPixels++;

	if (level == 7)
	{
		if (numLeafPixels == 0) numLeafNodes++;

		numLeafPixels++;
		sumLeafR += r;
		sumLeafG += g;
		sumLeafB += b;
	}
	else
	{
		//
		// This node's colour space can be divided into 8 sub-spaces by
		// cutting along the midpoint of R, G and B.
		//
		// The pixel should go into one of those 8 sub-spaces, so let's
		// determine which one it falls into.
		//
		int	i = 0;

		if (r > rmid) i += 1;
		if (g > gmid) i += 2;
		if (b > bmid) i += 4;

		//
		// If we haven't yet created a node for that sub-space, do so now.
		//
		if (child[i] == NULL)
		{
			unsigned char crmin = rmin;
			unsigned char crmax = rmid;
			unsigned char cgmin = gmin;
			unsigned char cgmax = gmid;
			unsigned char cbmin = bmin;
			unsigned char cbmax = bmid;

			if (r > rmid)
			{
				crmin = rmid + 1;
				crmax = rmax;
			}

			if (g > rmid)
			{
				cgmin = gmid + 1;
				cgmax = gmax;
			}

			if (b > bmid)
			{
				cbmin = bmid + 1;
				cbmax = bmax;
			}

			child[i] = new Node(crmin, cgmin, cbmin, crmax, cgmax, cbmax);
		}

		child[i]->doAddPixel(r, g, b, level+1);
	}

	//
	// Add this pixel's quantizing error to the overall error for this
	// node.
	//
	int	rd = (int)r - (int)rmid;
	int	gd = (int)g - (int)gmid;
	int	bd = (int)b - (int)bmid;

	quantError += (unsigned long)(rd*rd + gd*gd + bd*bd);
}


unsigned long shaveXPM::Node::prune(unsigned long minError)
{
	newMinError = 0;
	head->doPrune(minError, false);

	return newMinError;
}


bool shaveXPM::Node::doPrune(unsigned long minError, bool force)
{
	unsigned int i;

	//
	// If our quantizing error is within that being pruned, then force all
	// of our children to be pruned as well, since it is possible that they
	// have higher errors and would not otherwise be pruned.
	//
	if (quantError <= minError) force = true;

	for (i = 0; i < 8; i++)
	{
		if (child[i] && child[i]->doPrune(minError, force))
		{
			//
			// The child is being pruned, so add its pixels to our own and
			// delete the child.
			//
			// Note that moving pixels into this node makes it a leaf node.
			// If it was not already a leaf node, then we must increment
			// the leaf node count.
			//
			if (numLeafPixels == 0) numLeafNodes++;

			numLeafPixels += child[i]->numLeafPixels;
			sumLeafR += child[i]->sumLeafR;
			sumLeafG += child[i]->sumLeafG;
			sumLeafB += child[i]->sumLeafB;

			delete child[i];

			child[i] = NULL;

			//
			// Decrement the leaf node count to take account of the child
			// we just deleted.
			//
			numLeafNodes--;
		}
	}

	if (force) return true;

	if ((newMinError == 0) || (quantError < newMinError))
		newMinError = quantError;

	return false;
}


void shaveXPM::Node::getColours(struct Colour* colours)
{
	nextIndex = 0;
	head->doGetColours(colours);
}


void shaveXPM::Node::doGetColours(struct Colour* colours)
{
	if (numLeafPixels > 0)
	{
		index = nextIndex++;

		unsigned int	round = numLeafPixels / 2;

		colours[index].r = (unsigned char)((sumLeafR + round) / numLeafPixels);
		colours[index].g = (unsigned char)((sumLeafG + round) / numLeafPixels);
		colours[index].b = (unsigned char)((sumLeafB + round) / numLeafPixels);
	}

	unsigned int i;

	for (i= 0; i < 8; i++)
		if (child[i]) child[i]->doGetColours(colours);
}


unsigned long shaveXPM::Node::getIndex(
	unsigned char r, unsigned char g, unsigned char b
)
{
	return head->doGetIndex(r, g, b);
}


unsigned long shaveXPM::Node::doGetIndex(
	unsigned char r, unsigned char g, unsigned char b
)
{
	int	i = 0;

	if (r > rmid) i += 1;
	if (g > gmid) i += 2;
	if (b > bmid) i += 4;

	if (child[i] == NULL) return index;

	return child[i]->doGetIndex(r, g, b);
}


MString shaveXPM::toHex(unsigned char val)
{
	static char*	hexChars = "0123456789abcdef";
	static char		hexStr[3];

	hexStr[0] = hexChars[val >> 4];
	hexStr[1] = hexChars[val & 0x0f];
	hexStr[2] = '\0';

	return MString(hexStr);
}


bool shaveXPM::write(
	int width, int height, const unsigned char* rgba, const MString& filename
)
{
	ofstream	f(filename.asChar());

	if (!f.good()) return false;

	//
	// Which chars can we use for colour indices in the XPM file?
	//
	const char*	validChars =	"abcdefghijklmnopqrstuvwxyz"
								"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
								"0123456789"
								"-_=+[]{}';:|/?.,>< `~!@#$%^&*()";
	unsigned int	numValidChars = (unsigned int)strlen(validChars);

	//
	// Sort the pixels in the image into an 8-way colour tree.
	//
	int	numPixels = width * height;
	int	numPixelElements = numPixels * 4;
	int	i;

	Node::init();

	for (i = 0; i < numPixelElements; i += 4)
	{
		Node::addPixel(rgba[i], rgba[i+1], rgba[i+2]);
	}

	//
	// The XPM funcs that Maya uses on Windows don't allow more than 2
	// chars per colour index.  So that marks the limit on the number of
	// colours we can use.
	//
	unsigned long	maxColours = numValidChars * numValidChars;
	unsigned long	minError = 0;

	//
	// Iteratively prune the colour tree until we're under our limit.
	//
	while (Node::numColours() > maxColours)
	{
		minError = Node::prune(minError);
	}

	unsigned int	numColours = Node::numColours();

	struct Colour* colours = new struct Colour[numColours];

	Node::getColours(colours);

	//
	// Get colour indices for all of the pixels in the image.
	//
	unsigned int*	indices = new unsigned int[numPixels];

	for (i = 0; i < numPixelElements; i += 4)
	{
		indices[i >> 2] = Node::getIndex(rgba[i], rgba[i+1], rgba[i+2]);
	}

	Node::cleanup();

	//
	// Output XPM header info.
	//
	// Note that on Windows, Maya will reject an XPM file if the opening
	// brace is not on the same line as the 'static char' declaration.  I
	// kid you not.
	//
	f << "/* XPM */" << endl;
	f << "static char* swatch[] = {" << endl;

	f << "/* width height ncolours charsPerPixel */" << endl;
	f << "\"" << width << " " << height << " " << numColours << " 2\"," << endl;

	//
	// Generate and output the colour database.  That is, the association
	// between character sequences and the colours they will represent in
	// the 'pixels' section.
	//
	MStringArray	colourCodes;
	char			code[3];
	f << "/* colours */" << endl;

	for (i = 0; i < (int)numColours; i++)
	{
		unsigned int	index = i;

		code[0] = validChars[index / numValidChars];
		code[1] = validChars[index % numValidChars];
		code[2] = '\0';
		colourCodes.append(code);

#if defined(OSMac_) && (MAYA_API_VERSION >= 201600) && (MAYA_API_VERSION < 201700)
		f << "\"" << code << " c #" << toHex(colours[i].r).asChar()
		  << toHex(colours[i].g).asChar() << toHex(colours[i].b).asChar()
		  << "\"," << endl;
#else
		f << "\"" << code << " c #" << toHex(colours[i].r)
		  << toHex(colours[i].g) << toHex(colours[i].b) << "\"," << endl;
#endif
	}

	delete [] colours;

	//
	// Output the pixel data.  Note that image data starts at the *bottom*
	// of the image, but XPM format starts at the top.  So we have to flip
	// things around.
	//
	int	col;

	f << "/* pixels */" << endl;

	for (i = numPixels - width; i >= 0; i -= width)
	{
		f << "\"";

		for (col = 0; col < width; col++) {
#if defined(OSMac_) && (MAYA_API_VERSION >= 201600) && (MAYA_API_VERSION < 201700)
			f << colourCodes[indices[i+col]].asChar();
#else
			f << colourCodes[indices[i+col]];
#endif
		}

		f << "\"," << endl;
	}

	f << "};" << endl;

	f.close();

	return true;
}