/*
 * Huffman coding optimization.
 *
 * This actually is optimization, in the sense that we find the best possible
 * Huffman table(s) for the given data.  We first scan the supplied data and
 * count the number of uses of each symbol that is to be Huffman-coded.
 * (This process must agree with the code above.)  Then we build an
 * optimal Huffman coding tree for the observed counts.
 */

#ifdef ENTROPY_OPT_SUPPORTED


/* These are static so htest_one_block can find 'em */
static long * dc_count_ptrs[NUM_HUFF_TBLS];
static long * ac_count_ptrs[NUM_HUFF_TBLS];


LOCAL void
gen_huff_coding (compress_info_ptr cinfo, HUFF_TBL *htbl, long freq[])
/* Generate the optimal coding for the given counts */
{
#define MAX_CLEN 32     /* assumed maximum initial code length */
  UINT8 bits[MAX_CLEN+1];   /* bits[k] = # of symbols with code length k */
  short codesize[257];      /* codesize[k] = code length of symbol k */
  short others[257];        /* next symbol in current branch of tree */
  int c1, c2;
  int p, i, j;
  long v;

  /* This algorithm is explained in section K.2 of the JPEG standard */

  MEMZERO(bits, SIZEOF(bits));
  MEMZERO(codesize, SIZEOF(codesize));
  for (i = 0; i < 257; i++)
	others[i] = -1;     /* init links to empty */
  
  freq[256] = 1;        /* make sure there is a nonzero count */
  /* including the pseudo-symbol 256 in the Huffman procedure guarantees
   * that no real symbol is given code-value of all ones, because 256
   * will be placed in the largest codeword category.
   */

  /* Huffman's basic algorithm to assign optimal code lengths to symbols */

  for (;;) {
	/* Find the smallest nonzero frequency, set c1 = its symbol */
	/* In case of ties, take the larger symbol number */
	c1 = -1;
	v = 1000000000L;
	for (i = 0; i <= 256; i++) {
	  if (freq[i] && freq[i] <= v) {
	v = freq[i];
	c1 = i;
	  }
	}

	/* Find the next smallest nonzero frequency, set c2 = its symbol */
	/* In case of ties, take the larger symbol number */
	c2 = -1;
	v = 1000000000L;
	for (i = 0; i <= 256; i++) {
	  if (freq[i] && freq[i] <= v && i != c1) {
	v = freq[i];
	c2 = i;
	  }
	}

	/* Done if we've merged everything into one frequency */
	if (c2 < 0)
	  break;
	
	/* Else merge the two counts/trees */
	freq[c1] += freq[c2];
	freq[c2] = 0;

	/* Increment the codesize of everything in c1's tree branch */
	codesize[c1]++;
	while (others[c1] >= 0) {
	  c1 = others[c1];
	  codesize[c1]++;
	}
	
	others[c1] = c2;        /* chain c2 onto c1's tree branch */
	
	/* Increment the codesize of everything in c2's tree branch */
	codesize[c2]++;
	while (others[c2] >= 0) {
	  c2 = others[c2];
	  codesize[c2]++;
	}
  }

  /* Now count the number of symbols of each code length */
  for (i = 0; i <= 256; i++) {
	if (codesize[i]) {
	  /* The JPEG standard seems to think that this can't happen, */
	  /* but I'm paranoid... */
	  if (codesize[i] > MAX_CLEN)
	ERREXIT(cinfo->emethods);
	/*(cinfo->emethods, "Huffman code size table overflow");*/

	  bits[codesize[i]]++;
	}
  }

  /* JPEG doesn't allow symbols with code lengths over 16 bits, so if the pure
   * Huffman procedure assigned any such lengths, we must adjust the coding.
   * Here is what the JPEG spec says about how this next bit works:
   * Since symbols are paired for the longest Huffman code, the symbols are
   * removed from this length category two at a time.  The prefix for the pair
   * (which is one bit shorter) is allocated to one of the pair; then,
   * skipping the BITS entry for that prefix length, a code word from the next
   * shortest nonzero BITS entry is converted into a prefix for two code words
   * one bit longer.
   */
  
  for (i = MAX_CLEN; i > 16; i--) {
	while (bits[i] > 0) {
	  j = i - 2;        /* find length of new prefix to be used */
	  while (bits[j] == 0)
	j--;
	  
	  bits[i] -= 2;     /* remove two symbols */
	  bits[i-1]++;      /* one goes in this length */
	  bits[j+1] += 2;       /* two new symbols in this length */
	  bits[j]--;        /* symbol of this length is now a prefix */
	}
  }

  /* Remove the count for the pseudo-symbol 256 from the largest codelength */
  while (bits[i] == 0)      /* find largest codelength still in use */
	i--;
  bits[i]--;
  
  /* Return final symbol counts (only for lengths 0..16) */
  MEMCOPY(htbl->bits, bits, SIZEOF(htbl->bits));
  
  /* Return a list of the symbols sorted by code length */
  /* It's not real clear to me why we don't need to consider the codelength
   * changes made above, but the JPEG spec seems to think this works.
   */
  p = 0;
  for (i = 1; i <= MAX_CLEN; i++) {
	for (j = 0; j <= 255; j++) {
	  if (codesize[j] == i) {
	htbl->huffval[p] = (UINT8) j;
	p++;
	  }
	}
  }
}


