/*********************************************************/
/* This source code copyright (c) 1991-2001, Aaron Giles */
/* See the Read Me file for licensing information.       */
/* Contact email: mac@aarongiles.com                     */
/*********************************************************/

//=====================================================================================
// Generic includes for Macintosh headers
//=====================================================================================

#if THINK_C
#include "THINK.Header"
#include <stddef.h>
#elif applec
#pragma load ":Headers:MPW.Header"
#elif __MWERKS__
//#include "MW.Header"
#else
#include "JPEGView.h"
#endif

//=====================================================================================
// Includes specific to this module
//=====================================================================================

#include "Quantize.h"

//=====================================================================================
// Constants local to this module
//=====================================================================================

enum {
	kRedScale = 4,
	kGreenScale = 5,
	kBlueScale = 3
};

//=====================================================================================
// Box: structure for holding information about color reduction boxes
//=====================================================================================

typedef struct Box {
	ushort c0min, c0max, c1min, c1max, c2min, c2max;
	ulong colorCount;
} Box, *BoxPtr, **BoxHandle;

//=====================================================================================
// FillData: structure for passing information between C and assembly portions of the
// histogram filling routine
//=====================================================================================

typedef struct FillData {
	ulong height, width;
	ulong rowBytes;
	ushort *histogram;
	uchar *baseAdr;
} FillData, *FillDataPtr;

//=====================================================================================
// Global variables local to this module
//=====================================================================================

static ushort gNumBoxes, *gHistogram;
static Box gBoxList[256];

//=====================================================================================
// Prototypes for functions local to this module
//=====================================================================================

static BoxPtr FindBiggestPixelPop(void);
static BoxPtr FindBiggestVolume(void);
static void UpdateBox(BoxPtr boxp);
static void MedianCut(short desiredColors);
static void ComputeColor(BoxPtr boxp, int icolor, PaletteHandle thePalette);

#if THINK_C || !USE_ASM
static void DoFillHistogram32(FillDataPtr theData);
static void DoFillHistogram16(FillDataPtr theData);
#else
extern void DoFillHistogram32(FillDataPtr theData);
extern void DoFillHistogram16(FillDataPtr theData);
#endif

//=====================================================================================
// OSErr DoQuantize(ImageHandle theImage, long numColors, PaletteHandle *thePalette)
//
// Calculates two-pass color quantization for a given image, reducing it to the
// specified number of colors and returning the result as a palette in the specified
// handle.
//
// External calls: KeepSpinning, CheckedNewHandle, ClearMem, DrawImage, KillGWorld
//=====================================================================================

extern OSErr DoQuantize(ImageHandle theImage, long numColors, PaletteHandle *thePalette,
	NestedProgressPtr progProc)
{
	Boolean hadGWorld = ((*theImage)->gworld != nil);
	OSErr theErr = noErr;
	Handle theHandle;
	Str32 theNum;

	*thePalette = nil;
	KeepSpinning();
	if (theHandle = CheckedNewHandle(32768L * sizeof(short), true)) {
		NumToString(numColors, theNum);
		SetSlideControlsText(gString[strSlideReducing], (*theImage)->file.name, gNullString, gNullString);
		UpdateParamText(gString[strReducingTo], theNum, gNullString, (*theImage)->file.name);
		HLockHi(theHandle);
		ClearMem(*theHandle, 32768L * sizeof(short));
		gHistogram = (ushort *)StripAddress(*theHandle);
		if (!GWOrigSize(theImage)) KillGWorld(theImage);
		KeepSpinning();
		theErr = DrawImage(theImage, bfQuantize + (gSlideShow ? bfUseSlideProgress : 0),
			progProc);
		HUnlock(theHandle);
		if (theErr == noErr) 
			theErr = QuantizeFromHistogram(theHandle, numColors, thePalette);
		DisposeHandle(theHandle);
	} else gIntError = errNoQuantMemory, theErr = memFullErr;
	if (theErr != noErr) {
		if (!hadGWorld && (*theImage)->gworld) KillGWorld(theImage);
	}
	return theErr;
}

