/*
 *  wfalib.c:		Library functions both for encoding 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/07/18 15:57:28 $
 *  $Author: hafner $
 *  $Revision: 5.5 $
 *  $State: Exp $
 */

#define _DEFAULT_SOURCE 1 /* New name for SVID & BSD source defines */
#define _BSD_SOURCE 1   /* Make sure strdup() is in string.h */
#define _XOPEN_SOURCE 500  /* Make sure strdup() is in string.h */

#include "config.h"

#include <stdlib.h>
#include <string.h>

#include "pm_c_util.h"

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

#include "wfa.h"
#include "misc.h"
#include "wfalib.h"

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

				prototypes
  
*****************************************************************************/

static unsigned
xy_to_address (unsigned x, unsigned y, unsigned level, unsigned n);

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

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

wfa_t *
alloc_wfa (bool_t coding)
/*
 *  WFA constructor:
 *  Initialize the WFA structure 'wfa' and allocate memory.
 *  Flag 'coding' indicates whether WFA is used for coding or decoding.
 *
 *  Return value:
 *	pointer to the new WFA structure
 */
{
   wfa_t *wfa = Calloc (1, sizeof (wfa_t));
		 
   /*
    *  Allocate memory
    */
   wfa->final_distribution = Calloc (MAXSTATES, sizeof (real_t));
   wfa->level_of_state     = Calloc (MAXSTATES, sizeof (byte_t));
   wfa->domain_type        = Calloc (MAXSTATES, sizeof (byte_t));
   wfa->delta_state        = Calloc (MAXSTATES, sizeof (bool_t));
   wfa->tree               = Calloc (MAXSTATES * MAXLABELS, sizeof (word_t));
   wfa->x                  = Calloc (MAXSTATES * MAXLABELS, sizeof (word_t));
   wfa->y                  = Calloc (MAXSTATES * MAXLABELS, sizeof (word_t));
   wfa->mv_tree            = Calloc (MAXSTATES * MAXLABELS, sizeof (mv_t));
   wfa->y_state            = Calloc (MAXSTATES * MAXLABELS, sizeof (word_t));
   wfa->into               = Calloc (MAXSTATES * MAXLABELS * (MAXEDGES + 1),
				     sizeof (word_t));
   wfa->weight             = Calloc (MAXSTATES * MAXLABELS * (MAXEDGES + 1),
				     sizeof (real_t));
   wfa->int_weight         = Calloc (MAXSTATES * MAXLABELS * (MAXEDGES + 1),
				     sizeof (word_t));
   wfa->wfainfo            = Calloc (1, sizeof (wfa_info_t));;
   wfa->prediction         = Calloc (MAXSTATES * MAXLABELS, sizeof (byte_t));

   wfa->wfainfo->wfa_name   = NULL;
   wfa->wfainfo->basis_name = NULL;
   wfa->wfainfo->title 	    = strdup ("");
   wfa->wfainfo->comment    = strdup ("");

   /*
    *  Initialize structure
    */
   {
      unsigned  state, label;

      wfa->states       = 0;
      wfa->basis_states = 0;
      wfa->root_state   = 0;
      for (state = 0; state < MAXSTATES; state++) 
      {
	 wfa->final_distribution [state] = 0;
	 wfa->domain_type [state]        = 0;
	 for (label = 0; label < MAXLABELS; label++)
	 {
	    wfa->into [state][label][0] = NO_EDGE;
	    wfa->tree [state][label]    = RANGE;
	    wfa->y_state [state][label] = RANGE;
	 }
      }
   }

   if (coding)				/* initialize additional variables */
      wfa->y_column = Calloc (MAXSTATES * MAXLABELS, sizeof (byte_t));
   else
      wfa->y_column = NULL;
   
   return wfa;
}