/* Process a single block's worth of coefficients */
/* Note that the DC coefficient has already been converted to a difference */

LOCAL void
htest_one_block (JBLOCK block, JCOEF block0,
		 long dc_counts[], long ac_counts[])
{
  register INT32 temp;
  register int nbits;
  register int k, r;
  
  /* Encode the DC coefficient difference per section F.1.2.1 */
  
  /* Find the number of bits needed for the magnitude of the coefficient */
  temp = block0;
  if (temp < 0) temp = -temp;
  
  for (nbits = 0; temp; nbits++)
	temp >>= 1;
  
  /* Count the Huffman symbol for the number of bits */
  dc_counts[nbits]++;
  
  /* Encode the AC coefficients per section F.1.2.2 */
  
  r = 0;            /* r = run length of zeros */
  
  for (k = 1; k < DCTSIZE2; k++) {
	if ((temp = block[k]) == 0) {
	  r++;
	} else {
	  /* if run length > 15, must emit special run-length-16 codes (0xF0) */
	  while (r > 15) {
	ac_counts[0xF0]++;
	r -= 16;
	  }
	  
	  /* Find the number of bits needed for the magnitude of the coefficient */
	  if (temp < 0) temp = -temp;
	  
	  for (nbits = 0; temp; nbits++)
	temp >>= 1;
	  
	  /* Count Huffman symbol for run length / number of bits */
	  ac_counts[(r << 4) + nbits]++;
	  
	  r = 0;
	}
  }

  /* If the last coef(s) were zero, emit an end-of-block code */
  if (r > 0)
	ac_counts[0]++;
}



/*
 * Trial-encode one MCU's worth of Huffman-compressed coefficients.
 */

LOCAL void
htest_encode (compress_info_ptr cinfo, JBLOCK *MCU_data)
{
  short blkn, ci;
  jpeg_component_info * compptr;

  /* Take care of restart intervals if needed */
  if (cinfo->restart_interval) {
	if (cinfo->restarts_to_go == 0) {
	  /* Re-initialize DC predictions to 0 */
	  for (ci = 0; ci < cinfo->comps_in_scan; ci++)
	cinfo->last_dc_val[ci] = 0;
	  /* Update restart state */
	  cinfo->restarts_to_go = cinfo->restart_interval;
	}
	cinfo->restarts_to_go--;
  }

  for (blkn = 0; blkn < cinfo->blocks_in_MCU; blkn++) {
	ci = cinfo->MCU_membership[blkn];
	compptr = cinfo->cur_comp_info[ci];
	/* NB: unlike the real entropy encoder, we may not change the input data */
	htest_one_block(MCU_data[blkn],
			(JCOEF) (MCU_data[blkn][0] - cinfo->last_dc_val[ci]),
			dc_count_ptrs[compptr->dc_tbl_no],
			ac_count_ptrs[compptr->ac_tbl_no]);
	cinfo->last_dc_val[ci] = MCU_data[blkn][0];
  }
}



/*
 * Find the best coding parameters for a Huffman-coded scan.
 * When called, the scan data has already been converted to a sequence of
 * MCU groups of quantized coefficients, which are stored in a "big" array.
 * The source_method knows how to iterate through that array.
 * On return, the MCU data is unmodified, but the Huffman tables referenced
 * by the scan components may have been altered.
 */

