/******************************************************************************
File:		bits.c

Authors: 	John Carpinelli   (johnfc@ecr.mu.oz.au)
	 	Wayne Salamonsen  (wbs@mundil.cs.mu.oz.au)

Purpose:	Data compression using a nth order binary model and revised 
		arithmetic coding method.

Based on: 	A. Moffat, R. Neal, I.H. Witten, "Arithmetic Coding Revisited",
		Proc. IEEE Data Compression Conference, Snowbird, Utah, 
		March 1995.


Copyright 1995 John Carpinelli and Wayne Salamonsen, All Rights Reserved.

These programs are supplied free of charge for research purposes only,
and may not sold or incorporated into any commercial product.  There is
ABSOLUTELY NO WARRANTY of any sort, nor any undertaking that they are
fit for ANY PURPOSE WHATSOEVER.  Use them at your own risk.  If you do
happen to find a bug, or have modifications to suggest, please report
the same to Alistair Moffat, alistair@cs.mu.oz.au.  The copyright
notice above and this statement of conditions must remain an integral
part of each and every copy made of these files.

******************************************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include "stats.h"
#include "coder.h"
#ifdef SYSV
#include <sys/times.h>
#include <limits.h>
#endif

#define MAX_CONTEXT_BITS	20    	/* max. number of bits for context */
#define MIN_CONTEXT_BITS	0	/* min. number of bits for context */
#define DEFAULT_BITS_CONTEXT 	16	/* default value for bits_context */
#define ENCODE          	0
#define DECODE          	1
#define MAGICNO         	"123b"  /* Magic Number for files */
#define MAGICNO_LENGTH		4	/* length of magic number */
#define BREAK_INTERVAL		10000	/* every round off to bit boundary */
#define MEGABYTE		(1<<20)	/* size of one megabyte */
#define NOMEMLEFT		-1	/* flag set when mem runs out */
#define DEFAULT_MEM		1	/* default 1 megabyte limit */
#define MIN_MBYTES        	1	/* minimum allowable memory size */
#define MAX_MBYTES        	255	/* maximum no for 8 bit int */


/* function prototypes */
void *do_malloc(size_t size);
int get_memory(size_t size);
void encode_file(unsigned char tempstring[], int templength);
void decode_file();
void print_results(int operation);


/* global variables */
int verbose = 0;			/* flag for printing statistics  */
int bits_context=DEFAULT_BITS_CONTEXT;	/* no. bits in binary contexts */
unsigned int non_null_contexts = 0;	/* count of contexts used */
unsigned int total_memory = 0;		/* total memory used by all contexts */
int mbytes = DEFAULT_MEM;		/* memory limit in megabytes */


/* 
 * parse command line arguments. Decide whether to decode or encode
 * and optional memory size. Also sets filename to stdin if none specified 
 */