//=====================================================================================
// QuantizeFromHistogram(Handle theHistogram, short numColors, 
//     PaletteHandle *thePalette)
//
// Performs two-pass color reduction, given a histogram of color counts.
//
// External calls: KeepSpinning
//=====================================================================================

extern OSErr QuantizeFromHistogram(Handle theHistogram, long numColors, 
	PaletteHandle *thePalette)
{
	char hState = HGetState(theHistogram);
	OSErr theErr = noErr;
	RGBColor theColor;
	ushort i;

	HLock(theHistogram);
	gHistogram = (ushort *)StripAddress(*theHistogram);
	*thePalette = nil;
	if (*thePalette = NewPalette(numColors, nil, pmTolerant, 0)) {
		KeepSpinning();
		gNumBoxes = 1;
		gBoxList[0].c0min = gBoxList[0].c1min = gBoxList[0].c2min = 0;
		gBoxList[0].c0max = gBoxList[0].c1max = gBoxList[0].c2max = 31;
		UpdateBox(&gBoxList[0]);
		KeepSpinning();
		MedianCut(numColors - 2);
		KeepSpinning();
		for (i = 0; i < gNumBoxes; i++) ComputeColor(&gBoxList[i], i, *thePalette);
		KeepSpinning();
		for ( ; i < numColors; i++) {
			GetEntryColor(gGreyPalette[8], i, &theColor);
			SetEntryColor(*thePalette, i, &theColor);
		}
		KeepSpinning();
	} else gIntError = errNoQuantMemory, theErr = memFullErr;
	HSetState(theHistogram, hState);
	return theErr;
}

//=====================================================================================
// void ShrinkGWorld(ImageHandle theImage)
//
// Attempts to create a properly-sized GWorld from an overly-deep GWorld.
//
// External calls: PushPort, SlideProgress, MySetPort, MakeImageGWorld, KillGWorld,
//		PopPort, NewBits, KeepSpinning, OldBits
//=====================================================================================

extern OSErr ShrinkGWorld(ImageHandle theImage)
{
	GWorldPtr oldGWorld = (*theImage)->gworld, newGWorld;
	JVDrawParamsHandle theParams;
	Rect srcRect, dstRect;
	OSErr theErr = noErr;

	// first we try to allocate a proper-sized GWorld
	KeepSpinning();
	(*theImage)->gworld = nil;
	(*theImage)->flags &= ~ifGWOrigSize;
	if (!MakeImageGWorld(theImage, false)) {
	
		// if it didn't work, restore things and kill the oversized GWorld
		(*theImage)->gworld = oldGWorld;
		KillGWorld(theImage);
		return noErr;
	}
	
	// calculate source and destination rectangles
	srcRect = (GWOrigSize(theImage)) ? (*theImage)->crect : oldGWorld->portRect;
	dstRect = (*theImage)->gworld->portRect;

	// allocate the drawing parameters
	if (theParams = NewDrawParams(&(*theImage)->gworld->portRect, (*theImage)->quality,
				false, nil, (*theImage)->privateData)) {

		// set up the destination port and preflight our drawing
		SetUpDrawPort(theParams, kOffscreenPort1, (*theImage)->gworld, &srcRect, &dstRect, true);
		theErr = PreflightDrawing(theParams);
		
		// if everything worked, copy from GWorld to GWorld, and postflight when done
		if (theErr == noErr) {
			SpinIndef();
			CopyGWorldToGWorld(oldGWorld, (*theParams)->dummy, &srcRect, &srcRect, srcCopy, nil);
			theErr = PostflightDrawing(theParams);
		}

		// now kill the old GWorld through the proper channels
		KeepSpinning();
		newGWorld = (*theImage)->gworld;
		(*theImage)->gworld = oldGWorld;
		KillGWorld(theImage);
		(*theImage)->gworld = newGWorld;
		
		// dispose of everything and clean up
		DisposeHandle((Handle)theParams);
	} else theErr = MemError();
	return theErr;
}