void
free_wfa (wfa_t *wfa)
/*
 *  WFA destructor:
 *  Free memory of given 'wfa'.
 *
 *  No return value.
 *
 *  Side effects:
 *	'wfa' struct is discarded.
 */
{
   if (wfa->wfainfo->wfa_name)
      Free (wfa->wfainfo->wfa_name);
   if (wfa->wfainfo->basis_name)
      Free (wfa->wfainfo->basis_name);
   if (wfa->wfainfo->title)
      Free (wfa->wfainfo->title);
   if (wfa->wfainfo->comment)
      Free (wfa->wfainfo->comment);

   Free (wfa->final_distribution);
   Free (wfa->level_of_state);
   Free (wfa->domain_type);
   Free (wfa->tree);
   Free (wfa->x);
   Free (wfa->y);
   Free (wfa->mv_tree);
   Free (wfa->y_state);
   Free (wfa->into);
   Free (wfa->weight);
   Free (wfa->int_weight);
   Free (wfa->wfainfo);
   Free (wfa->prediction);
   Free (wfa->delta_state);
   if (wfa->y_column)
      Free (wfa->y_column);
   Free (wfa);
}

real_t 
compute_final_distribution (unsigned state, const wfa_t *wfa)
/*
 *  Compute the final distribution of the given 'state'.
 *  Uses the fact that the generated 'wfa' is average preserving.
 *
 *  Return value:
 *	final distribution
 */
{
   unsigned label;
   real_t   final = 0;

   for (label = 0; label < MAXLABELS; label++)
   {
      unsigned edge;
      int      domain;
      
      if (ischild (domain = wfa->tree [state][label]))
	 final += wfa->final_distribution [domain];
      for (edge = 0; isedge (domain = wfa->into [state][label][edge]); edge++)
	 final += wfa->weight [state][label][edge]
		  * wfa->final_distribution [domain];
   }
   
   return final / MAXLABELS;
}

word_t *
compute_hits (unsigned from, unsigned to, unsigned n, const wfa_t *wfa)
/*
 *  Selects the 'n' most popular domain images of the given 'wfa'.
 *  Consider only linear combinations of state images
 *  {i | 'from' <= i <= 'to'}. I.e. domains are in {i | from <= i < 'to'}
 *  Always ensure that state 0 is among selected states even though from
 *  may be > 0.
 *  
 *  Return value:
 *	pointer to array of the most popular state images
 *	sorted by increasing state numbers and terminated by -1
 */
{
   word_t   *domains;
   unsigned  state, label, edge;
   int       domain;
   pair_t   *hits = Calloc (to, sizeof (pair_t));

   for (domain = 0; domain < (int) to; domain++)
   {
      hits [domain].value = domain;
      hits [domain].key   = 0;
   }
   
   for (state = from; state <= to; state++)
      for (label = 0; label < MAXLABELS; label++)
	 for (edge = 0; isedge (domain = wfa->into [state][label][edge]);
	      edge++)
	    hits [domain].key++;

   qsort (hits + 1, to - 1, sizeof (pair_t), sort_desc_pair);

   n       = MIN(to, n);
   domains = Calloc (n + 1, sizeof (word_t));

   for (domain = 0; domain < (int) n && (!domain || hits [domain].key);
	domain++)
      domains [domain] = hits [domain].value;
   if (n != domain)
      debug_message ("Only %d domains have been used in the luminance.",
		     domain);
   n = domain;
   qsort (domains, n, sizeof (word_t), sort_asc_word);
   domains [n] = -1;
   
   Free (hits);
   
   return domains;
}

void
append_edge (unsigned from, unsigned into, real_t weight,
	     unsigned label, wfa_t *wfa)
/*
 *  Append an edge from state 'from' to state 'into' with
 *  the given 'label' and 'weight' to the 'wfa'.
 *
 *  No return value.
 *
 *  Side effects:
 *	'wfa' structure is changed.
 */
{
   unsigned new;			/* position of the new edge */
   unsigned edge;

   /*
    *  First look where to insert the new edge:
    *  edges are sorted by increasing 'into' values
    */
   for (new = 0; (isedge (wfa->into [from][label][new])
		  && wfa->into [from][label][new] < (int) into); new++)
      ;
   /*
    *  Move the edges 'n' to position 'n+1', for n = max, ..., 'new'
    */
   for (edge = new; isedge (wfa->into [from][label][edge]); edge++)
      ;
   for (edge++; edge != new; edge--)
   {
      wfa->into [from][label][edge]    = wfa->into [from][label][edge - 1];
      wfa->weight [from][label][edge]  = wfa->weight [from][label][edge - 1];
      wfa->int_weight [from][label][edge]
	 = wfa->int_weight [from][label][edge - 1];
   }
   /*
    *  Insert the new edge
    */
   wfa->into [from][label][edge]       = into;
   wfa->weight [from][label][edge]     = weight;
   wfa->int_weight [from][label][edge] = weight * 512 + 0.5;
}