int main(int argc, char *argv[])
{	
    int i;			
    int what = ENCODE;		/* flag as to whether to encode or decode */
    unsigned char tempstore[MAGICNO_LENGTH];	/* stores magic number */
    int templength=0;		/* number of magic number bytes read */
    int	selected = -1;		/* stores if decode set at command line */

    for (i=1; i<argc; ) 
    {
	if (argv[i][0] == '-') 
	{
	    switch(argv[i][1]) {
	      case 'e':		/* do encode */
		selected = ENCODE;
		i++;
		break;
	      case 'd':		/* do decode */
		selected = DECODE;
		i++;
		break;
	      case 'm':		/* set memory size */
		i++;
		mbytes = atoi(argv[i]);
		i++;
		break;
	      case 'v':		/* set verbose flag to print stats */
		verbose = 1;
		i++;
		break;
	      case 'f':		/* set number of F bits */
		i++;
		f_bits = atoi(argv[i]);
		max_frequency = 1 << f_bits;
		i++;
		break;
	      case 'c':
		i++;
		bits_context = atoi(argv[i]);
		i++;
		break;
	      default:		/* incorrect args */
		fprintf(stderr, "usage: %s [-e [-m n] | -d] [-v] [-f n] [-c n] [file]\n", argv[0]);
		exit(1);
	    }
	}
	else if (freopen(argv[i++], "r", stdin) == (FILE *)NULL) 
	{
	    fprintf(stderr, "%s: cannot read %s\n",
		    argv[0], argv[--i]);
	    exit(1);
	}
    }

    /* check if bits_context is within allowable range */
    if (bits_context < MIN_CONTEXT_BITS || bits_context > MAX_CONTEXT_BITS)
    {
	fprintf(stderr, "%s: context bits must be between %d and %d\n",
		argv[0], MIN_CONTEXT_BITS, MAX_CONTEXT_BITS);
	exit(1);
    }
    
    /* check if f_bits is within allowable range */
    if (f_bits < MIN_F_BITS || f_bits > MAX_F_BITS)
    {
	fprintf(stderr, "%s: number of f bits must be between %d and %d\n",
		argv[0] ,MIN_F_BITS, MAX_F_BITS);
	exit(1);
    }

    /* check if memory limit is within allowable range */ 
    if (mbytes < MIN_MBYTES || mbytes > MAX_MBYTES)
    {
	fprintf(stderr, "memory limit must be between %d and %d\n", 
		MIN_MBYTES, MAX_MBYTES);
	exit(1);
    }


    /* Check input file for Magic Number. */
    if (selected != ENCODE)
    {
 	templength=fread(tempstore, 1, MAGICNO_LENGTH, stdin);
	bytes_input += templength;
	if (memcmp(tempstore, MAGICNO, MAGICNO_LENGTH) == 0)
	    what = DECODE;
	else if (selected == DECODE)
	{
	    fprintf(stderr, 
		 "bits: bad magic number... incompatible compressed file\n");
	    exit(1);
	}
    }
	
    if (what == ENCODE)	
    {
	/* print magic number in output */
	fwrite(MAGICNO, 1, MAGICNO_LENGTH, stdout);
	bytes_output += MAGICNO_LENGTH;

	/* store memory limit being used in output */
	putc(mbytes, stdout);
	bytes_output += 1;
	
	/* print number of f_bits being used in output */
	putc(f_bits, stdout);
	bytes_output += 1;

	/* print bits_context in output */
	putc(bits_context, stdout);
	bytes_output += 1;
	encode_file(tempstore, templength);
    }
    else
    {
	/* read memory limit to be used and store in mbytes */
	mbytes = getc(stdin);
        bytes_input += 1;

	/* get number of f_bits to be used and store in f_bits */
	f_bits = getc(stdin);
	max_frequency = 1<<f_bits;
	bytes_input++;

	/* get number of bits_context to be used and store */
	bits_context = getc(stdin);
	bytes_input++;

	decode_file();
    }
    if (verbose)
	print_results(what);
    return 0;			/* exited cleanly */
}



/*
 *
 * print the results of compressing/decompressing a file
 *
 */
void print_results(int operation)
{
    fprintf(stderr, "Context bits           : %10d\n", bits_context);
    fprintf(stderr, "Number of contexts     : %10u\n", 
	    non_null_contexts);
    fprintf(stderr, "Memory used            : %10.1f Kbytes\n", 
	total_memory*1.0/1024);
    if (operation == ENCODE)
	fprintf(stderr, "Input file size        : %10u bytes\n", bytes_input);
    fprintf(stderr, "Output file size       : %10u bytes\n", bytes_output);
    if ((operation == ENCODE) && (bytes_input > 0))
	fprintf(stderr, "Compression rate       : %10.3f bpc (%0.2f%%) \n", 
		8.0 * bytes_output / bytes_input, 
		(float)bytes_output/bytes_input*100);

    /* only provide timing details if times function is available */
#ifdef SYSV
{
    struct tms cpu_usage;
    float cpu_used, comp_rate;

    times(&cpu_usage);    	/* determine the cpu time used */
    cpu_used = ((float) cpu_usage.tms_utime) / sysconf(_SC_CLK_TCK);

    if (cpu_used == 0)
	comp_rate = 0;
    else
    {
	if (operation == ENCODE)
	    comp_rate = ((float) bytes_input) / (1024 * cpu_used);
	else
	    comp_rate = ((float) bytes_output) / (1024 * cpu_used);
    }
    fprintf(stderr, "Compression time       : %10.2f seconds (%0.2f Kb/s)\n",
	    cpu_used, comp_rate);
}
#endif
}





/*
 *
 * call the standard C function malloc after checking against the memory
 * limit. If the limit is exceeded return NULL
 *
 */
void 
*do_malloc(size_t size)
{
    total_memory += size;
    if ((total_memory / MEGABYTE) >= mbytes)
	return NULL;
    else
        return (malloc(size));
}


/*
 *
 * adds specified memory to current memory count
 * returns 0 if successful, NOMEMLEFT if memory limit is reached
 */
int 
get_memory(size_t size)
{
    total_memory += size;
    if ((total_memory / MEGABYTE) >= mbytes)
	return NOMEMLEFT;
    else
	return 0;
      
}

/*
 *
 * compress data from stdin to a file on stdout by encoding
 * each bit with probabilities derived from the previous bits
 *
 */
