/******************************************************************************
 *
 * Copyright Jill R. Goldschneider, 1998
 *
 * This work was partially supported by a NASA Graduate Student
 * Fellowship in Global Change Research, an NSF Young Investigator Award,
 * and U.~S. Army Research Grant DAAH004-96-1-0255.
 *
 *****************************************************************************/

/******************************************************************************
 *
 * FILE          wpt_bitalloc_util.c
 * AUTHOR        Jill R. Goldschneider
 * DATE          February 1997
 * REVISIONS     February 1998 - update documentation
 * DESCRIPTION   WPT tree pruning functions
 *               void      initialize_tree(TreeNode *node)
 *               TreeNode *min_slope_node(TreeNode *node)
 *               double    prune(TreeNode *node)
 *               void      wpt_bitalloc(TreeNode *root, double stoppingrate,
 *			                double *distortion,  double *rate)
 *
 *****************************************************************************/
#include "wpt.h"

static double minimum(double num1, double num2);

/******************************************************************************
 * DOCUMENTATION ********************************************************
 ******************************************************************************
   NAME initialize_tree

   DESCRIPTION This routine initializes the WPT data structure for
   systematic joint best-bases selection and optimal bit allocation.

   ARGUMENTS
      IOARG  node   the root node of the WPT tree

   RETURN

   ALGORITHM initialize_tree does a post-order traversal of the WPT.
   The node values that are initialized are those representing the
   {\em branch} distortion (node$\rightarrow$distortion), rate
   (node$\rightarrow$rate), minimum slope
   (node$\rightarrow\lambda_{\min}$), and best-basis decision
   (node$\rightarrow$split).  The minimum branch slope computation
   will find those slopes which bridge the two convex hulls of the
   node's data and the combined best basis of the node's children.

   AUTHOR Jill R. Goldschneider

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

void initialize_tree(TreeNode *node)
{
  int    i;
  double temp_dist;
  double temp_rate;
  double temp_cost1;
  double temp_cost2;
  double temp_cost_local1;
  double temp_cost_local2;
  double slope1;
  double slope2;

  /* the node is a leaf node */
  if (node->depth == treelevel) {
    node->distortion = node->data->distortion;
    node->rate = node->data->rate;
    node->lambda_min = node->data->lambda;
    node->split = FALSE;
  }

  /* node is internal */
  else {
    /* initialize the node's children */
    for (i = 0; i < treedim; i++) {
      initialize_tree(node->child[i]);
    }

    /* initialize internal node */
    /* find minimum slope */
    node->lambda_min = node->data->lambda;
    for (i = 0; i < treedim; i++) {
      node->lambda_min = minimum(node->lambda_min, node->child[i]->lambda_min);
    }

    /* compute distortion, rate, and cost (assume lambda's are 0) */
    slope1 = 0.0;
    slope2 = node->lambda_min;
    temp_dist = 0.0;
    temp_rate = 0.0;
    temp_cost1 = 0.0;
    temp_cost2 = 0.0;
    for (i = 0; i < treedim; i++) {
      temp_rate += node->child[i]->rate;
      temp_dist += node->child[i]->distortion;
      temp_cost1 += node->child[i]->distortion +
	slope1 * node->child[i]->rate;
      temp_cost2 += node->child[i]->distortion +
	slope2 * node->child[i]->rate;
    }
    temp_cost_local1 = node->data->distortion
      + slope1 * node->data->rate;
    temp_cost_local2 = node->data->distortion
      + slope2 * node->data->rate;

    /* if costs are the same, use smaller rate, i.e. take only convex vertex */
    if(temp_cost1 == temp_cost_local1) {
      if(temp_rate < node->data->rate) {
	node->split = TRUE;
      }
      else {
	node->split = FALSE;
      }
    }
    else if(temp_cost1 < temp_cost_local1) {
      node->split = TRUE;
      /* if the second costs are reversed, a split is to occur,
       compute the missing link lambda value */
      if(temp_cost_local2 <= temp_cost2) {
	node->lambda_min = -1.0 * ((temp_dist - node->data->distortion) /
	  (temp_rate - node->data->rate));
      }
    }
    else {
      node->split = FALSE;
      /* if the second costs are reversed, a split is to occur,
       compute the missing link lambda value */
      if(temp_cost2 <= temp_cost_local2) {
	node->lambda_min = -1.0 * ((temp_dist - node->data->distortion) /
	  (temp_rate - node->data->rate));
      }
    }

    /* it is better to encode using children */
    if(node->split == TRUE) {
      node->rate = temp_rate;
      node->distortion = temp_dist;
    }
    /* it is better to encode using this node */
    else {
      node->rate = node->data->rate;
      node->distortion = node->data->distortion;
    }
  }
}