METHODDEF void
huff_optimize (compress_info_ptr cinfo, MCU_output_caller_ptr source_method)
/* Optimize Huffman-coding parameters (Huffman symbol table) */
{
  int i, tbl;
  HUFF_TBL **htblptr;

  /* Allocate and zero the count tables */
  /* Note that gen_huff_coding expects 257 entries in each table! */

  for (i = 0; i < NUM_HUFF_TBLS; i++) {
	dc_count_ptrs[i] = NULL;
	ac_count_ptrs[i] = NULL;
  }

  for (i = 0; i < cinfo->comps_in_scan; i++) {
	/* Create DC table */
	tbl = cinfo->cur_comp_info[i]->dc_tbl_no;
	if (dc_count_ptrs[tbl] == NULL) {
	  dc_count_ptrs[tbl] = (long *) (*cinfo->emethods->alloc_small)
					(257 * SIZEOF(long));
	  MEMZERO(dc_count_ptrs[tbl], 257 * SIZEOF(long));
	}
	/* Create AC table */
	tbl = cinfo->cur_comp_info[i]->ac_tbl_no;
	if (ac_count_ptrs[tbl] == NULL) {
	  ac_count_ptrs[tbl] = (long *) (*cinfo->emethods->alloc_small)
					(257 * SIZEOF(long));
	  MEMZERO(ac_count_ptrs[tbl], 257 * SIZEOF(long));
	}
  }

  /* Initialize DC predictions to 0 */
  for (i = 0; i < cinfo->comps_in_scan; i++) {
	cinfo->last_dc_val[i] = 0;
  }
  /* Initialize restart stuff */
  cinfo->restarts_to_go = cinfo->restart_interval;

  /* Scan the MCU data, count symbol uses */
  (*source_method) (cinfo, htest_encode);

  /* Now generate optimal Huffman tables */
  for (tbl = 0; tbl < NUM_HUFF_TBLS; tbl++) {
	if (dc_count_ptrs[tbl] != NULL) {
	  htblptr = & cinfo->dc_huff_tbl_ptrs[tbl];
	  if (*htblptr == NULL)
	*htblptr = (HUFF_TBL *) (*cinfo->emethods->alloc_small) (SIZEOF(HUFF_TBL));
	  /* Set sent_table FALSE so updated table will be written to JPEG file. */
	  (*htblptr)->sent_table = FALSE;
	  /* Compute the optimal Huffman encoding */
	  gen_huff_coding(cinfo, *htblptr, dc_count_ptrs[tbl]);
	  /* Release the count table */
	  (*cinfo->emethods->free_small) ((void *) dc_count_ptrs[tbl]);
	}
	if (ac_count_ptrs[tbl] != NULL) {
	  htblptr = & cinfo->ac_huff_tbl_ptrs[tbl];
	  if (*htblptr == NULL)
	*htblptr = (HUFF_TBL *) (*cinfo->emethods->alloc_small) (SIZEOF(HUFF_TBL));
	  /* Set sent_table FALSE so updated table will be written to JPEG file. */
	  (*htblptr)->sent_table = FALSE;
	  /* Compute the optimal Huffman encoding */
	  gen_huff_coding(cinfo, *htblptr, ac_count_ptrs[tbl]);
	  /* Release the count table */
	  (*cinfo->emethods->free_small) ((void *) ac_count_ptrs[tbl]);
	}
  }
}


#endif /* ENTROPY_OPT_SUPPORTED */


/*
 * The method selection routine for Huffman entropy encoding.
 */

GLOBAL void
jselchuffman (compress_info_ptr cinfo)
{
  if (! cinfo->arith_code) {
	cinfo->methods->entropy_encode_init = huff_init;
	cinfo->methods->entropy_encode = huff_encode;
	cinfo->methods->entropy_encode_term = huff_term;
#ifdef ENTROPY_OPT_SUPPORTED
	cinfo->methods->entropy_optimize = huff_optimize;
	/* The standard Huffman tables are only valid for 8-bit data precision.
	 * If the precision is higher, force optimization on so that usable
	 * tables will be computed.  This test can be removed if default tables
	 * are supplied that are valid for the desired precision.
	 */
	if (cinfo->data_precision > 8)
	  cinfo->optimize_coding = TRUE;
	if (cinfo->optimize_coding)
	  cinfo->total_passes++;    /* one pass needed for entropy optimization */
#endif
  }
}