/*
 *  mc.c:       Input of motion compensation
 *
 *  written by: Michael Unger
 *              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:13 $
 *  $Author: hafner $
 *  $Revision: 5.1 $
 *  $State: Exp $
 */

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

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

#include "wfa.h"
#include "bit-io.h"
#include "misc.h"
#include "mvcode.h"

#include "mc.h"

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

                             local variables

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

typedef struct huff_node
{
   int               code_index;        /* leaf if index >= 0 */
   struct huff_node *left;              /* follow if '0' bit read */
   struct huff_node *right;             /* follow if '1' bit read */
   int               index_set [34];
} huff_node_t;

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

                                prototypes

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

static void
decode_mc_tree (frame_type_e frame_type, unsigned max_state,
                wfa_t *wfa, bitfile_t *input);
static void
decode_mc_coords (unsigned max_state, wfa_t *wfa, bitfile_t *input);
static int
get_mv (int f_code, huff_node_t *hn, bitfile_t *input);
static huff_node_t *
create_huff_tree (void);
static void
create_huff_node (huff_node_t *hn, int bits_processed);

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

                                public code

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

void
read_mc (frame_type_e frame_type, wfa_t *wfa, bitfile_t *input)
/*
 *  Read motion compensation information of the 'input' stream.
 *  Depending on 'frame_type' different decoding methods are used.
 *
 *  No return value.
 *
 *  Side effects:
 *      'wfa->mv_tree' is filled with the decoded values.
 */
{
   unsigned max_state = wfa->wfainfo->color
                        ? wfa->tree [wfa->tree [wfa->root_state][0]][0]
                        : wfa->states;

   decode_mc_tree (frame_type, max_state, wfa, input);
   decode_mc_coords (max_state, wfa, input);
}

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

                                private code

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

static void
decode_mc_tree (frame_type_e frame_type, unsigned max_state,
                wfa_t *wfa, bitfile_t *input)
/*
 *  Read tree of motion compensation decisions of the 'input' stream.
 *  Depending on 'frame_type' different decoding methods are used.
 *  'max_state' is the last state with motion compensation infos.
 *
 *  No return value.
 *
 *  Side effects:
 *      'wfa->mv_tree' is filled with decoded values.
 */
{
   unsigned  state;                     /* current state */
   unsigned *queue;                     /* states in breadth first order */
   unsigned  last;                      /* last node
                                           (update for each new node) */

   /*
    *  Traverse tree in breadth first order (starting at level
    *  'wfa->wfainfo->p_max_level'). Use a queue to store the children
    *  of each node ('last' is the next free queue element).
    */
   queue = Calloc (MAXSTATES, sizeof (unsigned));
   for (last = 0, state = wfa->basis_states; state < max_state; state++)
      if (wfa->level_of_state [state] - 1 == (int) wfa->wfainfo->p_max_level)
         queue [last++] = state;        /* init level 'p_max_level' */

   if (frame_type == P_FRAME)
   {
      unsigned label;                   /* current label */
      unsigned current;                 /* current node to process */

      for (current = 0; current < last; current++)
         for (label = 0; label < MAXLABELS; label++)
         {
            state = queue[current];
            if (wfa->x [state][label]   /* process visible states only */
                +  width_of_level (wfa->level_of_state [state] - 1)
                <= wfa->wfainfo->width
                &&
                wfa->y [state][label]
                +  height_of_level (wfa->level_of_state [state] - 1)
                <= wfa->wfainfo->height)
            {
               wfa->mv_tree [state][label].type
                  = get_bit (input) ? NONE : FORWARD;
            }
            else
               wfa->mv_tree [state][label].type = NONE;
            if (wfa->mv_tree [state][label].type == NONE &&
                !isrange (wfa->tree [state][label]) &&
                wfa->level_of_state [state] - 1 >=
                (int) wfa->wfainfo->p_min_level)
               queue [last++] = wfa->tree [state][label]; /* append child  */
         }
   }
   else
   {
      unsigned label;                   /* current label */
      unsigned current;                 /* current node to process */

      for (current = 0; current < last; current++)
         for (label = 0; label < MAXLABELS; label++)
         {
            state = queue[current];
            if (wfa->x [state][label]   /* process visible states only */
                + width_of_level (wfa->level_of_state [state] - 1)
                > wfa->wfainfo->width
                ||
                wfa->y [state][label]
                + height_of_level (wfa->level_of_state [state] - 1)
                > wfa->wfainfo->height)
               wfa->mv_tree[state][label].type = NONE;
            else if (get_bit (input))   /* 1   */
               wfa->mv_tree[state][label].type = NONE;
            else if (get_bit (input))   /* 01  */
               wfa->mv_tree[state][label].type = INTERPOLATED;
            else if (get_bit (input))   /* 001 */
               wfa->mv_tree[state][label].type = BACKWARD;
            else                        /* 000 */
               wfa->mv_tree[state][label].type = FORWARD;
            if (wfa->mv_tree[state][label].type == NONE &&
                !isrange (wfa->tree[state][label]) &&
                wfa->level_of_state[state] - 1
                >= (int) wfa->wfainfo->p_min_level)
               queue[last++] = wfa->tree[state][label]; /* append child  */
         }
   }

   INPUT_BYTE_ALIGN (input);
   Free (queue);
}