/******************************************************************************
 * DOCUMENTATION ********************************************************
 ******************************************************************************
   NAME min_slope_node

   DESCRIPTION min_slope_node finds the WPT node with the next minimum
   slope.

   ARGUMENTS
      IARG  node  root of tree

   RETURN min_slope_node returns a pointer to the minimum slope WPT
   node.

   ALGORITHM min_slope_node follows the minimum slope path through the
   tree to find the minimum slope node which satisfies
   $\text{node}\rightarrow\lambda_{\min} =
   \text{node}\rightarrow\text{data}\rightarrow\lambda$.

   AUTHOR Jill R. Goldschneider

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

TreeNode *min_slope_node(TreeNode *node)
{
  int      i;

  /* current node has minimum slope, return this node */
  if(node->lambda_min == node->data->lambda) {
    return(node);
  }
  /* current node is a leaf slope, return this node */
  if(!(node->child[0])) {
    /* this should never happen, since leaf nodes are picked up above */
    printf("shouldn't get here! \n");
    return(node);
  }

  for (i = 0; i < treedim; i++) {
    /* go to first child that has the same lambda_min */
    if(node->lambda_min == node->child[i]->lambda_min) {
      node = node->child[i];
      return(min_slope_node(node));
    }
  }

  /* otherwise a missing link */
  /* printf("depth = %d, node->lambda_min = %f, node->data->lambda = %f\n",
            node->depth, node->lambda_min, node->data->lambda); */

  return(node);
}

/******************************************************************************
 * DOCUMENTATION ********************************************************
 ******************************************************************************
   NAME prune

   DESCRIPTION prune implements one iteration of systematic joint
   best-basis selection and optimal bit allocation.  It changes from
   the current vertex on the global lower convex hull to the next
   vertex on the global lower convex hull.

   ARGUMENTS
      IOARG  node    the node to be pruned

   RETURN The minimum cost $J = D + \lambda R$ associated with the
   best basis/optimal bit allocation is returned.

   ALGORITHM If the minimum slope is equal to the minimum slope of the
   node's data's quantizer, then the node's data's quantizer is
   pruned.  Otherwise, the slope represents a bridge between the two
   convex hulls, and there will be a change in best basis.  Based on
   this information, the best-basis decision is made for the given
   node and all of its ancestors.  While this is done, the next
   minimum slope is propagated to the root of the tree.

   AUTHOR Jill R. Goldschneider

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

double prune(TreeNode *node)
{
  int       i;
  double    temp_dist;
  double    temp_rate;
  double    temp_cost1;
  double    temp_cost2;
  double    temp_cost_local1;
  double    temp_cost_local2;
  double    slope1;
  double    slope2;

  /* prune at this node only if the link is not a missing link*/
  slope1 = node->lambda_min;
  if (slope1 == node->data->lambda) {
    node->data = node->data->next; /* prune */
  }
  /* link is a missing link, do not prune, but follow the missing link */
  else {
    /* update this node */
    /* the costs of slope1 must be the same, so choose calculate
       the next minimum slope and choolse the lower rate quantizer */

    if(node->depth == treelevel) { /* this should never happen */
      fprintf(stderr, "%s: warning this should not occur \n", programname);
    }

    /* find minimum slope */
    node->lambda_min = node->data->lambda;
    for (i = 0; i < treedim; i++) {
      node->lambda_min = minimum(node->lambda_min, node->child[i]->lambda_min);
    }
    slope2 = node->lambda_min;

    /* compute distortion, rate, and cost */;
    temp_dist = 0.0;
    temp_rate = 0.0;
    for (i = 0; i < treedim; i++) {
      temp_rate += node->child[i]->rate;
      temp_dist += node->child[i]->distortion;
    }

    /* the costs are the same for slope1, choose the lower slope node */
    /* it is better to encode using children */
    if(temp_rate < node->data->rate) {
      node->split = TRUE;
      node->rate = temp_rate;
      node->distortion = temp_dist;
    }
    /* it is better to encode using this node */
    else {
      node->split = FALSE;
      node->rate = node->data->rate;
      node->distortion = node->data->distortion;
    }

    /* exit if node is root, no ancestors to update */
    if(node->depth == 0) {
      return(node->distortion + slope1 * node->rate);
    }
    /* if node is not root, go to parent */
    node = node->parent;
  }

  /* node is terminal, cannot be a missing link node */
  if(node->depth == treelevel) {
    node->distortion = node->data->distortion;
    node->rate = node->data->rate;
    node->split = FALSE;
    node->lambda_min = node->data->lambda;
    /* exit if node is root, no ancestors to update */
    if(node->depth == 0) {
      return(node->distortion + slope1 * node->rate);
    }
    /* if node is not root, go to parent */
    node = node->parent;
  }

  /* update the node and its ancestors */
  while(node->depth >= 0) {
    /* find minimum slope */
    node->lambda_min = node->data->lambda;
    for (i = 0; i < treedim; i++) {
      node->lambda_min = minimum(node->lambda_min, node->child[i]->lambda_min);
    }

    slope2 = node->lambda_min;

    /* compute distortion, rate, and cost */;
    temp_dist = 0.0;
    temp_rate = 0.0;
    temp_cost1 = 0.0;
    temp_cost2 = 0.0;
    for (i = 0; i < treedim; i++) {
      temp_rate += node->child[i]->rate;
      temp_dist += node->child[i]->distortion;
      temp_cost1 += node->child[i]->distortion +
	slope1 * node->child[i]->rate;
      temp_cost2 += node->child[i]->distortion +
	slope2 * node->child[i]->rate;
    }
    temp_cost_local1 = node->data->distortion +
      slope1 * node->data->rate;
    temp_cost_local2 = node->data->distortion +
      slope2 * node->data->rate;

    /* if costs are the same, use smaller rate, i.e. take only convex vertex */
    if(temp_cost1 == temp_cost_local1) {
      if(temp_rate < node->data->rate) {
	node->split = TRUE;
      }
      else {
	node->split = FALSE;
      }
    }
    else if(temp_cost1 < temp_cost_local1) {
      node->split = TRUE;
      if(temp_cost_local2 <= temp_cost2) {
	node->lambda_min = -1.0 * ((temp_dist - node->data->distortion) /
				   (temp_rate - node->data->rate));
      }
    }
    else {
      node->split = FALSE;
      /* if the second costs are reversed, a split is to occur,
	 compute the missing link lambda value */
      if(temp_cost2 <= temp_cost_local2) {
	node->lambda_min = -1.0 * ((temp_dist - node->data->distortion) /
				   (temp_rate - node->data->rate));
      }
    }

    /* it is better to encode using children */
    if(node->split == TRUE) {
      node->rate = temp_rate;
      node->distortion = temp_dist;
    }
    /* it is better to encode using this node */
    else {
      node->rate = node->data->rate;
      node->distortion = node->data->distortion;
    }

    /* exit if node is root, no ancestors to update */
    if(node->depth == 0) {
      return(minimum(temp_cost1, temp_cost_local1));
    }
    /* if node is not root, go to parent */
    else {
      node = node->parent;
    }
  }

  return(minimum(temp_cost1, temp_cost_local1));
}