void 
remove_states (unsigned from, wfa_t *wfa)
/* 
 *  Remove 'wfa' states 'wfa->basis_states',...,'wfa->states' - 1.
 *
 *  No return value.
 *
 *  Side effects:
 *	'wfa' structure is cleared for the given states.
 */
{
   unsigned state;

   for (state = from; state < wfa->states; state++)
   {
      unsigned label;
      
      for (label = 0; label < MAXLABELS; label++) 
      {
	 wfa->into [state][label][0]      = NO_EDGE;
	 wfa->tree [state][label]         = RANGE;
	 wfa->prediction [state][label]   = FALSE;
	 wfa->y_state [state][label]      = RANGE;
	 wfa->mv_tree [state][label].type = NONE;
	 wfa->mv_tree [state][label].fx   = 0;
	 wfa->mv_tree [state][label].fy   = 0;
	 wfa->mv_tree [state][label].bx   = 0;
	 wfa->mv_tree [state][label].by   = 0;
      }
      wfa->domain_type [state] = 0;
      wfa->delta_state [state] = FALSE;
   }

   wfa->states = from;
}

void
copy_wfa (wfa_t *dst, const wfa_t *src)
/*
 *  Copy WFA struct 'src' to WFA struct 'dst'.
 *
 *  No return value.
 *
 *  Side effects:
 *	'dst' is filled with same data as 'src'
 *
 *  NOTE: size of WFA 'dst' must be at least size of WFA 'src'
 */
{
   unsigned state;

   memset (dst->final_distribution, 0, MAXSTATES * sizeof (real_t));
   memset (dst->level_of_state, 0, MAXSTATES * sizeof (byte_t));
   memset (dst->domain_type, 0, MAXSTATES * sizeof (byte_t));
   memset (dst->mv_tree, 0, MAXSTATES * MAXLABELS * sizeof (mv_t));
   memset (dst->tree, 0, MAXSTATES * MAXLABELS * sizeof (word_t));
   memset (dst->x, 0, MAXSTATES * MAXLABELS * sizeof (word_t));
   memset (dst->y, 0, MAXSTATES * MAXLABELS * sizeof (word_t));
   memset (dst->y_state, 0, MAXSTATES * MAXLABELS * sizeof (word_t));
   memset (dst->into, NO_EDGE,
	   MAXSTATES * MAXLABELS * (MAXEDGES + 1) * sizeof (word_t));
   memset (dst->weight, 0,
	   MAXSTATES * MAXLABELS * (MAXEDGES + 1) * sizeof (real_t));
   memset (dst->int_weight, 0,
	   MAXSTATES * MAXLABELS * (MAXEDGES + 1) * sizeof (word_t));
   memset (dst->prediction, 0, MAXSTATES * MAXLABELS * sizeof (byte_t));
   memset (dst->delta_state, 0, MAXSTATES * sizeof (bool_t));
   if (dst->y_column)
      memset (dst->y_column, 0, MAXSTATES * MAXLABELS * sizeof (byte_t));

   for (state = 0; state < MAXSTATES; state++) /* clear WFA struct */
   {
      unsigned label;
      
      for (label = 0; label < MAXLABELS; label++)
      {
	 dst->into [state][label][0]      = NO_EDGE;
	 dst->tree [state][label]         = RANGE;
	 dst->mv_tree [state][label].type = NONE;
	 dst->y_state[state][label]       = RANGE;
      }
      dst->delta_state [state] = NO;
      dst->domain_type [state] = 0;
   }
   
   dst->frame_type   = src->frame_type;
   dst->states 	     = src->states;
   dst->basis_states = src->basis_states;
   dst->root_state   = src->root_state;

   memcpy (dst->wfainfo, src->wfainfo, sizeof (wfa_info_t));

   if (dst->states == 0)		/* nothing to do */
      return;

   memcpy (dst->final_distribution, src->final_distribution,
	   src->states * sizeof (real_t));
   memcpy (dst->level_of_state, src->level_of_state,
	   src->states * sizeof (byte_t));
   memcpy (dst->domain_type, src->domain_type,
	   src->states * sizeof (byte_t));
   memcpy (dst->delta_state, src->delta_state,
	   src->states * sizeof (bool_t));
   memcpy (dst->mv_tree, src->mv_tree,
	   src->states * MAXLABELS * sizeof (mv_t));
   memcpy (dst->tree, src->tree,
	   src->states * MAXLABELS * sizeof (word_t));
   memcpy (dst->x, src->x,
	   src->states * MAXLABELS * sizeof (word_t));
   memcpy (dst->y, src->y,
	   src->states * MAXLABELS * sizeof (word_t));
   memcpy (dst->y_state, src->y_state,
	   src->states * MAXLABELS * sizeof (word_t));
   memcpy (dst->into, src->into,
	   src->states * MAXLABELS * (MAXEDGES + 1) * sizeof (word_t));
   memcpy (dst->weight, src->weight,
	   src->states * MAXLABELS * (MAXEDGES + 1) * sizeof (real_t));
   memcpy (dst->int_weight, src->int_weight,
	   src->states * MAXLABELS * (MAXEDGES + 1) * sizeof (word_t));
   memcpy (dst->prediction, src->prediction,
	   src->states * MAXLABELS * sizeof (byte_t));
   if (dst->y_column)
      memcpy (dst->y_column, src->y_column,
	      src->states * MAXLABELS * sizeof (byte_t));
}

