/*----------------------------------------------------------------------------*/	
/*----------------------------------------------------------------------------*/	
/* EZW2 - Embedded Zerotree Wavelet Coder
 * 
 * Author : Mow-Song, Ng
 * Data   : 20-07-2002
 *
 * Last update : 31-08-2002
 *
 * A large portion/idea of the codes come from:
 * C. Valens - Embedded Zerotree Wavelet Encoder Tutorial.
 * G. Davis - Baseline Wavelet Transform Coder Construction Kit
 *
 * The authors copyrights to the codes are acknowledged.
 * The EZW coder is patented by Jerome Shapiro. As far as I know, these are
 * the patent numbers 
 * 5315670 (Issued May 24, 1994), 5321776 (Issued June 14, 1994) and 
 * 5412741 (Issued May 2, 1995).
 *
 *
 * Well, I wrote the coder with some one else codes, add some lines, remove 
 * some lines, sometimes just verbatim. Any way, if you feel that I am worth
 * the acknowledgement, please do so, or perhaps drop me an e-mail.
 *
 * Please refer to my website for more info.
 *
 * My contact:
 * msng@mmu.edu.my
 * http://www.pesona.mmu.edu.my/~msng
 *
 */
/*----------------------------------------------------------------------------*/	
/*----------------------------------------------------------------------------*/	

#include "ezw.h"

/*----------------------------------------------------------------------------*/	
/*----------------------------------------------------------------------------*/	
void EZW2Decode(WTRANSFORM *transform, ArithDecoder *decoder, int DecodeBytes,
					 int Threshold, int HisCount, int ScanOrder);

void DecodeDominantPassMortonScan(DLNode *DomList, SLNode *SubList,
											 int *SubListSize, WTRANSFORM *transform, 
											 MAP *map, int threshold, 
											 ArithDecoder *decoder,
											 BasicCoder **DominantPassCoder, 
											 int DecodeBytes);
void DecodeNode(SLNode *SubList, int *SubListSize, WTRANSFORM *transform, 
					 DLNode *node, int threshold);

void DecodeSubordinatePass(SLNode *SubList, int SubListSize, int threshold,
									WTRANSFORM *transform, ArithDecoder *decoder,
									BasicCoder *SubordinateListCoder, int DecodeBytes);

void UNEZW2Error(char *fmt, ...);
void UNEZW2Warning(char *fmt, ...);

/* termination flag */
Boolean EndDecoding;