void encode_file(unsigned char tempstring[], int templength)
{
    unsigned int	i, j, cur_context, closest_context;
    unsigned int 	mask, bit, next_break, buffer;
    int			bits_to_go;
    binary_context	**contexts;
    binary_context	*still_coding;		

    /* initialise context array */
    contexts = (binary_context **)do_malloc(sizeof(binary_context *) * 
					  (1 << bits_context));
    if (contexts == NULL)
    {
	fprintf(stderr, "bits: not enough memory to allocate context array\n");
	exit(1);
    }
    still_coding = create_binary_context();
    
    /* initialise contexts to NULL */
    for (i=0; i < 1 << bits_context; i++)
	contexts[i] = NULL;

    /* initalise variables */
    cur_context = 0;
    mask = (1 << bits_context) - 1;
    next_break = BREAK_INTERVAL;
    start_encode();
    startoutputtingbits();

    /* encode the first characters read for testing MAGIC_NO */
    for (i=0; ; i++)
    {
	if (i < templength)
	    buffer = tempstring[i];
	else 
	{
	    if ((buffer = getchar()) == EOF)
		break;
	    bytes_input++;
	}
	binary_encode(still_coding, 1);
	   
	for (bits_to_go = 7; bits_to_go >= 0; bits_to_go--)
	{
	    if (contexts[cur_context] == NULL)
	    {
		if (get_memory(sizeof(binary_context)) != NOMEMLEFT)
		{
		    contexts[cur_context] = create_binary_context();
		    non_null_contexts++;
		}
		else 
		    /* 
		     * determine closest existing context to current one
		     * by stripping off the leading bits of the context
		     * guaranteed to get contexts[0] if nothing closer
		     */
		{
		    closest_context = cur_context;
		    j = 1;
		    do
		    {
			closest_context &= (mask >> j);
			j++;
		    } while (contexts[closest_context] == NULL);
		    contexts[cur_context] = contexts[closest_context];
		}
	    }
	    bit = (buffer >> bits_to_go) & 1;
	    binary_encode(contexts[cur_context], bit);
	    cur_context = ((cur_context << 1) | bit) & mask;
	}
	if (next_break-- == 0)
	{
	    finish_encode();
	    start_encode();
	    next_break = BREAK_INTERVAL;
	}
    
    }
    /* encode end of message flag */
    binary_encode(still_coding, 0);
    finish_encode();
    doneoutputtingbits();
}


/*
 *
 * decode a compressed file to stdout
 *
 */
void decode_file()
{
    int			i, j, cur_context, closest_context, buffer, bits_to_go;
    int 		mask, bit, next_break;
    binary_context	**contexts;
    binary_context	*still_coding;		

    /* initialise context array */
    contexts = (binary_context **)do_malloc(sizeof(binary_context *) * 
					  (1 << bits_context));
    if (contexts == NULL)
    {
	fprintf(stderr, "bits: unable to malloc %d bytes\n",
		sizeof(binary_context *) * (1 << bits_context)); 
	exit(1);
    }
    still_coding = create_binary_context();
    
    /* initialise contexts to NULL */
    for (i=0; i < 1 << bits_context; i++)
	contexts[i] = (binary_context *)NULL;

    /* initalise variables */
    cur_context = 0;
    mask = (1 << bits_context) - 1;
    next_break = BREAK_INTERVAL;

    start_decode();
    startinputtingbits();

    /* decode the file */
    while (binary_decode(still_coding))
    {
	buffer = 0;
	for (bits_to_go = 7; bits_to_go >= 0; bits_to_go--)
	{
	    if (contexts[cur_context] == (binary_context *)NULL)
	    {
		if (get_memory(sizeof(binary_context)) != NOMEMLEFT)
		{
		    contexts[cur_context] = create_binary_context();
		    non_null_contexts++;
		}
		else 
		{   /* 
		     * determine closest existing context to current one
		     * by stripping off the leading bits of the context
		     * guaranteed to get contexts[0] if nothing closer
		     */
		    closest_context = cur_context;
		    j = 1;
		    do
		    {
			closest_context &= (mask >> j);
			j++;
		    } while (contexts[closest_context] == NULL);
		    contexts[cur_context] = contexts[closest_context];
		}
	    }
	    bit = binary_decode(contexts[cur_context]);
	    buffer = (buffer << 1) | bit;
	    cur_context = ((cur_context << 1) | bit) & mask;
	}	    
	putchar(buffer);
	bytes_output++;

	if (next_break-- == 0)
	{
	    finish_decode();
	    start_decode();
	    next_break = BREAK_INTERVAL;
	}
    }
    finish_decode();
    doneinputtingbits();
}