/*
 *  coeff.c:		Matching pursuit coefficients probability model
 *
 *  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:50:51 $
 *  $Author: hafner $
 *  $Revision: 5.1 $
 *  $State: Exp $
 */

#include <string.h>
#include "config.h"

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

#include "rpf.h"
#include "misc.h"
#include "coeff.h"

/*
 *  Coefficient model interface:
 *  Implementing the coefficients model interface requires the
 *  following steps: 
 *  - Add a constructor that initializes the coeff_t structure
 *  - Allocate new model with default_alloc() 
 *  - Fill the c_array_t coeff_models[] array with constructor and name
 *  - Write code for methods bits() and update()
 *  - Either use default functions for remaining methods or override them
 *  The new model is automatically registered at the command line.
 */

/*****************************************************************************
		  uniform distribution coefficients model
*****************************************************************************/

static coeff_t *
alloc_uniform_coeff_model (rpf_t *rpf, rpf_t *dc_rpf,
			   unsigned min_level, unsigned max_level);
static void
uniform_update (const real_t *used_coeff, const word_t *used_states,
		unsigned level, coeff_t *coeff);
static real_t
uniform_bits (const real_t *used_coeff, const word_t *used_states,
	      unsigned level, const coeff_t *coeff);

/*****************************************************************************
			  default functions
*****************************************************************************/

static void
default_model_free (void *model);
static void *
default_model_duplicate (const coeff_t *coeff, const void *model);
static void
default_free (coeff_t *coeff);
static coeff_t *
default_alloc (rpf_t *rpf, rpf_t *dc_rpf,
	       unsigned min_level, unsigned max_level);

/*****************************************************************************
		adaptive arithmetic coding model
*****************************************************************************/

static coeff_t *
alloc_aac_coeff_model (rpf_t *rpf, rpf_t *dc_rpf,
		       unsigned min_level, unsigned max_level);
static void
aac_model_free (void *model);
static void *
aac_model_alloc (const coeff_t *coeff);
static void *
aac_model_duplicate (const coeff_t *coeff, const void *model);
static void
aac_update (const real_t *used_coeff, const word_t *used_states,
	    unsigned level, coeff_t *coeff);
static real_t
aac_bits (const real_t *used_coeff, const word_t *used_states,
	  unsigned level, const coeff_t *coeff);

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

				public code
  
*****************************************************************************/

typedef struct c_array
{
   const char *identifier;
   coeff_t    *(*function) (rpf_t *rpf, rpf_t *dc_rpf,
			    unsigned min_level, unsigned max_level);
} c_array_t;

c_array_t coeff_models[] = {{"adaptive", alloc_aac_coeff_model},
			    {"uniform",	 alloc_uniform_coeff_model},
			    {NULL,	 NULL}};

coeff_t *
alloc_coeff_model (const char *coeff_model_name, rpf_t *rpf, rpf_t *dc_rpf,
		   unsigned min_level, unsigned max_level)
/*
 *  Allocate a new coefficients model which is identified by the string
 *  'coeff_model_name'.  'rpf' and 'dc_rpf' define the reduced
 *  precision formats the should be used to quantize normal and DC
 *  components, respectively. 'min_level' and 'max_level' define the
 *  range of range approximations.
 * 
 *  Return value:
 *	pointer to the allocated coefficients model
 *
 *  Note:
 *      Refer to 'coeff.h' for a short description of the member functions.  */
{
   unsigned n;
   
   for (n = 0; coeff_models [n].identifier; n++) /* step through all id's */
      if (strcaseeq (coeff_models [n].identifier, coeff_model_name)) 
	 return coeff_models [n].function (rpf, dc_rpf, min_level, max_level);

   warning ("Can't initialize coefficients model '%s'. "
	    "Using default value '%s'.",
	    coeff_model_name, coeff_models [0].identifier);

   return coeff_models [0].function (rpf, dc_rpf, min_level, max_level);
}

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

				private code
  
*****************************************************************************/

/*****************************************************************************
		  uniform distribution coefficients model
*****************************************************************************/

static coeff_t *
alloc_uniform_coeff_model (rpf_t *rpf, rpf_t *dc_rpf,
			   unsigned min_level, unsigned max_level)
/*
 *  Underlying probability model: uniform distribution.
 *  I.e. each coefficient is written as such with
 *  (1 + exponent_bits + mantissa_bits) bits.
 */
{
   coeff_t *coeff = default_alloc (rpf, dc_rpf, min_level, max_level);

   coeff->bits   = uniform_bits;
   coeff->update = uniform_update;
   
   return coeff;
}

static real_t
uniform_bits (const real_t *used_coeff, const word_t *used_states,
	      unsigned level, const coeff_t *coeff)
{
   unsigned edge;
   real_t   bits = 0;			/* #bits to store coefficients */
   
   for (edge = 0; isedge (used_states [edge]); edge++)
   {
      rpf_t *rpf = used_states [edge] ? coeff->rpf : coeff->dc_rpf;
      
      bits += rpf->mantissa_bits + 1;
   }

   return bits;
}

static void
uniform_update (const real_t *used_coeff, const word_t *used_states,
		unsigned level, coeff_t *coeff)
{
   return;				/* nothing to do */
}

/*****************************************************************************
		adaptive arithmetic coding  model
*****************************************************************************/

typedef struct aac_model
{
   word_t *counts;
   word_t *totals;
} aac_model_t;