static void
decode_mc_coords (unsigned max_state, wfa_t *wfa, bitfile_t *input)
/*
 *  Read motion vector coordinates of the 'input' stream. They are stored
 *  with the static Huffman code of the MPEG and H.263 standards.
 *  'max_state' is the last state with motion compensation infos.
 *
 *  No return value.
 *
 *  Side effects:
 *      'wfa->mv_tree' is filled with decoded values.
 */
{
   unsigned            label;           /* current label */
   unsigned            state;           /* current state */
   mv_t               *mv;              /* current motion vector */
   static huff_node_t *huff_mv_root = NULL; /* root of huffman tree */

   if (huff_mv_root == NULL)
      huff_mv_root = create_huff_tree ();

   for (state = wfa->basis_states; state < max_state; state++)
      for (label = 0; label < MAXLABELS; label++)
      {
         mv = &wfa->mv_tree[state][label];
         switch (mv->type)
         {
            case NONE:
               break;
            case FORWARD:
               mv->fx = get_mv (1, huff_mv_root, input);
               mv->fy = get_mv (1, huff_mv_root, input);
               break;
            case BACKWARD:
               mv->bx = get_mv (1, huff_mv_root, input);
               mv->by = get_mv (1, huff_mv_root, input);
               break;
            case INTERPOLATED:
               mv->fx = get_mv (1, huff_mv_root, input);
               mv->fy = get_mv (1, huff_mv_root, input);
               mv->bx = get_mv (1, huff_mv_root, input);
               mv->by = get_mv (1, huff_mv_root, input);
               break;
         }
      }

   INPUT_BYTE_ALIGN (input);
}

static int
get_mv (int f_code, huff_node_t *hn, bitfile_t *input)
/*
 *  Decode next motion vector component in bitstream
 *  by traversing the huffman tree.
 */
{
   int vlc_code, vlc_code_magnitude, residual, diffvec;

   while (hn->code_index < 0)
   {
      if (hn->code_index == -2)
         error ("wrong huffman code !");
      if (get_bit (input))
         hn = hn->right;
      else
         hn = hn->left;
   }
   vlc_code = hn->code_index - 16;
   if (vlc_code == 0 || f_code == 1)
      return vlc_code;

   vlc_code_magnitude = abs (vlc_code) - 1;
   if (f_code <= 1)
      residual = 0;
   else
      residual = get_bits (input, f_code - 1);
   diffvec = (vlc_code_magnitude << (f_code - 1)) + residual + 1;

   return vlc_code > 0 ? diffvec : - diffvec;
}

static huff_node_t *
create_huff_tree (void)
/*
 *  Construct huffman tree from code table
 */
{
   unsigned     i;
   huff_node_t *huff_root = Calloc (1, sizeof (huff_node_t));

   /*
    *  The nodes' index set contains indices of all codewords that are
    *  still decodable by traversing further down from the node.
    *  (The root node has the full index set.)
    */

   for (i = 0; i < 33; i++)
      huff_root->index_set [i] = i;
   huff_root->index_set [i] = -1;       /* end marker */

   create_huff_node (huff_root, 0);

   return huff_root;
}

static void
create_huff_node (huff_node_t *hn, int bits_processed)
/*
 *  Create one node in the huffman tree
 */
{
   int lind = 0;                        /* next index of left huff_node */
   int rind = 0;                        /* next index of right huff_node */
   int code_len, i, ind;

   hn->code_index = -1;
   if (hn->index_set [0] < 0)           /* empty index set ? */
   {
      hn->code_index = -2;              /* error */
      return;
   }
   hn->left  = Calloc (1, sizeof (huff_node_t));
   hn->right = Calloc (1, sizeof (huff_node_t));

   for (i = 0; (ind = hn->index_set[i]) >= 0; i++)
   {
      code_len = mv_code_table[ind][1];
      if (code_len == bits_processed)   /* generate leaf */
      {
         hn->code_index = ind;
         Free (hn->left);
         Free (hn->right);
         return;
      }
      if (mv_code_table[ind][0] & (1 << (code_len - 1 - bits_processed)))
         hn->right->index_set[rind++] = ind;
      else
         hn->left->index_set[lind++] = ind;
   }
   hn->right->index_set[rind] = -1;     /* set end markers */
   hn->left->index_set[lind]  = -1;
   create_huff_node (hn->left, bits_processed + 1);
   create_huff_node (hn->right, bits_processed + 1);
}