/******************************************************************************
 * DOCUMENTATION ********************************************************
 ******************************************************************************
   NAME minimum

   DESCRIPTION The minimum value is returned.

   ARGUMENTS
      IARG  num1   first argument
      IARG  num2   second argument

   RETURN The minimum value $\min(\text{num1}, \text{num2})$ is returned.

   ALGORITHM

   AUTHOR Jill R. Goldschneider

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

static double minimum(double num1, double num2)
{
  if (num1 < num2) {
      return(num1);
  }
  else {
    return(num2);
  }
}

/******************************************************************************
 * DOCUMENTATION ********************************************************
 ******************************************************************************
   NAME wpt_bitalloc

   DESCRIPTION Systematic joint best-basis selection and optimal bit
   allocation algorithm.

   ARGUMENTS
      IOARG root            the root of the codebook tree
      IARG  stoppingrate    stopping rate
      IOARG rate            achieved rate
      IOARG distortion      achieved distortion

   RETURN

   ALGORITHM  While the rate is greater than the stopping rate, then
   the minimum slope WPT node is found using min_slope_node, and pruned
   using prune.

   AUTHOR Jill R. Goldschneider

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

void wpt_bitalloc(TreeNode *root, double stoppingrate,
		  double *rate,  double *distortion)
{
  double    min_slope;            /* current minimum R-D slope */
  double    cost;                 /* Lagrange functional cost */
  TreeNode *node;                 /* a node of codebook tree */


  /* prepare to begin */
  min_slope = 0.0;
  *rate = root->rate;
  *distortion = root->distortion;
  cost = *distortion;

  printf("\n");
  printf("         LAMBDA             RATE       DISTORTION"
	 "             COST\n");
  printf("%15f  %15f  %15f  %15f\n", min_slope, *rate, *distortion, cost);
  fflush(stdout);

  /* until the stopping rate is achieved, prune the tree by ... */

  /* can have zero rate, but a small lambda since other nodes may have lower
     lambda values */
  /* while ((root->lambda_min < HUGE) && (*rate > stoppingrate)) { */
  while (*rate > stoppingrate) {
    /* do a single iteration on the tree and output values */
    node = root;
    min_slope = root->lambda_min;
    node = min_slope_node(node);
    cost = prune(node);
    /* there may be no change in rate if PTSVQ's not on the WPT global
       lower convex hull are pruned */
    if(*rate != root->rate) {
      *rate = root->rate;
      *distortion = root->distortion;
      /* print data */
      printf("%15f  %15f  %15f  %15f\n", min_slope, *rate, *distortion, cost);
      fflush(stdout);
    }
  }
}