static coeff_t *
alloc_aac_coeff_model (rpf_t *rpf, rpf_t *dc_rpf,
		       unsigned min_level, unsigned max_level)
/*
 *  Underlying probability model: adaptive arithmetic coding using
 *  the level of a range as context.
 */
{
   coeff_t *coeff = default_alloc (rpf, dc_rpf, min_level, max_level);
   
   coeff->bits            = aac_bits;
   coeff->update          = aac_update;
   coeff->model_free      = aac_model_free;
   coeff->model_duplicate = aac_model_duplicate;
   coeff->model		  = aac_model_alloc (coeff);
   
   return coeff;
}

static real_t
aac_bits (const real_t *used_coeff, const word_t *used_states,
	  unsigned level, const coeff_t *coeff)
{
   real_t	bits  = 0;		/* # bits to store coefficients */
   unsigned	edge;
   int		state;
   word_t      *counts;
   aac_model_t *model = (aac_model_t *) coeff->model;

   counts = model->counts
	    + (1 << (1 + coeff->dc_rpf->mantissa_bits))
	    + ((level - coeff->min_level)
	       * (1 << (1 + coeff->rpf->mantissa_bits)));
   
   for (edge = 0; isedge (state = used_states [edge]); edge++)
      if (state)
	 bits -= log2 (counts [rtob (used_coeff [edge], coeff->rpf)]
		       / (real_t) model->totals [level
						- coeff->min_level + 1]);
      else
	 bits -= log2 (model->counts [rtob (used_coeff [edge], coeff->dc_rpf)]
		       / (real_t) model->totals [0]);
   
   return bits;
}

static void
aac_update (const real_t *used_coeff, const word_t *used_states,
	    unsigned level, coeff_t *coeff)
{
   unsigned	edge;
   int		state;
   word_t      *counts;
   aac_model_t *model = (aac_model_t *) coeff->model;

   counts = model->counts
	    + (1 << (1 + coeff->dc_rpf->mantissa_bits))
	    + ((level - coeff->min_level)
	       * (1 << (1 + coeff->rpf->mantissa_bits)));

   for (edge = 0; isedge (state = used_states [edge]); edge++)
      if (state)
      {
	 counts [rtob (used_coeff [edge], coeff->rpf)]++;
	 model->totals [level - coeff->min_level + 1]++;
      }
      else
      {
	 model->counts [rtob (used_coeff [edge], coeff->dc_rpf)]++;
	 model->totals [0]++;
      }
}

static void *
aac_model_duplicate (const coeff_t *coeff, const void *model)
{
   aac_model_t *src = (aac_model_t *) model;
   aac_model_t *dst = aac_model_alloc (coeff);

   memcpy (dst->counts, src->counts,
	   sizeof (word_t) * ((coeff->max_level - coeff->min_level + 1)
			      * (1 << (1 + coeff->rpf->mantissa_bits))
			      + (1 << (1 + coeff->dc_rpf->mantissa_bits))));
   memcpy (dst->totals, src->totals,
	   sizeof (word_t) * (coeff->max_level - coeff->min_level + 1 + 1));
   
   return dst;
}

static void *
aac_model_alloc (const coeff_t *coeff)
{
   aac_model_t *model;
   unsigned	size = (coeff->max_level - coeff->min_level + 1)
		       * (1 << (1 + coeff->rpf->mantissa_bits))
		       + (1 << (1 + coeff->dc_rpf->mantissa_bits));
   
   model 	 = Calloc (1, sizeof (aac_model_t));
   model->counts = Calloc (size, sizeof (word_t));
   model->totals = Calloc (coeff->max_level - coeff->min_level + 1 + 1,
			   sizeof (word_t));
   /*
    *  Initialize model
    */
   {
      unsigned  n;
      word_t   *ptr = model->counts;
      
      for (n = size; n; n--)
	 *ptr++ = 1;
      model->totals [0] = 1 << (1 + coeff->dc_rpf->mantissa_bits);
      for (n = coeff->min_level; n <= coeff->max_level; n++)
	 model->totals [n - coeff->min_level + 1]
	    = 1 << (1 + coeff->rpf->mantissa_bits);
   }
   
   return (void *) model;
}

static void
aac_model_free (void *model)
{
   aac_model_t	*aac_model = (aac_model_t *) model;

   if (aac_model)
   {
      Free (aac_model->counts);
      Free (aac_model->totals);
      Free (aac_model);
   }
}

/*****************************************************************************
				default functions
*****************************************************************************/

static coeff_t *
default_alloc (rpf_t *rpf, rpf_t *dc_rpf,
	       unsigned min_level, unsigned max_level)
{
   coeff_t *coeff = Calloc (1, sizeof (coeff_t));

   coeff->rpf 	      	  = rpf;
   coeff->dc_rpf       	  = dc_rpf;
   coeff->min_level	  = min_level;
   coeff->max_level	  = max_level;
   coeff->model	      	  = NULL;
   coeff->bits	      	  = NULL;
   coeff->update      	  = NULL;
   coeff->free	      	  = default_free;
   coeff->model_free  	  = default_model_free;
   coeff->model_duplicate = default_model_duplicate;
   
   return coeff;
}

static void
default_free (coeff_t *coeff)
{
   coeff->model_free (coeff->model);
   Free (coeff);
}

static void *
default_model_duplicate (const coeff_t *coeff, const void *model)
{
   return NULL;
}

static void
default_model_free (void *model)
{
   if (model)
      Free (model);
}