//=====================================================================================
// void FillHistogram(PixMapPtr thePixMap)
//
// Fills the local histogram according to the data in the specified PixMap.  Works
// for 16-bit and 32-bit deep PixMaps.  Both C and 68020 assembly versions are
// provided.
//
// External calls: none
//=====================================================================================

extern void FillHistogram(PixMapPtr thePixMap)
{
	register short theDepth = thePixMap->pixelSize;
	char mmuMode = true32b;
	FillData theData;

	theData.height = Height(&thePixMap->bounds);
	theData.width = Width(&thePixMap->bounds);
	theData.rowBytes = thePixMap->rowBytes & 0x3fff;
	theData.histogram = gHistogram;
	if (thePixMap->pmVersion == 4) theData.baseAdr = (uchar *)thePixMap->baseAddr;
	else theData.baseAdr = (uchar *)StripAddress(thePixMap->baseAddr);
	SwapMMUMode(&mmuMode);
	switch (theDepth) {
		case 32:
			DoFillHistogram32(&theData);
			break;
		case 16:
			DoFillHistogram16(&theData);
			break;
	}
	SwapMMUMode(&mmuMode);
}

//=====================================================================================
// BoxPtr FindBiggestPixelPop(void)
//
// Searches the box list for the box with the largest number of pixels.
//
// External calls: none
//=====================================================================================

static BoxPtr FindBiggestPixelPop(void)
{
	register ulong max = 0;
	register BoxPtr boxp;
	BoxPtr which = nil;
	register ushort i;
	
	for (i = 0, boxp = gBoxList; i < gNumBoxes; i++, boxp++)
		if (boxp->colorCount > max)
			if ((boxp->c0max > boxp->c0min) || 
				(boxp->c1max > boxp->c1min) ||
				(boxp->c2max > boxp->c2min)) {
				which = boxp;
				max = boxp->colorCount;
			}
	return which;
}

//=====================================================================================
// BoxPtr FindBiggestVolume(void)
//
// Searches the box list for the box with the largest volume (actually largest norm).
//
// External calls: none
//=====================================================================================

static BoxPtr FindBiggestVolume(void)
{
	register ulong norm, c0, c1, c2;
	register ulong max = 0;
	register BoxPtr boxp;
	BoxPtr which = nil;
	register ushort i;
	
	for (i = 0, boxp = gBoxList; i < gNumBoxes; i++, boxp++) {
		c0 = (boxp->c0max - boxp->c0min) * kRedScale;
		c1 = (boxp->c1max - boxp->c1min) * kGreenScale;
		c2 = (boxp->c2max - boxp->c2min) * kBlueScale;
		norm = c0*c0 + c1*c1 + c2*c2;
		if (norm > max) {
			which = boxp;
			max = norm;
		}
	}
	return which;
}

//=====================================================================================
// void UpdateBox(BoxPtr boxp)
//
// Shrinks a given box to its minimum size and recalculates the number of pixels
// contained therein.
//
// External calls: none
//=====================================================================================

