/*
 *  arith.c:		Adaptive arithmetic coding and decoding
 *
 *  Written by:		Ullrich Hafner
 *
 *  This file is part of FIASCO (Fractal Image And Sequence COdec)
 *  Copyright (C) 1994-2000 Ullrich Hafner
 */

/*
 *  $Date: 2000/06/14 20:49:37 $
 *  $Author: hafner $
 *  $Revision: 5.1 $
 *  $State: Exp $
 */

#include "config.h"

#include "types.h"
#include "macros.h"
#include "error.h"

#include "bit-io.h"
#include "misc.h"
#include "arith.h"

/******************************************************************************

				public code

******************************************************************************/

arith_t *
alloc_encoder (bitfile_t *output)
/*
 *  Arithmetic coder constructor:
 *  Initialize the arithmetic coder.
 *
 *  Return value:
 *	A pointer to the new coder structure
 */
{
   arith_t *arith = Calloc (1, sizeof (arith_t));

   assert (output);

   arith->low       = LOW;
   arith->high      = HIGH;
   arith->underflow = 0;
   arith->file 	    = output;

   return arith;
}

void
free_encoder (arith_t *arith)
/*
 *  Arithmetic encoder destructor.
 *  Flush the arithmetic coder. Append all remaining bits to the
 *  output stream. Append zero bits to get the output file byte aligned.
 *
 *  No return value.
 */
{
   u_word_t   low;			/* start of the current code range  */
   u_word_t   high;			/* end of the current code range    */
   u_word_t   underflow;		/* number of underflow bits pending */
   bitfile_t *output;

   assert (arith);

   low       = arith->low;
   high      = arith->high;
   underflow = arith->underflow;
   output    = arith->file;

   low = high;

   RESCALE_OUTPUT_INTERVAL;

   OUTPUT_BYTE_ALIGN (output);

   Free (arith);
}

real_t
encode_symbol (unsigned symbol, arith_t *arith, model_t *model)
/*
 *  Encode the given 'symbol' using the given probability 'model'.
 *  The current state of the arithmetic coder is given by 'arith'.
 *  Output bits are appended to the stream 'output'.
 *
 *  The model is updated after encoding the symbol (if necessary the
 *  symbol counts are rescaled).
 *
 *  Return value:
 *	information content of the encoded symbol.
 *
 *  Side effects:
 *	'model' is updated (probability distribution)
 *	'arith' is updated (coder state)
 */
{
   u_word_t   low_count;		/* lower bound of 'symbol' interval */
   u_word_t   high_count;		/* upper bound of 'symbol' interval */
   u_word_t   scale;			/* range of all 'm' symbol intervals */
   unsigned   range;			/* range of current interval */
   unsigned   index;			/* index of probability model */
   u_word_t   low;			/* start of the current code range  */
   u_word_t   high;			/* end of the current code range    */
   u_word_t   underflow;		/* number of underflow bits pending */
   bitfile_t *output;			/* output file */

   assert (model && arith);

   /*
    * Get interval values
    */
   low       = arith->low;
   high      = arith->high;
   underflow = arith->underflow;
   output    = arith->file;

   assert (high > low);

   if (model->order > 0)		/* order-'n' model*/
   {
      unsigned power;			/* multiplicator */
      unsigned i;

      /*
       *  Compute index of the probability model to use.
       *  See init_model() for more details.
       */
      power = 1;			/* multiplicator */
      index = 0;			/* address of prob. model */

      for (i = 0; i < model->order; i++) /* generate a M-nary number */
      {
	 index += model->context [i] * power;
	 power *= model->symbols;
      }

      index *= model->symbols + 1;	/* we need space for M + 1 elements */

      for (i = 0; i < model->order - 1; i++)
	 model->context [i] = model->context [i + 1];
      model->context [i] = symbol;
   }
   else
      index = 0;

   scale      = model->totals [index + model->symbols];
   low_count  = model->totals [index + symbol];
   high_count = model->totals [index + symbol + 1];

   /*
    *  Compute the new interval depending on the input 'symbol'.
    */
   range = (high - low) + 1;
   high  = low + (u_word_t) ((range * high_count) / scale - 1);
   low   = low + (u_word_t) ((range * low_count) / scale);

   RESCALE_OUTPUT_INTERVAL;

   if (model->scale > 0)		/* adaptive model */
   {
      unsigned i;

      /*
       *  Update probability model
       */
      for (i = symbol + 1; i <= model->symbols; i++)
	 model->totals [index + i]++;
      if (model->totals [index + model->symbols] > model->scale) /* scaling */
      {
	 for (i = 1; i <= model->symbols; i++)
	 {
	    model->totals [index + i] >>= 1;
	    if (model->totals [index + i] <= model->totals [index + i - 1])
	       model->totals [index + i] = model->totals [index + i - 1] + 1;
	 }
      }
   }

   /*
    *  Store interval values
    */
   arith->low  	    = low;
   arith->high 	    = high;
   arith->underflow = underflow;

   return - log2 ((high_count - low_count) / (real_t) scale);
}

