// --------------- MULTI-CHANNEL ROUTINES ---------------
// wavelet image compression
// channel balancing & quality enforcement
// (c) 2001-2002 Daniel Vollmer (maven@maven.de)

/* 	This program is free software; you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation; either version 2 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program; if not, write to the Free Software
	Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA	02111-1307	USA
*/

#include "wavelet.h"
#include "mchannel.h"


// colour-space conversion (using 16.16 fixed point)
// this is the same transform as used in JPEG, and the results still need
// to be rounded and shifted / divided

#define RGB_Y(r, g, b) ((r) * 19595 + (g) * 38470 + (b) * 7471)
#define RGB_Cb(r, g, b) ((r) * -11059 + (g) * -21709 + (b) * 32768)
#define RGB_Cr(r, g, b) ((r) * 32768 + (g) * -27439 + (b) * -5329)

// the results of these need to be clamped to the range [0..255] after rounding
// and shifting
#define YCbCr_R(y, cb, cr) (/*(y) * 65536 + (cb) * 0 +*/ (cr) * 91881)
#define YCbCr_G(y, cb, cr) (/*(y) * 65536 +*/ (cb) * -22554 + (cr) * -46802)
#define YCbCr_B(y, cb, cr) (/*(y) * 65536 +*/ (cb) * 116130 /*+ (cr) * 0*/)

//----------------------------------------------------------------------------
// wv_rgb_to_ycbcr	converts an array of pixels
//				(given as seperate bitplanes) from RGB into YCbCr
//				returns the number of pixels converted
// Num			numbers of consecutive pixels to be converted
// R, G, B		pointers to the red, green and blue arrays, respectively (R->Y, G->Cb, B->Cr)
//----------------------------------------------------------------------------
WAVELET_DLL_API int WAVELET_DLL_CC wv_rgb_to_ycbcr(const int Num, wv_pel* R, wv_pel* G, wv_pel* B)
{
	if (Num > 0 && R && G && B)
	{
		int i;

		for (i = 0; i < Num; i++)
		{
			int rr = R[i], gg = G[i], bb = B[i];

			R[i] = (wv_pel)((RGB_Y(rr, gg, bb) + 32768) >> 16); // Y
			G[i] = (wv_pel)(((RGB_Cb(rr, gg, bb) + 32767) >> 16)); // Cb
			B[i] = (wv_pel)(((RGB_Cr(rr, gg, bb) + 32767) >> 16)); // Cr
		}
		return Num;
	}
	return 0;
}


//----------------------------------------------------------------------------
// wv_ycbcr_to_rgb	converts an array of pixels
//				(given as seperate bitplanes) from YCbCr into RGB
//				returns the number of pixels converted
// Num			numbers of consecutive pixels to be converted
// Y, Cb, Cr	pointers to the Y, Cb, and Cr arrays, respectively (Y->R, Cb->G, Cr->B)
//----------------------------------------------------------------------------
WAVELET_DLL_API int WAVELET_DLL_CC wv_ycbcr_to_rgb(const int Num, wv_pel* Y, wv_pel* Cb, wv_pel* Cr)
{
	if (Num > 0 && Y && Cb && Cr)
	{
		int i;

		for (i = 0; i < Num; i++)
		{
			int yy = Y[i], cb = Cb[i], cr = Cr[i];
			
			Y[i] = yy + ((YCbCr_R(yy, cb, cr) + 32768) >> 16); Y[i] = min(255, max(0, Y[i])); // R
			Cb[i] = yy + ((YCbCr_G(yy, cb, cr) + 32768) >> 16); Cb[i] = min(255, max(0, Cb[i])); // G
			Cr[i] = yy + ((YCbCr_B(yy, cb, cr) + 32768) >> 16); Cr[i] = min(255, max(0, Cr[i])); // B
		}
		return Num;
	}
	return 0;
}