void
locate_subimage (unsigned orig_level, unsigned level, unsigned bintree,
		 unsigned *x, unsigned *y, unsigned *width, unsigned *height)
/*
 *  Compute pixel coordinates of the subimage which 'bintree' address is given.
 *  The level of the original image is 'orig_level' and the level of the
 *  subimage is 'level'.
 *
 *  No return value.
 *
 *  Side effects:
 *	'*x', '*y'		coordinates of the upper left corner
 *      '*width', '*height'	size of image
 */
{
   /*
    *  Compute coordinates of the subimage
    */
   *x = *y = 0;				/* start at NW corner */
   *width  = width_of_level (level);
   *height = height_of_level (level);

   if (level > orig_level)
   {
      error ("size of tile must be less or equal than image size.");
      return;
   }
   else if (bintree >= (unsigned) (1 << (orig_level - level)))
   {
      error ("address out of bounds.");
      return;
   }
   else if (level < orig_level)
   {
      unsigned mask;			/* mask for bintree -> xy conversion */
      bool_t   hor;			/* 1 next subdivision is horizontal
					   0 next subdivision is vertical */
      unsigned l = orig_level - 1;	/* current level */
      
      hor = orig_level % 2;		/* start with vertival subdivision
					   for square image and vice versa */
   
      for (mask = 1 << (orig_level - level - 1); mask; mask >>= 1, hor = !hor)
      {
	 if (bintree & mask)		/* change coordinates */
	 {
	    if (hor)			/* horizontal subdivision */
	       *y += height_of_level (l);
	    else			/* vertical subdivision */
	       *x += width_of_level (l);
	 }
	 l--;
      }
   }
}

void
compute_spiral (int *vorder, unsigned image_width, unsigned image_height,
		unsigned tiling_exp, bool_t inc_spiral)