void
encode_array (bitfile_t *output, const unsigned *data, const unsigned *context,
	      const unsigned *c_symbols, unsigned n_context, unsigned n_data,
	      unsigned scaling)
/*
 *  Arithmetic coding of #'n_data' symbols given in the array 'data'.
 *  If 'n_context' > 1 then a number (context [n]) is assigned to every
 *  data element n, specifying which context (i.e. number of symbols given by
 *  c_symbols [context [n]] and adaptive probability model) must be used.
 *  Rescale probability models if range > 'scaling'.
 *
 *  No return value.
 */
{
   u_word_t **totals;			/* probability model */

   if (!n_context)
      n_context = 1;			/* always use one context */

   assert (output && c_symbols && data);
   assert (n_context == 1 || context);

   /*
    *  Allocate probability models, start with uniform distribution
    */
   totals = Calloc (n_context, sizeof (u_word_t *));
   {
      unsigned c;

      for (c = 0; c < n_context; c++)
      {
	 unsigned i;

	 totals [c]    = Calloc (c_symbols [c] + 1, sizeof (u_word_t));
	 totals [c][0] = 0;

	 for (i = 0; i < c_symbols [c]; i++)
	    totals [c][i + 1] = totals [c][i] + 1;
      }
   }

   /*
    *  Encode array elements
    */
   {
      u_word_t low  	 = 0;		/* Start of the current code range */
      u_word_t high 	 = 0xffff;	/* End of the current code range */
      u_word_t underflow = 0;		/* Number of underflow bits pending */
      unsigned n;

      for (n = 0; n < n_data; n++)
      {
	 u_word_t low_count;		/* lower bound of 'symbol' interval */
	 u_word_t high_count;		/* upper bound of 'symbol' interval */
	 u_word_t scale;		/* range of all 'm' symbol intervals */
	 unsigned range;		/* current range */
	 int	  d;			/* current data symbol */
	 int	  c;			/* context of current data symbol */

	 d = data [n];
	 c = n_context > 1 ? context [n] : 0;

	 scale	    = totals [c][c_symbols [c]];
	 low_count  = totals [c][d];
	 high_count = totals [c][d + 1];

	 /*
	  * Rescale high and low for the new symbol.
	  */
	 range = (high - low) + 1;
	 high  = low + (u_word_t) ((range * high_count) / scale - 1);
	 low   = low + (u_word_t) ((range * low_count) / scale);
	 RESCALE_OUTPUT_INTERVAL;

	 /*
	  *  Update probability models
	  */
	 {
	    unsigned i;

	    for (i = d + 1; i < c_symbols [c] + 1; i++)
	       totals [c][i]++;

	    if (totals [c][c_symbols [c]] > scaling) /* scaling */
	       for (i = 1; i < c_symbols [c] + 1; i++)
	       {
		  totals [c][i] >>= 1;
		  if (totals [c][i] <= totals [c][i - 1])
		     totals [c][i] = totals [c][i - 1] + 1;
	       }
	 }
      }
      /*
       *  Flush arithmetic encoder
       */
      low = high;
      RESCALE_OUTPUT_INTERVAL;
      OUTPUT_BYTE_ALIGN (output);
   }

   /*
    *  Cleanup ...
    */
   {
      unsigned c;
      for (c = 0; c < n_context; c++)
	 Free (totals [c]);
      Free (totals);
   }
}