/*----------------------------------------------------------------------------*/	
/*----------------------------------------------------------------------------*/	
void main(int argc, char **argv)
{
	FIMAGE *DecodeImage;
	IMAGE  *DecodeImage2, *OriginalImage;
   WAVELET *wavelet;
	WTRANSFORM *transform;
	char EncodeImageFileName[_MAX_PATH];
	char DecodeImageFileName[_MAX_PATH];
	char OriginalImageFileName[_MAX_PATH];
	Boolean DecodeFileNameSpecify = FALSE;
	int NScale;
	ArithDecoder *decoder;
	BIT_FILE *EncodeStream;
	float AbsMaxCoeff;
	float ImageMean;
	int Threshold;
	int ImageXSize, ImageYSize;
	int MaxHistogramCount;
	int i;
	int Temp;
	char *ProgramName, *InputImageName, *OriginalImageName;
	Boolean ComparePSNR = FALSE;	
	float PSNR;
 	int DecodeBytes=0, EncodedImageFileSize;
   struct stat statistics; 
	int WaveletIndex;
	FILTERSET *Filter[16];
	int ScanOrder;

   EndDecoding  = FALSE;

   fprintf(stdout, "Embedded Zerotree Wavelet Coder - EZW2 (Decoder)\n");
	
	InitializeFilterSets();
	/* Initialize the filters used */
	Filter[0] = Antonini;
	Filter[1] = Adelson;
	Filter[2] = Odegard;
	Filter[3] = Villa1810;
	Filter[4] = Brislawn;
	Filter[5] = Brislawn2;
	Filter[6] = Haar;
	Filter[7] = Daub4;
	Filter[8] = Daub6;
	Filter[9] = Daub8;
	Filter[10] = Villa1;
	Filter[11] = Villa2;
	Filter[12] = Villa3;
	Filter[13] = Villa4;
	Filter[14] = Villa5;
	Filter[15] = Villa6;

	ProgramName = argv[0];

	if(argc < 2){
	   fprintf(stderr, "Usage: %s [encode image]\n", ProgramName);
		fprintf(stderr, "\t -d <decode filename>\n");
		fprintf(stderr, "\t -o <original image>\n");
		fprintf(stderr, "\t -b <decode bytes>\n");
		return;
	}
	
	InputImageName = argv[1];
	argc-=2; argv+=2;
   
	while(argc>0){
		if (!strcmp("-o", *argv)){
			argv++; argc--;
			OriginalImageName = *argv;
			argv++; argc--;
			sprintf(OriginalImageFileName, "%s.pgm", OriginalImageName);
			ComparePSNR = TRUE;
		}
		else if (!strcmp("-d", *argv)){
			argv++; argc--;
			sprintf(DecodeImageFileName, "%s.unezw.pgm", *argv);
			DecodeFileNameSpecify = TRUE;
			argv++; argc--; 
		}
		else if (!strcmp("-b", *argv)){
			argv++; argc--;
			DecodeBytes = atoi(*argv);
			if (DecodeBytes < 0){
				DecodeBytes = 0;
			}
			argv++; argc--;

		}
		else{
			fprintf(stderr, "Usage: %s [encode image]\n", ProgramName);
			fprintf(stderr, "\t -d <decode filename>\n");
			fprintf(stderr, "\t -o <original image>\n");
			fprintf(stderr, "\t -b <decode bytes>\n");
			return;
		}
	}

	sprintf(EncodeImageFileName, "%s.ezw", InputImageName);
	
	if(stat(EncodeImageFileName,&statistics) == -1){
		EncodedImageFileSize = 0;
	}
	else{
		EncodedImageFileSize = (int)statistics.st_size;
	}
	
	if (DecodeBytes == 0 || DecodeBytes>EncodedImageFileSize){
		DecodeBytes = EncodedImageFileSize;
	}

	if (!DecodeFileNameSpecify){
		if (DecodeBytes==EncodedImageFileSize){
			sprintf(DecodeImageFileName, "%s.unezw.pgm", InputImageName);	
		}
		else{
			sprintf(DecodeImageFileName, "%s_%d.unezw.pgm", InputImageName, DecodeBytes);
		}
	}

	/* open encode stream */
	EncodeStream = OpenInputBitFile(EncodeImageFileName);
	if (EncodeStream==NULL){
		UNEZW2Error("Fail to read encoded image.\n");
	}
	
	decoder = ArithDecoderAlloc(EncodeStream);
	ArithDecoderStart(decoder);
	
	/* read header */
	ImageXSize	 = BasicCoderReadNBits(decoder, ImageXSizeBits);
	ImageYSize	 = BasicCoderReadNBits(decoder, ImageYSizeBits);
	ScanOrder	 = BasicCoderReadNBits(decoder, ScanOrderBits);
	NScale		 = BasicCoderReadNBits(decoder, ScaleBits);
	WaveletIndex = BasicCoderReadNBits(decoder, WaveletIndexBits);
	Temp			 = BasicCoderReadNBits(decoder, ImageMeanBits);	
	ImageMean	 = *((float *)(&Temp));
	Threshold	 = 1<<(BasicCoderReadNBits(decoder, ThresholdBits));					
	MaxHistogramCount = BasicCoderReadNBits(decoder, MaxHistoCountBits);									
	
	/* initialize wavelet filter */
	wavelet = WaveletAlloc(Filter[WaveletIndex]);
	
	transform = WaveletTransformAllocBlank(wavelet, ImageXSize, ImageYSize, NScale, -1);
	
	if(ComparePSNR){
		fprintf(stdout, "Original Image: %s\nEncoded Image : %s\nDecoded Image : %s\n",
			OriginalImageFileName, EncodeImageFileName, DecodeImageFileName);
	}
	else{
		fprintf(stdout, "Encoded Image : %s\nDecoded Image : %s\n",
		EncodeImageFileName, DecodeImageFileName);
	}
	
	fprintf(stdout, "Width: %d   Height: %d   ", transform->hsize, transform->vsize);
	fprintf(stdout, "Mean: %.3f\n", ImageMean);
	fprintf(stdout, "Initial threshold: %d   Levels: %d   Wavelet: %s   Scan: %d\n",
		Threshold, NScale, FilterName[WaveletIndex], ScanOrder);
	fprintf(stdout, "Decode bytes: %d\n", DecodeBytes);

	/* decode */
	EZW2Decode(transform, decoder, DecodeBytes, Threshold, MaxHistogramCount, ScanOrder);
	
	DecodeImage = FImageAlloc(transform->hsize, transform->vsize);
	WaveletTransformInvert(transform, DecodeImage);
	
	/* put back mean */
	for (i=0; i<DecodeImage->xsize*DecodeImage->ysize; i++){
		DecodeImage->pixelLinear[i] += ImageMean;
	}
	
	WriteFloatToPGM(DecodeImage, DecodeImageFileName);
   
	/* change to char type 21/09/2002 */
	if (ComparePSNR){
		if ((OriginalImage = ReadPGM(OriginalImageFileName))==NULL){
			UNEZW2Error("Fail to open original image.\n");
		}

		if ((DecodeImage2 = ReadPGM(DecodeImageFileName))==NULL){
			UNEZW2Error("Fail to open decode image.\n");
		}

		PSNR = (float)ImageComparePSNR(OriginalImage, DecodeImage2);

		fprintf(stdout, "\nPSNR = %.4f dB\n", PSNR);
		ImageFree(OriginalImage);
		ImageFree(DecodeImage2);
	}

	ArithDecoderDealloc(decoder);
	FImageFree(DecodeImage);
	WaveletDealloc(wavelet);
	WaveletTransformDealloc(transform);
	RemoveFilterSets();
	CloseInputBitFile(EncodeStream);
	PrintLeaks();
	return;
}