/*
 *  Compute image tiling with spiral order.
 *  'inc_spiral' specifies whether the spiral starts in the middle
 *  of the image (TRUE) or at the border (FALSE).
 *  'image_width'x'image_height' define the size of the image.
 *  The image is split into 'tiling->exp' tiles.
 *
 *  No return value.
 *
 *  Side effects:
 *	vorder[] is filled with tiling permutation
 */
{
   unsigned x, y;			/* current position */
   unsigned xmin, xmax, ymin, ymax;	/* boundaries for current line */
   unsigned width, height;		/* offset for each tile */
   unsigned lx, ly, level;		/* level x and y */
   unsigned tiles;			/* total number of tiles */
   unsigned address;			/* bintree address */
   
   lx     = log2 (image_width - 1) + 1;
   ly     = log2 (image_height - 1) + 1;
   level  = MAX(lx, ly) * 2 - ((ly == lx + 1) ? 1 : 0);
   tiles  = 1 << tiling_exp;		/* Number of image tiles */
   width  = width_of_level (level - tiling_exp);
   height = height_of_level (level - tiling_exp);
   for (address = 0; address < tiles; address++)
   {
      unsigned x0, y0, width, height;
      
      locate_subimage (level, level - tiling_exp, address,
		       &x0, &y0, &width, &height);
      vorder [address] = (x0 < image_width && y0 < image_height) ? 0 : -1;
   }

   xmin    = 0;
   xmax    = width_of_level (level);
   ymin    = 0;
   ymax    = height_of_level (level);
   address = 0;

   /*
    *  1234
    *  CDE5  Traverse image in spiral order 
    *  BGF6  starting at the top left corner
    *  A987
    */
   while (TRUE)
   {
      for (x = xmin, y = ymin; x < xmax; x += width) /* W>E */
      {
	 while (vorder [address] == -1)
	    address++;
	 if (x < image_width && y < image_height) /* valid range */
	    vorder [address++] = xy_to_address (x, y, level, tiling_exp);
	 while (address < tiles && vorder [address] == -1)
	    address++;
      }
      ymin += height;

      if (address >= tiles)
	 break;
      
      for (x = xmax - width, y = ymin; y < ymax; y += height) /* N>S  */
      {
	 while (vorder [address] == -1)
	    address++;
	 if (x <= image_width && y <= image_height) /* valid range */
	    vorder [address++] = xy_to_address (x, y, level, tiling_exp);
	 while (address < tiles && vorder [address] == -1)
	    address++;
      }
      xmax -= width;

      if (address >= tiles)
	 break;

      for (x = xmax - width, y = ymax - width; x >= xmin; x -= width) /* E<W */
      {
	 while (vorder [address] == -1)
	    address++;
	 if (x <= image_width && y <= image_height) /* valid range */
	    vorder [address++] = xy_to_address (x, y, level, tiling_exp);
	 while (address < tiles && vorder [address] == -1)
	    address++;
      }
      ymax -= height;

      if (address >= tiles)
	 break;

      for (x = xmin, y = ymax - height; y >= ymin; y -= height)	/* S>N */
      {
	 while (vorder [address] == -1)
	    address++;
	 if (x <= image_width && y <= image_height) /* valid range */
	    vorder [address++] = xy_to_address (x, y, level, tiling_exp);
	 while (address < tiles && vorder [address] == -1)
	    address++;
      }
      xmin += width;
	 
      if (address >= tiles)
	 break;
   }

   if (inc_spiral)
   {
      int i = 0, j = tiles - 1;

      while (i < j)
      {
	 int tmp;
	    
	 while (vorder [i] == -1)
	    i++;
	 while (vorder [j] == -1)
	    j--;
	    
	 tmp 	       = vorder [i];
	 vorder [i] = vorder [j];
	 vorder [j] = tmp;
	 i++;
	 j--;
      }
   }
   /*
    *  Print tiling info
    */
   {
      unsigned number;
      
      for (number = 0, address = 0; address < tiles; address++)
	 if (vorder [address] != -1)
	    debug_message ("number %d: address %d",
			   number++, vorder [address]);
   }
}

bool_t
find_range (unsigned x, unsigned y, unsigned band,
	    const wfa_t *wfa, unsigned *range_state, unsigned *range_label)