arith_t *
alloc_decoder (bitfile_t *input)
/*
 *  Arithmetic decoder constructor:
 *  Initialize the arithmetic decoder with the first
 *  16 input bits from the stream 'input'.
 *
 *  Return value:
 *	A pointer to the new decoder structure
 */

{
   arith_t *arith = Calloc (1, sizeof (arith_t));

   assert (input);

   arith->low  = LOW;
   arith->high = HIGH;
   arith->code = get_bits (input, 16);
   arith->file = input;

   return arith;
}

void
free_decoder (arith_t *arith)
/*
 *  Arithmetic decoder destructor:
 *  Flush the arithmetic decoder, i.e., read bits to get the input
 *  file byte aligned.
 *
 *  No return value.
 *
 *  Side effects:
 *	structure 'arith' is discarded.
 */
{
   assert (arith);

   INPUT_BYTE_ALIGN (arith->file);

   Free (arith);
}

unsigned
decode_symbol (arith_t *arith, model_t *model)
/*
 *  Decode the next symbol - the state of the arithmetic decoder
 *  is given in 'arith'. Read refinement bits from the stream 'input'
 *  and use the given probability 'model'. Update the probability model after
 *  decoding the symbol (if necessary also rescale the symbol counts).
 *
 *  Return value:
 *	decoded symbol
 *
 *  Side effects:
 *	'model' is updated (probability distribution)
 *	'arith' is updated (decoder state)
 */
{
   unsigned   range;			/* range of current interval */
   unsigned   count;			/* value in the current interval */
   unsigned   index;			/* index of probability model */
   unsigned   symbol;			/* decoded symbol */
   u_word_t   scale;			/* range of all 'm' symbol intervals */
   u_word_t   low;			/* start of the current code range  */
   u_word_t   high;			/* end of the current code range    */
   u_word_t   code;			/* the present input code value */
   bitfile_t *input;			/* input file */

   assert (arith && model);

   /*
    * Get interval values
    */
   low   = arith->low;
   high  = arith->high;
   code  = arith->code;
   input = arith->file;

   assert (high > low);

   if (model->order > 0)		/* order-'n' model */
   {
      unsigned power;			/* multiplicator */
      unsigned i;

      /*
       *  Compute index of the probability model to use.
       *  See init_model() for more details.
       */
      power = 1;			/* multiplicator */
      index = 0;			/* address of prob. model */

      for (i = 0; i < model->order; i++) /* generate a m-nary number */
      {
	 index += model->context[i] * power;
	 power *= model->symbols;
      }

      index *= model->symbols + 1;	/* we need space for m + 1 elements */
   }
   else
      index = 0;

   scale = model->totals [index + model->symbols];
   range = (high - low) + 1;
   count = ((code - low + 1) * scale - 1) / range;

   for (symbol = model->symbols; count < model->totals [index + symbol];
	symbol--)
      ;

   if (model->order > 0)		/* order-'n' model */
   {
      unsigned i;

      for (i = 0; i < model->order - 1; i++)
	 model->context [i] = model->context [i + 1];
      model->context [i] = symbol;
   }

   /*
    *  Compute interval boundaries
    */
   {
      u_word_t low_count;		/* lower bound of 'symbol' interval */
      u_word_t high_count;		/* upper bound of 'symbol' interval */

      low_count  = model->totals [index + symbol];
      high_count = model->totals [index + symbol + 1];
      high       = low + (u_word_t) ((range * high_count) / scale - 1 );
      low        = low + (u_word_t) ((range * low_count) / scale );
   }

   RESCALE_INPUT_INTERVAL;

   if (model->scale > 0)		/* adaptive model */
   {
      unsigned i;

      /*
       *  Update probability model
       */
      for (i = symbol + 1; i <= model->symbols; i++)
	 model->totals [index + i]++;
      if (model->totals [index + model->symbols] > model->scale) /* scaling */
      {
	 for (i = 1; i <= model->symbols; i++)
	 {
	    model->totals [index + i] >>= 1;
	    if (model->totals [index + i] <= model->totals [index + i - 1])
	       model->totals [index + i] = model->totals [index + i - 1] + 1;
	 }
      }
   }

   /*
    *  Store interval values
    */
   arith->low  = low;
   arith->high = high;
   arith->code = code;

   return symbol;
}