/*----------------------------------------------------------------------------*/	
/*----------------------------------------------------------------------------*/	
void EZW2Decode(WTRANSFORM *transform, ArithDecoder *decoder, int DecodeBytes,
					 int threshold, int HistoCount, int ScanOrder)
{
	BasicCoder *DominantPassCoder[4];
	BasicCoder *SubordinateListCoder;
	int i, nPass=0;
	MAP *map;
	int BitCount;
	DLNode *DomList;
	SLNode *SubList;
	int SubListSize;

	/* dominant list */
	if ((DomList = (DLNode *)malloc(sizeof(DLNode)
						*transform->vsize*transform->hsize))==NULL){
		UNEZW2Error("Fail to allocate dominant list.\n");
	}
		
	/* subordinate list */
	if ((SubList = (SLNode *)malloc(sizeof(SLNode)
						*transform->vsize*transform->hsize))==NULL){
		UNEZW2Error("Fail to allocate subordinate list.\n");
	}
	SubListSize = 0;
	
	/* 4 context of dominant list coding */
	for (i=0; i< 4; i++){
		if ((DominantPassCoder[i] = BasicCoderAlloc(4, HistoCount)) == NULL){
			UNEZW2Error("Fail to allocate dominanat pass coder.\n");
		}
	}
	
	/* 1 context of subordinate list coding */
	if ((SubordinateListCoder = BasicCoderAlloc(2, HistoCount)) == NULL){
		UNEZW2Error("Fail to allocate subordinate list coder.\n");
	}
	
	/* map to keep track of coded significant coefficients */
	if ((map=MapAlloc(transform->hsize, transform->vsize, 
			transform->nsteps))==NULL){
		UNEZW2Error("Fail to allocate map.\n");
	}
	
	/* bits input so far */
	BitCount = ArithDecoderNBitsInput(decoder);

	while ((threshold!=0) && (EndDecoding!=TRUE) ){
		printf("Pass #%2d - ", ++nPass);

		if (ScanOrder==0){
			DecodeDominantPassMortonScan(DomList, SubList, &SubListSize, transform, 
						map, threshold, decoder, DominantPassCoder, DecodeBytes);
		}
		else{
			/* no implementation */
		}
		
		DecodeSubordinatePass(SubList, SubListSize, threshold>>1, transform, 
						decoder, SubordinateListCoder, DecodeBytes);

		BitCount = ArithDecoderNBitsInput(decoder);
		
		printf("Subordinate list size: %6d  (%5d)  Total bytes: %d\n", 
			SubListSize, threshold, BitCount/8);
		
		if (EndDecoding!=TRUE){
			threshold >>= 1;
			
			for (i=0; i<4; i++){
				BasicCoderReset(DominantPassCoder[i]);
			}
			BasicCoderReset(SubordinateListCoder);
			//SortSubordinateList(list);		
		}
	}