//----------------------------------------------------------------------------
// wv_init_multi_channels	tries to find the matching settings for
//				a number of channels (i.e. satisfy all the constraints given)
//				returns the bits needed (never larger than MaxBits),
//				0 on failure to satisfy contraints
// MaxBits		number of bits to be used by all the channels,
//				or 0 for quality based encoding
// Threshold	how closely it tries to equalise quality (of the channels),
//				should be inbetween 0.0f..1.0f for good results, use 0.0f
//				if you don't mind waiting slightly longer...
// NumChannels	# of channels given in Channels to be compressed
// Channels		is the array of channel-parameters with NumChannels entries
//				(need a valid t_wv_cchannel and max_mse),
//				when encoding for quality each channel's max_mse gives the
//				maximum target error,
//				when encoding for size each channel's max_mse gives a
//				a target error relative to the other channels
// Sets			array with NumChannels entries where the best settings are
//				stored (on success). These have to be freed by the calling
//				application using wv_done_channel_settings.
//----------------------------------------------------------------------------
WAVELET_DLL_API int WAVELET_DLL_CC wv_init_multi_channels(const int MaxBits, const float Threshold, const int NumChannels, t_wv_mchannel_params* Channels, t_wv_csettings** Sets)
{
	int bits = 0;

	if (NumChannels > 0 && Channels && Sets)
	{
		int i;

		for (i = 0; i < NumChannels; i++)
			Sets[i] = NULL;
		if (MaxBits <= 0)
		{ // simply enforce mse constraints
			for (i = 0; i < NumChannels; i++)
			{
				bits = wv_init_channel_settings(Channels[i].cc, 0, Channels[i].max_mse, &Sets[i]);
				if (bits <= 0)
					break;
			}
		}
		else
		{ // balance the channels so that MSEs are about equal (including max_mse interpreted as delta_mse)
			int num_iter = 0, adjustment;
			int omin_idx = -1, omax_idx = -1;
			t_wv_csettings* lsets[wv_MAX_CHANNELS];
			float closest_match = NumChannels * 65536.0f;

			for (i = 0; i < NumChannels; i++)
			{
				Channels[i].bits_unused = Channels[i].old_bits = 0;
				Sets[i] = NULL;
				lsets[i] = NULL;
			}
		
			// distribute the bits equally as inital state
			adjustment = 0;
			for (i = 1; i < NumChannels; i++)
			{
				Channels[i].bits_alloced = MaxBits / NumChannels;
				adjustment += Channels[i].bits_alloced;
			}
			Channels[0].bits_alloced = MaxBits - adjustment; // gets all the bits left over
			adjustment = 1;
			
			do
			{
				float bmin_mse = 65536.0f, bmax_mse = -1.0f;
				float avg_mse = 0.0f;
				int bmin_idx = -1, bmax_idx = -1;
				int bits_available;

				for (i = 0; i < NumChannels; i++)
				{
					if (lsets[i] && Channels[i].bits_alloced != Channels[i].old_bits)
					{
						wv_done_channel_settings(lsets[i]);
						lsets[i] = NULL;
					}
					if (!lsets[i])
						bits = wv_init_channel_settings(Channels[i].cc, Channels[i].bits_alloced, 65536.0f, &lsets[i]);
					else
						bits = lsets[i]->num_bits;
					if (bits > 0 && lsets[i])
					{
//						fprintf(f, "%i bits_a: %i bits_u: %i error: %f\n", i, Channels[i].bits_alloced, lsets[i]->num_bits, lsets[i]->emse);
						Channels[i].old_bits = bits; // as we are going to redistribute the unused ones							
						Channels[i].bits_unused = Channels[i].bits_alloced - bits;
						// find the best and the worst channel (including relative error adjustments)
						if (lsets[i]->emse - Channels[i].max_mse <= bmin_mse)
						{
							bmin_mse = lsets[i]->emse - Channels[i].max_mse;
							bmin_idx = i;
						}
						if (lsets[i]->emse - Channels[i].max_mse >= bmax_mse)
						{
							bmax_mse = lsets[i]->emse - Channels[i].max_mse;
							bmax_idx = i;
						}
						avg_mse += lsets[i]->emse;
					}
					else
						break;
				}
						
				if (i != NumChannels || bmin_idx == -1 || bmax_idx == -1)
				{ // uh-oh, shouldn't happen
					bits = 0;
					break; // bits == 0 -> abort
				}
/*				// equalise difference between best & worst mse
				if (bmax_mse - bmin_mse <= closest_match)
				{ // this is better than our best-one so far, so keep it
					for (i = 0; i < NumChannels; i++)
					{
						Sets[i] = lsets[i]; // copy the set
						lsets[i] = NULL; // and don't let it be deleted
					}
					closest_match = bmax_mse - bmin_mse;
				}*/
				if (avg_mse <= closest_match)
				{ // this is better than our best-one so far, so keep it
					for (i = 0; i < NumChannels; i++)
					{
						Sets[i] = lsets[i]; // copy the set
						lsets[i] = NULL; // and don't let it be deleted
					}
					closest_match = avg_mse;
				}
				if ((bmin_idx == bmax_idx) || (bmax_mse - bmin_mse <= Threshold))
					break; // balanced or only one channel

				if (bmin_idx == omax_idx && bmax_idx == omin_idx)
				{
					adjustment++;
					if (adjustment > 16)
						break; // only flipping between two channels
				}

				// redistribute bits							
				bits_available = 0; // add up all the unused bits
				for (i = 0; i < NumChannels; i++)
					if (i != bmin_idx && i != bmax_idx)
					{
						bits_available += Channels[i].bits_unused;
						Channels[i].bits_alloced -= Channels[i].bits_unused;
					}
				
				// redistribute the unused bits, plus the one from the best & worst channel among them
				bits_available += Channels[bmin_idx].bits_alloced + Channels[bmax_idx].bits_alloced;

				Channels[bmin_idx].bits_alloced = (Channels[bmin_idx].bits_alloced * adjustment) / (adjustment + 1);
				Channels[bmax_idx].bits_alloced = bits_available - Channels[bmin_idx].bits_alloced; // gets the rest

				omin_idx = bmin_idx;
				omax_idx = bmax_idx;

				num_iter++;
			} while (num_iter < 256);

			// free the settings (we've copied the best ones to Sets)
			for (i = 0; i < NumChannels; i++)
				if (lsets[i] != NULL)
					wv_done_channel_settings(lsets[i]);
		}
		
		// recompute the final number of bits used
		bits = 0;
		for (i = 0; i < NumChannels; i++)
			if (Sets[i])
				bits += Sets[i]->num_bits;
	}
	return bits;
}