unsigned *
decode_array (bitfile_t *input, const unsigned *context,
	      const unsigned *c_symbols, unsigned n_context,
	      unsigned n_data, unsigned scaling)
/*
 *  Arithmetic decoding of #'n_data' symbols.
 *  If 'n_context' > 1 then a number (context [n]) is assigned to every
 *  data element n, specifying which context (i.e. number of symbols given by
 *  c_symbols [context [n]] and adaptive probability model) must be used.
 *  Rescale probability models if range > 'scaling'.
 *
 *  Return value:
 *	pointer to array containing the decoded symbols
 */
{
   unsigned  *data;			/* array to store decoded symbols */
   u_word_t **totals;			/* probability model */

   if (n_context < 1)
      n_context = 1;			/* always use one context */
   assert (input && c_symbols);
   assert (n_context == 1 || context);

   data = Calloc (n_data, sizeof (unsigned));

   /*
    *  Allocate probability models, start with uniform distribution
    */
   totals = Calloc (n_context, sizeof (u_word_t *));
   {
      unsigned c;

      for (c = 0; c < n_context; c++)
      {
	 unsigned i;

	 totals [c]    = Calloc (c_symbols [c] + 1, sizeof (u_word_t));
	 totals [c][0] = 0;

	 for (i = 0; i < c_symbols [c]; i++)
	    totals [c][i + 1] = totals [c][i] + 1;
      }
   }

   /*
    *  Fill array 'data' with decoded values
    */
   {
      u_word_t code = get_bits (input, 16); /* The present input code value */
      u_word_t low  = 0;		/* Start of the current code range */
      u_word_t high = 0xffff;		/* End of the current code range */
      unsigned n;

      for (n = 0; n < n_data; n++)
      {
	 u_word_t scale;		/* range of all 'm' symbol intervals */
	 u_word_t low_count;		/* lower bound of 'symbol' interval */
	 u_word_t high_count;		/* upper bound of 'symbol' interval */
	 unsigned count;		/* value in the current interval */
	 unsigned range;		/* current interval range */
	 unsigned d;			/* current data symbol */
	 unsigned c;			/* context of current data symbol */

	 c = n_context > 1 ? context [n] : 0;

	 assert (high > low);
	 scale = totals [c][c_symbols [c]];
	 range = (high - low) + 1;
	 count = (((code - low) + 1 ) * scale - 1) / range;

	 for (d = c_symbols [c]; count < totals [c][d]; d--) /* next symbol */
	    ;
	 low_count  = totals [c][d];
	 high_count = totals [c][d + 1];

	 high = low + (u_word_t) ((range * high_count) / scale - 1 );
	 low  = low + (u_word_t) ((range * low_count) / scale );
	 RESCALE_INPUT_INTERVAL;

	 /*
	  *  Updata probability models
	  */
	 {
	    unsigned i;

	    for (i = d + 1; i < c_symbols [c] + 1; i++)
	       totals [c][i]++;

	    if (totals [c][c_symbols [c]] > scaling) /* scaling */
	       for (i = 1; i < c_symbols [c] + 1; i++)
	       {
		  totals [c][i] >>= 1;
		  if (totals [c][i] <= totals [c][i - 1])
		     totals [c][i] = totals [c][i - 1] + 1;
	       }
	 }
	 data [n] = d;
      }
      INPUT_BYTE_ALIGN (input);
   }

   /*
    *  Cleanup ...
    */
   {
      unsigned c;

      for (c = 0; c < n_context; c++)
	 Free (totals [c]);
      Free (totals);
   }

   return data;
}