static void UpdateBox(BoxPtr boxp)
{
	short c0min, c0max, c1min, c1max, c2min, c2max;
	register short c0, c1, c2;
	register ushort *hist;
	ulong ccount;
	
	c0min = boxp->c0min;	c0max = boxp->c0max;
	c1min = boxp->c1min;	c1max = boxp->c1max;
	c2min = boxp->c2min;	c2max = boxp->c2max;

	if (c0max > c0min)
		for (c0 = c0min; c0 <= c0max; c0++)
			for (c1 = c1min; c1 <= c1max; c1++) {
				hist = gHistogram + (((c0 << 5) + c1) << 5) + c2min;
				for (c2 = c2min; c2 <= c2max; c2++)
					if (*hist++) {
						boxp->c0min = c0min = c0;
						goto haveC0Min;
					}
			}
haveC0Min:
	if (c0max > c0min)
		for (c0 = c0max; c0 >= c0min; c0--)
			for (c1 = c1min; c1 <= c1max; c1++) {
				hist = gHistogram + (((c0 << 5) + c1) << 5) + c2min;
				for (c2 = c2min; c2 <= c2max; c2++)
					if (*hist++) {
						boxp->c0max = c0max = c0;
						goto haveC0Max;
					}
			}
haveC0Max:
	if (c1max > c1min)
		for (c1 = c1min; c1 <= c1max; c1++)
			for (c0 = c0min; c0 <= c0max; c0++) {
				hist = gHistogram + (((c0 << 5) + c1) << 5) + c2min;
				for (c2 = c2min; c2 <= c2max; c2++)
					if (*hist++) {
						boxp->c1min = c1min = c1;
						goto haveC1Min;
					}
			}
haveC1Min:
	if (c1max > c1min)
		for (c1 = c1max; c1 >= c1min; c1--)
			for (c0 = c0min; c0 <= c0max; c0++) {
				hist = gHistogram + (((c0 << 5) + c1) << 5) + c2min;
				for (c2 = c2min; c2 <= c2max; c2++)
					if (*hist++) {
						boxp->c1max = c1max = c1;
						goto haveC1Max;
					}
			}
haveC1Max:
	if (c2max > c2min)
		for (c2 = c2min; c2 <= c2max; c2++)
			for (c0 = c0min; c0 <= c0max; c0++) {
				hist = gHistogram + (((c0 << 5) + c1min) << 5) + c2;
				for (c1 = c1min; c1 <= c1max; c1++, hist += 32)
					if (*hist) {
						boxp->c2min = c2min = c2;
						goto haveC2Min;
					}
			}
haveC2Min:
	if (c2max > c2min)
		for (c2 = c2max; c2 >= c2min; c2--)
			for (c0 = c0min; c0 <= c0max; c0++) {
				hist = gHistogram + (((c0 << 5) + c1min) << 5) + c2;
				for (c1 = c1min; c1 <= c1max; c1++, hist += 32)
					if (*hist) {
						boxp->c2max = c2max = c2;
						goto haveC2Max;
					}
			}
haveC2Max:
	ccount = 0;
	for (c0 = c0min; c0 <= c0max; c0++)
		for (c1 = c1min; c1 <= c1max; c1++) {
			hist = gHistogram + (((c0 << 5) + c1) << 5) + c2min;
			for (c2 = c2min; c2 <= c2max; c2++)
				ccount += *hist++;
		}
	boxp->colorCount = ccount;
}

//=====================================================================================
// void MedianCut(short desiredColors)
//
// Repeatedly divides the original box into smaller boxes until we have as many boxes
//     as there are desired colors.
//
// External calls: KeepSpinning
//=====================================================================================

static void MedianCut(short desiredColors)
{
	long c0min, c0max, c1min, c1max, c2min, c2max;
	long c0, c1, c2, cmax;
	register BoxPtr b1, b2;
	ushort n, lb;
	
	while (gNumBoxes < desiredColors) {
		KeepSpinning();
		if ((gNumBoxes << 1) <= desiredColors) b1 = FindBiggestPixelPop();
		else b1 = FindBiggestVolume();
		if (!b1) break;
		c0min = b1->c0min;	c0max = b1->c0max;
		c1min = b1->c1min;	c1max = b1->c1max;
		c2min = b1->c2min;	c2max = b1->c2max;
		b2 = &gBoxList[gNumBoxes];
		b2->c0max = c0max; b2->c1max = c1max; b2->c2max = c2max;
		b2->c0min = c0min; b2->c1min = c1min; b2->c2min = c2min;
		c0 = (c0max - c0min) * kRedScale;
		c1 = (c1max - c1min) * kGreenScale;
		c2 = (c2max - c2min) * kBlueScale;
		cmax = c0; n = 0;
		if (c1 > cmax) { cmax = c1; n = 1; }
		if (c2 > cmax) n = 2;
		switch (n) {
			case 0:
				lb = (c0max + c0min) >> 1;
				b1->c0max = lb;
				b2->c0min = lb + 1;
				break;
			case 1:
				lb = (c1max + c1min) >> 1;
				b1->c1max = lb;
				b2->c1min = lb + 1;
				break;
			case 2:
				lb = (c2max + c2min) >> 1;
				b1->c2max = lb;
				b2->c2min = lb + 1;
				break;
		}
		UpdateBox(b1);
		UpdateBox(b2);
		gNumBoxes++;
	}
}