/*
 *  Find a range ('*range_state', '*range_label') that contains
 *  pixel ('x', 'y') in the iven color 'band'.
 *
 *  Return value:
 *	TRUE on success, or FALSE if there is no such range
 *
 *  Side effects:
 *	'*range_state' and '*range_label' are modified on success.
 */
{
   unsigned state, label;
   unsigned first_state, last_state;
   bool_t   success = NO;
   
   first_state = wfa->basis_states;
   last_state  = wfa->states;
   if (wfa->wfainfo->color)
      switch (band)
      {
	 case Y:
	    first_state = wfa->basis_states;
	    last_state  = wfa->tree [wfa->tree [wfa->root_state][0]][0];
	    break;
	 case Cb:
	    first_state = wfa->tree [wfa->tree [wfa->root_state][0]][0] + 1;
	    last_state  = wfa->tree [wfa->tree [wfa->root_state][0]][1];
	    break;
	 case Cr:
	    first_state = wfa->tree [wfa->tree [wfa->root_state][0]][1] + 1;
	    last_state  = wfa->states;
	    break;
	 default:
	    error ("unknown color component.");
      }

   for (state = first_state; state < last_state; state++)
      for (label = 0; label < MAXLABELS; label++)
	 if (isrange (wfa->tree [state][label]))
	    if (x >= wfa->x [state][label] && y >= wfa->y [state][label]
		&& x < (unsigned) (wfa->x [state][label]
			+ width_of_level (wfa->level_of_state [state] - 1))
		&& y < (unsigned) (wfa->y [state][label]
			+ height_of_level (wfa->level_of_state [state] - 1))) 
	    {
	       success      = YES;
	       *range_state = state;
	       *range_label = label;

	       return success;
	    }

   return success;
}

void
sort_ranges (unsigned state, unsigned *domain,
	     range_sort_t *rs, const wfa_t *wfa)
/*
 *  Generate list of ranges in coder order.
 *  'state' is the current state of the call tree while 'domain' is the
 *  index of the last added WFA state.
 *
 *  Side effects:
 *	'domain' is incremented after recursion returns
 *	'rs'	 is filled accordingly
 *
 *  No return value.
 */
{
   unsigned label;
   
   for (label = 0; label < MAXLABELS; label++)
   {
      if (isrange (wfa->tree [state][label]))
	 rs->range_subdivided [rs->range_no] = NO;
      else
      {
	 sort_ranges (wfa->tree [state][label], domain, rs, wfa);
	 rs->range_subdivided [rs->range_no] = YES;
      }

      rs->range_state [rs->range_no]      = state;
      rs->range_label [rs->range_no]      = label;
      rs->range_max_domain [rs->range_no] = *domain;
      while (!usedomain (rs->range_max_domain [rs->range_no], wfa))
	 rs->range_max_domain [rs->range_no]--;

      if (label == 1 || !rs->range_subdivided [rs->range_no])
	 rs->range_no++;
   }
   
   (*domain)++;
}

bool_t
locate_delta_images (wfa_t *wfa)
/*
 *  Locate all WFA states that are part of a delta approximation.
 *  I.e., these states are assigned to ranges that have been predicted
 *  via MC or ND.
 *
 *  Return value:
 *	TRUE	at least one state is part of a delta approximation
 *	FALSE	no delta approximations in this WFA
 *
 *  Side effects:
 *	'wfa->delta [state][label]' is set accordingly.
 */
{
   unsigned state, label;
   bool_t   delta = NO;

   for (state = wfa->root_state; state >= wfa->basis_states; state--)
      wfa->delta_state [state] = NO;

   for (state = wfa->root_state; state >= wfa->basis_states; state--)
      for (label = 0; label < MAXLABELS; label++)
	 if (ischild (wfa->tree [state][label]))
	    if (wfa->mv_tree [state][label].type != NONE
		|| isedge (wfa->into [state][label][0])
		|| wfa->delta_state [state])
	    {
	       delta = YES;
	       wfa->delta_state [wfa->tree [state][label]] = YES;
	    }

   return delta;
}

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

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

static unsigned
xy_to_address (unsigned x, unsigned y, unsigned level, unsigned n)
/*
 *  Compute bintree address of subimage at coordinates ('x', 'y').
 *  Size of original image is determined by 'level'.
 *  'n' specifies number of iterations.
 *
 *  Return value:
 *	address of subimage
 */ 
{ 
   unsigned address = 0;

   while (n--)
   {
      address <<= 1;
      if (--level % 2) 
      {
	 if (x & width_of_level (level))
	    address++;
      }
      else
      {
	 if (y & height_of_level (level))
	    address++;
      }
   }
   
   return address;
}