model_t *
alloc_model (unsigned m, unsigned scale, unsigned n, unsigned *totals)
/*
 *  Model constructor:
 *  allocate and initialize an order-'n' probability model.
 *  The size of the source alphabet is 'm'. Rescale the symbol counts after
 *  'scale' symbols are encoded/decoded. The initial probability of every
 *  symbol is 1/m.
 *  If 'scale' = 0 then use static modeling (p = 1/n).
 *  If 'totals' is not NULL then use this array of 'm' values to set
 *  the initial counts.
 *
 *  Return value:
 *	a pointer to the new probability model structure.
 *
 *  Note: We recommend a small size of the alphabet because no escape codes
 *  are used to encode/decode previously unseen symbols.
 *
 */
{
   model_t  *model;			/* new probability model */
   unsigned  num;			/* number of contexts to allocate */
   bool_t    cont;			/* continue flag */
   bool_t    dec;			/* next order flag */
   unsigned  i;

   /*
    *  Allocate memory for the structure
    */
   model          = Calloc (1, sizeof (model_t));
   model->symbols = m;
   model->scale   = scale;
   model->order   = n;
   model->context = n > 0 ? Calloc (n, sizeof (unsigned)) : NULL;
   /*
    *  Allocate memory for the probability model.
    *  Each of the m^n different contexts requires its own probability model.
    */
   for (num = 1, i = 0; i < model->order; i++)
      num *= model->symbols;

   model->totals = Calloc (num * (model->symbols + 1), sizeof (unsigned));

   for (i = 0; i < model->order; i++)
      model->context[i] = 0;		/* start with context 0,0, .. ,0 */
   cont = YES;
   while (cont)				/* repeat while context != M ... M */
   {
      int	power;			/* multiplicator */
      int	index;			/* index of probability model */
      /*
       *  There are m^n different contexts:
       *  Let "context_1 context_2 ... context_n symbol" be the current input
       *  stream then the index of the probability model is given by:
       *  index = context_1 * M^0 + context_2 * M^1 + ... + context_n * M^(n-1)
       */
      power = 1;			/* multiplicator */
      index = 0;			/* address of prob. model */

      for (i = 0; i < model->order; i++)	/* generate a m-nary number */
      {
	 index += model->context[i] * power;
	 power *= model->symbols;
      }

      index *= model->symbols + 1;	/* size of each model is m + 1 */

      model->totals [index + 0] = 0;	/* always zero */

      for (i = 1; i <= model->symbols; i++) /* prob of each symbol is 1/m or
					       as given in totals */
	 model->totals[index + i] = model->totals [index + i - 1]
				    + (totals ? totals [i - 1] : 1);

      if (model->order == 0)		/* order-0 model */
	 cont = NO;
      else				/* try next context */
	 for (i = model->order - 1, dec = YES; dec; i--)
	 {
	    dec = NO;
	    model->context[i]++;
	    if (model->context[i] >= model->symbols)
	    {
	       /* change previous context */
	       model->context[i] = 0;
	       if (i > 0)		/* there's still a context remaining */
		  dec = YES;
	       else
		  cont = NO;		/* all context models initialized */
	    }
	 }
   }
   for (i = 0; i < model->order; i++)
      model->context[i] = 0;		/* start with context 0,0, .. ,0 */

   return model;
}

void
free_model (model_t *model)
/*
 *  Model destructor:
 *  Free memory allocated by the arithmetic 'model'.
 *
 *  No return value.
 *
 *  Side effects:
 *	struct 'model' is discarded
 */
{
   if (model != NULL)
   {
      if (model->context != NULL)
	 Free (model->context);
      Free (model->totals);
      Free (model);
   }
   else
      warning ("Can't free model <NULL>.");
}