//=====================================================================================
// void ComputeColor(BoxPtr boxp, int icolor, PaletteHandle thePalette)
//
// Calculates the final RGB color by doing a weighted average over pixels in the 
//     given box.
//
// External calls: none
//=====================================================================================

static void ComputeColor(BoxPtr boxp, int icolor, PaletteHandle thePalette)
{
	ulong count, total = 0, c0total = 0, c1total = 0, c2total = 0;
	short c0min, c0max, c1min, c1max, c2min, c2max;
	short c0, c1, c2, r, g, b;
	register ushort *hist;
	RGBColor theColor;
	
	if (!icolor) {
		theColor.red = theColor.blue = theColor.green = 0xffff;
		SetEntryColor(thePalette, 0, &theColor);
		theColor.red = theColor.blue = theColor.green = 0;
		SetEntryColor(thePalette, 1, &theColor);
	}
	
	c0min = boxp->c0min;	c0max = boxp->c0max;
	c1min = boxp->c1min;	c1max = boxp->c1max;
	c2min = boxp->c2min;	c2max = boxp->c2max;

	for (c0 = c0min; c0 <= c0max; c0++)
		for (c1 = c1min; c1 <= c1max; c1++) {
			hist = gHistogram + (((c0 << 5) + c1) << 5) + c2min;
			for (c2 = c2min; c2 <= c2max; c2++)
				if (count = *hist++) {
					total += count;
					c0total += ((c0 << 3) + 4) * count;
					c1total += ((c1 << 3) + 4) * count;
					c2total += ((c2 << 3) + 4) * count;
				}
		}

	if (total) {
		r = (c0total + (total >> 1)) / total;
		g = (c1total + (total >> 1)) / total;
		b = (c2total + (total >> 1)) / total;
		theColor.red = ((ushort)r << 8) | (ushort)r;
		theColor.green = ((ushort)g << 8) | (ushort)g;
		theColor.blue = ((ushort)b << 8) | (ushort)b;
	} else theColor.red = theColor.green = theColor.blue = 0;
	
	SetEntryColor(thePalette, icolor + 2, &theColor);
}

#if !USE_ASM

//=====================================================================================
// Generic C substitutes for the assembly functions defined below.
//=====================================================================================

static void DoFillHistogram32(FillDataPtr theData)
{
	#include "FillHistogram32.c"
}

static void DoFillHistogram16(FillDataPtr theData)
{
	#include "FillHistogram16.c"
}

#elif THINK_C

//=====================================================================================
// Assembly language interfaces for THINK C; here we define the offsets of all fields
// within the CopyData structure, and then provide C function wrappers around asm { }
// statements that #include the actual assembly code.  In MPW, we just need to compile
// Quantize.a to get the assembly routines.
//=====================================================================================

#define height offsetof(FillData, height)		/*    0 */
#define width offsetof(FillData, width)			/*    4 */
#define rowBytes offsetof(FillData, rowBytes)	/*    8 */
#define histogram offsetof(FillData, histogram)	/*   12 */
#define baseAdr offsetof(FillData, baseAdr)		/*   16 */

static void DoFillHistogram32(FillDataPtr theData)
{
	asm {
		#include "FillHistogram32.a"
	}
}

static void DoFillHistogram16(FillDataPtr theData)
{
	asm {
		#include "FillHistogram16.a"
	}
}

#endif