/*
 *         bpack - A Burrows-Wheeler Transform and Huffman Encoding
 *                            Implementation
 *
 *  Copyright (C) 2006      Angelo Arrifano <miknix@gmail.com>
 *  Copyright (C) 2006      Vasco Costa <the_papon@hotmail.com>
 *
 *  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 <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "bpack.h"
#include "huffman.h"

const unsigned int intsz=sizeof(int)*8;
#define CODESZ 32/sizeof(int)

//Byte frequency lookup table
struct freq_ltbl{
	unsigned char byte;
	unsigned int freq;
};

//Huffman tree node
struct htree_node{
	unsigned int freq;
	unsigned char byte;
	struct htree_node *le;
	struct htree_node *ri;
};

//Huffman tree node list
struct htree_node_lst{
	struct htree_node *node;
	struct htree_node_lst *ne;
};

//Huffman code lookup table
struct code_ltbl{
	unsigned int bcode[CODESZ]; //256 bit storage for codes, 8 ints for most of machines
	unsigned char len;
};


//Comparision function for freq table sorting
//Returns:
//  0: If b1 == b2
//  1: If b1 > b2
// -1: If b1 < b2 
int freq_ltbl_sort(struct freq_ltbl *b1, struct freq_ltbl *b2){
	if(b1->freq > b2->freq)
		return(1);
	else
		if(b1->freq == b2->freq)
			return(0);
		else
			return(-1);
}

//Function for building a huffman tree
//Returns:
//  Huffman tree ROOT
struct htree_node *build_htree(unsigned char *buffer, int size){
	struct freq_ltbl *ftbl;
	struct htree_node_lst *level=NULL, *nodel=NULL, *min_nodel=NULL;
	struct htree_node *htree=NULL;
	int i;
	unsigned int freq, min_freq=0;

	//Allocate (1.25kb) and initialize frequency lookup table
	if((ftbl=(struct freq_ltbl *)malloc(sizeof(struct freq_ltbl) * 256))==NULL) exit(10);

	for(i=0; i<256; i++){
		ftbl[i].byte=i;
		ftbl[i].freq=0;
	}	

	//Fill table with byte frequencies
	for(i=0; i<size; i++){
		ftbl[(int)buffer[i]].freq++;
	}

	//Sort freq table
	//FIXME: We sort all 256 bytes. Will be better to sort only existant bytes??
	qsort(ftbl, 256, sizeof(struct freq_ltbl), (int(*)(const void *,const void *))freq_ltbl_sort);

	//Create Huffman tree leafs and
	//link them in a list (bottom level).
	//Put lower frequencies in the end of list to increase
	//performance when searching for minimal sum of nodes (below)
	for(i=255; i>=0; i--){
		//Create leafs only for existant bytes
		if(ftbl[i].freq > 0){
			if((level=(struct htree_node_lst *)malloc(sizeof(struct htree_node_lst)))==NULL) exit(10);
			if((level->node=(struct htree_node *)malloc(sizeof(struct htree_node)))==NULL) exit(10);
			level->node->freq=ftbl[i].freq;
			level->node->byte=ftbl[i].byte;
			level->node->le=NULL;
			level->node->ri=NULL;
			level->ne=nodel;
			nodel=level;
		}else break;
	}

	//We do not need a frequency table anymore
	free(ftbl);
	
	//Do while we dont reach the top of the tree (top level)
	while(level->ne != NULL){
		//Find minimal sum of nodes in htree current level
		nodel=level;
		min_freq=nodel->node->freq + nodel->ne->node->freq + 1;
		while(nodel->ne != NULL){
			freq=nodel->node->freq + nodel->ne->node->freq;
			if(freq < min_freq){
				min_freq=freq;
				min_nodel=nodel;
			}else
			nodel=nodel->ne;
		}

		//Create a tree parent node with freq equal to the sum of freq
		//of its childs.
		if((htree=(struct htree_node *)malloc(sizeof(struct htree_node)))==NULL) exit(10);
		htree->freq=min_freq;
		htree->byte='_';
		htree->le=min_nodel->node;
		htree->ri=min_nodel->ne->node;
		//Replace two childs with their parent on current level.
		min_nodel->node=htree;
		nodel=min_nodel->ne;
		min_nodel->ne=nodel->ne;
		free(nodel);

		//Go back to the head of list
		nodel=level;
	}
	free(level);
	
	return(htree);
}


//Recursive function that runs over
//the entire tree and builds bitcodes
//for all bytes.
void build_code_ltbl_rec(struct htree_node *htree, struct code_ltbl *ctbl, struct code_ltbl code){
	unsigned char i;

	if(htree == NULL) exit_err(30);
	if(htree->le == NULL && htree->ri == NULL){
		memcpy(&ctbl[(int)htree->byte], &code, sizeof(struct code_ltbl));
		return;
	}
	i=code.len/intsz; //Get int position
	code.len++; //Increase bitlength
	code.bcode[i]<<=1; //Apend one zero to code
	build_code_ltbl_rec(htree->le, ctbl, code);
	code.bcode[i]|=1; //Change last bit to one
	build_code_ltbl_rec(htree->ri, ctbl, code);
}

int huff_code_output(struct code_ltbl *ctbl, unsigned char *ibuffer, int size){
	int i, j, k, blocksz;
	short int nrcodes;
	unsigned char *obuffer, obit, pos;
	struct code_ltbl code;

	obuffer=(char *)malloc(size);
	//FIXME: We need to check if compressed data grows
	//more than the original data.
	obuffer[0]=0;
	obit=0;
	for(i=size-1, k=size-1; i>=0; i--){
		memcpy(&code, &ctbl[ibuffer[i]], sizeof(struct code_ltbl));
		for(j=0; j<code.len; j++){
			if(obit == 8){
				obuffer[--k]=0;
				obit=0;
			}
			pos=j/intsz;
			obuffer[k]<<=1; 
			obit++;
			obuffer[k] |= code.bcode[pos] & 1;
			code.bcode[pos]>>=1;
		}
	}
	
	blocksz=size-k+1;

	//Write block size
	fwrite(&blocksz, sizeof(int), 1, stdout);

	//Write block
	fwrite(&obuffer[k], 1, blocksz, stdout);

	//Write number of used bits in last byte
	fwrite(&obit, 1, 1, stdout);

	//Find how many codes we have and write it to stdout
	//FIXME: Its better to do this during some previous operation
	nrcodes=0;
	for(i=0; i<256; i++){
		if(ctbl[i].len != 0) nrcodes++;
	}
	fwrite(&nrcodes, sizeof(short int), 1, stdout);

	//Write codes to stdout
	for(i=0; i<256; i++){
		if(ctbl[i].len != 0){
			fwrite((unsigned char *)&i, 1, 1, stdout);
			fwrite((unsigned char *)&ctbl[i].len, 1, 1, stdout);
			fwrite(ctbl[i].bcode, sizeof(int), ctbl[i].len/8+1, stdout);
		}
	}

	free(obuffer);

	return(blocksz);
}

//Main Function for Huffman Encoding
//Returns: Size of encoded data
int huff_encode(unsigned char *ibuffer, int size){
	struct htree_node *htree;
	struct code_ltbl *ctbl;
	struct code_ltbl initcode;
	int i, j;

	//Build huffman tree
	htree=build_htree(ibuffer, size); 

	//We only need 9kb for code lookup table.
	ctbl=(struct code_ltbl *)malloc(256 * sizeof(struct code_ltbl));
	for(i=0; i<256; i++){
		for(j=0; j<CODESZ; j++) ctbl[i].bcode[j]=0;
		ctbl[i].len=0;
	}

	//Call the recursive function with proper initial code
	initcode.len=0;
	for(i=0; i<CODESZ; i++) initcode.bcode[i]=0;
	build_code_ltbl_rec(htree, ctbl, initcode);
	free(htree); //We do not need this anymore.	
	
	size=huff_code_output(ctbl, ibuffer, size);

	free(ctbl);

	return(size);
}

void build_htree_rec(struct htree_node *htree, struct code_ltbl code, unsigned char byte){
unsigned char i;

	if(htree == NULL) exit_err(40);

	i=code.len/intsz;
	if(code.len == 1){
		if((code.bcode[i]&1) == 0){
			htree->le=(struct htree_node *)malloc(sizeof(struct htree_node));
			htree->le->le=NULL;
			htree->le->ri=NULL;
			htree->le->byte=byte;
		}else{
			htree->ri=(struct htree_node *)malloc(sizeof(struct htree_node));
			htree->ri->le=NULL;
			htree->ri->ri=NULL;
			htree->ri->byte=byte;
		}
		return;
	}
	
	if((code.bcode[i]&1) == 0){
		if(htree->le == NULL){
			htree->le=(struct htree_node *)malloc(sizeof(struct htree_node));
			htree->le->le=NULL;
			htree->le->ri=NULL;
		}
		code.bcode[i]>>=1;
		code.len--;
		build_htree_rec(htree->le, code, byte);
	}else{
		if(htree->ri == NULL){
			htree->ri=(struct htree_node *)malloc(sizeof(struct htree_node));
			htree->ri->le=NULL;
			htree->ri->ri=NULL;
		}
		code.bcode[i]>>=1;
		code.len--;
		build_htree_rec(htree->ri, code, byte);
	}	
}

void huff_decode(unsigned char *ibuffer, int isize, unsigned char *obuffer, int osize){
	unsigned char ubits, byte, posj, posk;
	short int nrcodes;
	unsigned int bcode[CODESZ];
	int i, j, k;
	struct code_ltbl hcode;
	struct htree_node *htree, *node;

	//Read number of used bits on last byte
	fread(&ubits, 1, 1, stdin);
	
	//Read number of codes
	fread(&nrcodes, sizeof(short int), 1, stdin);

	//Read codes and build Huffman tree
	htree=(struct htree_node *)malloc(sizeof(struct htree_node));
	for(i=0; i<nrcodes; i++){
		fread(&byte, 1, 1, stdin);
		fread(&hcode.len, 1, 1, stdin);
		fread(bcode, sizeof(int), hcode.len/8+1, stdin);

		//Invert bit code order
		memset(hcode.bcode, 0, CODESZ);
		for(j=0, k=hcode.len-1; j<hcode.len; j++, k--){
			posj=j/intsz;
			posk=k/intsz;
			hcode.bcode[posj]<<=1;
			hcode.bcode[posj]|= bcode[posk] & 1;		
			bcode[posk]>>=1;
		}
		build_htree_rec(htree, hcode, byte);
	}
	
	node=htree;
	for(i=0, k=0; i<isize; i++){
		for(j=0; j<ubits; j++){
			if(node == NULL) exit_err(50);
			if(node->le == NULL && node->ri == NULL){
				obuffer[k++]=node->byte;
				node=htree;
			}				
			if((ibuffer[i] & 1) == 0)
				node=node->le;
			else
				node=node->ri;
			ibuffer[i]>>=1;
		}
		ubits=8;
	}

	free(htree);
}

