/*
 * Hcodec.java
 * Building a Huffman Encoder-Decoder by making use of a TreeMap,
 * which collects all pre-calculated run-level and Huffman codewords.
*/

import java.util.Set;
import java.util.Map;
import java.util.TreeMap;
import java.util.Iterator;
import java.io.*;

class Hcodec
{
  private final int NSymbols = 256;   //maximum symbols allowed
  private final short ESC = 127;  //Escape run
  private final short ESC_HUF_CODE = 0x60;  //Escape code
  private final short EOS = 126;      //End of Stream 'symbol' ( run )
  private final short EOS_HUF_CODE = 0x0a;      //End of Stream 'symbol'
  private boolean tableNotBuilt = true;
  TreeMap<RunHuff, RunHuff>  htable = new TreeMap<RunHuff,RunHuff>();
  
  //use a map ( htable ) to collect all pre-calculated run-level 
  //  and Huffman codewords
  void build_htable ()
  {
    //N = number of pre-calculated codewords with positive levels
    //  In practice, N should be larger than 100
    short i, j, k, N = 11;    
    //lengths of Huffman codewords (not including sign-bit)
    byte hlen[] = { 2, 3, 4, 4, 4, 5, 5, 5, 6, 6,  7 };       
    //Huffman codewords, 0x60 for ESC, 0x3a for EOS; a codeword should NEVER be a prefix of another
    short hcode[] = {0x01,0x3,0x7,0xf,0xe, 0x16, 0x6, 0x1a, 0x2a, EOS_HUF_CODE, ESC_HUF_CODE};

    //data of 3D run-level tuples ( codewords )
    byte runs[] = {0, 1, 2, 0, 0, 3, 4, 5, 0, EOS, ESC};  //126 signifies end of stream, 127 signifies Escape
    short levels[] = {1, 1, 1, 2, 1, 1, 1, 1, 3, 1, 1};   //Note: don't set Escape level to 0
    byte  lasts[] = {0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0 };   //  otherwise => duplicate key          

    Run3D r = new Run3D();             //a 3D run-level codeword ( tuple )
    RunHuff rf[] = new RunHuff[128];   //table containing RunHuff objects

    //inserting RunHuff objects into rf[] 
    k = 0; j = 0;
    for ( i = 0; i < N; i++ ) {
      r.run = runs[i];
      r.level = levels[i];
      r.last = lasts[i];
      //construct a RunHuff object, positive level, so sign=0
      rf[k++] = new  RunHuff ( r, hcode[i] << 1,  hlen[i], j++ );
      //do the same thing for negative level, sign = 1 
      r.level = (short) -r.level;                                       
      rf[k++] = new RunHuff ( r, (short)((hcode[i]<<1) | 1), hlen[i], j++ );
    }
    //insert all (positive & negative levels) RunHuff objects into htable
    k = ( short ) ( 2 * N );
    for ( i = 0; i < k; i++ )
      htable.put( rf[i], rf[i] );
    tableNotBuilt = false;
  }

  void print_htable()
  {
    // Use an Iterator to traverse the mappings in the TreeMap.  Note
    // that the mappings are in sorted order (with respect to the keys). 
    Iterator itr = htable.entrySet().iterator();
    RunHuff rfs[] = new RunHuff[htable.size()];
    System.out.printf("table size=%d,", htable.size() );
     System.out.printf("\n(run, level, last), \tCodeword\tHuff Code Length, index");
     while ( itr.hasNext() )
     {
        Map.Entry entry = (Map.Entry) itr.next();
        RunHuff rhuf = (RunHuff) entry.getValue();
	   rfs[rhuf.index] = rhuf;
     }
     for ( int i = 0; i < htable.size(); i++ ) {
        RunHuff rhuf = rfs[i];
	System.out.printf("\n(%4d, %4d, %d), \t%8x \t%x\t    %d, \t%4d", rhuf.r.run, rhuf.r.level, rhuf.r.last, rhuf.codeword, rhuf.codeword >> 1, rhuf.hlen, rhuf.index );
     }
     System.out.printf("\n");
  }

  //Encode codeword unsing ESC 
  void escape_encode ( BitOutputStream outputs, Run3D r )
  {
    try {
      if ( r.level < 0 ) {           //value of level negative
        outputs.writeBit ( 1 );      //output sign-bit first
        r.level = (short) -r.level;  //change level value to positive 
      } else
        outputs.writeBit ( 0 );      //value of level positive
      outputs.writeBits ( ESC_HUF_CODE, 7 );  //ESCAPE code
      if ( r.run == 64 ) r.run = 63; //r.level differentiates between
                                     //  if last element nonzero 
      outputs.writeBits ( (int) r.run, 6 );      //6 bits for run value
      outputs.writeBits ( (int) r.level, 8 );    //8 bits for level value
      outputs.writeBit ( (int) r.last );         //1 bit for last value
    } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
    }
  }

  /*
    Inputs: 
      runs[] contains the 3D run-level tuples of a macroblock of quantized
      DCT coefficients
    Outputs:
      bitstreams of codewords ( Huffman + sign or ESCAPE + 3D run-level ) 
      to BitOutputStream outputs
    Note that data member htable contains all the pre-calculated Huffman 
      codewords of 3D run-level tuples
  */
  void huff_encode ( Run3D runs[], BitOutputStream outputs )
  {
    short i, j, k;
    if ( tableNotBuilt )
      build_htable();
    k = 0;  i = 0;
    while ( i < 64 ) {          //a macroblock has at most 64 samples
      try {
        //construct a RunHuff object; only runs[k] is relevant as it is
        //  used for searching (we've defined CompareTo in RunHuff class)
        RunHuff rf = new RunHuff( runs[k], 0, (byte) 0, (short) 0 ); 
        if ( htable.containsKey ( rf ) ){
          RunHuff rhuf = htable.get( rf );
          outputs.writeBits( (int) rhuf.codeword, rhuf.hlen + 1);
        } else {                          //not in table
          escape_encode( outputs, rf.r ); //need to 'escape encode' Run3D Object
       }
      } catch (IOException e) {
        e.printStackTrace();
        System.exit(0);
      }
      if ( runs[k].last != 0 ) break;   //end of run-level codewords
      i += ( runs[k].run + 1 );         //special case: whole run-block 0
      k++;
    }
  }

  //Encode End of Stream symbol
  void huff_encode_EOS ( BitOutputStream outputs )
  {
    Run3D r = new Run3D();         //a 3D run-level codeword ( tuple )
    r.last = 0;
    r.level = 1;
    r.run = EOS;
    try {
        //construct a RunHuff object; only run3D r is relevant as it is
        //  used for searching (we've defined CompareTo in RunHuff class)
        RunHuff rf = new RunHuff( r, 0, (byte) 0, (short) 0 ); 
        if ( htable.containsKey ( rf ) ){
          RunHuff rhuf = htable.get( rf );
          outputs.writeBits( (int) rhuf.codeword, rhuf.hlen + 1);
        } else {                   //not in table ( should not happen )
          System.out.printf("\nError: EOS not in Huffman table!\n");
       }
    } catch (IOException e) {
        e.printStackTrace();
        System.exit(0);
    }
    return; 
  }

  //Build  Huffman Tree using pre-calculated codewords.
  //  Return tree Dtable d.
  void build_huff_tree ( Dtables d )
  {
    int i, j, n0, free_slot, loc, loc0, root, ntotal;
    int mask, hcode;
  
    n0 = NSymbols;        //number of symbols (=number of leaves in tree)
    ntotal = 2 * n0 - 1;  //Huffman tree has  n0 - 1  internal nodes 
    root = 3 * n0 - 3;    //location of root, offset n0 has been added
    free_slot = root - 2; //next free table entry for filling in with a 
                          //  pointer or an index
                          //  (note:root has root_left, root_right)
    for ( i = 0; i < ntotal; ++i )   //initialize the table
      d.huf_tree[i] = -1;            //all entries empty

    //Java Map does not have direct iterator, so need entrySet() which returns a set.
    //  Java Set supports iterator.
    Iterator itr = htable.entrySet().iterator();
    RunHuff rfs[] = new RunHuff[htable.size()];
    while ( itr.hasNext() ) {  //user iterator to traverse Map ( in set veiw ).
        Map.Entry entry = (Map.Entry) itr.next();
        RunHuff rhuf = (RunHuff) entry.getValue();
	rfs[rhuf.index] = rhuf;  //put codewords in rfs[]
    }

    //use codewords to build tree and save it in d.huf_tree[] 
    for ( int ii = 0; ii < htable.size(); ii++ ) {
      RunHuff rhuf = rfs[ii];
      if ( rhuf.r.level < 0 ) continue;  //only save positive levels 
                                         //  of run-level codeword
      d.run_table[rhuf.index/2]=rhuf.r;  //save run-level codeword;index 
                                         //  divided by 2 as only postive levels saved
        loc = root;                      //always start from root
        mask = 0x01;                     //for examining bits of Huffman codeword
        hcode = rhuf.codeword >> 1;      //rightmost bit is sign-bit,not Huffman
        for (i = 0; i < rhuf.hlen; i++){ //traverse the Huffman codeword
          loc0 = loc - n0;               //everything shifted by offset n0
          if ( i == ( rhuf.hlen - 1 ) ){ //last bit, should point to leaf
            if ( (mask & hcode) == 0 )   //a 0, save it at 'left' leaf
              d.huf_tree[loc0] = (short) (rhuf.index/2);
            else                         //a 1, save it at 'right' leaf
              d.huf_tree[loc0-1] = (short) (rhuf.index/2);
            continue;  //get out of for-i for loop, consider next codeword 
          }
          if ( (mask & hcode) == 0 ){     //a 0 ( go left )
            if (d.huf_tree[loc0] == -1){ //slot empty
              d.huf_tree[loc0] = (short)free_slot;//point to left new child
              free_slot -= 2;            //next free table entry
            }                            //else : already has left child
            loc = d.huf_tree[loc0];      //follow the left child
          } else {                       //a 1 ( go right )
            if (d.huf_tree[loc0-1]== -1){//slot empty
              d.huf_tree[loc0-1]= (short)free_slot;//point to right new kid
              free_slot -= 2;
            }                            //else: already has right child
            loc = d.huf_tree[loc0-1];    //follow the right child
          }
          mask <<= 1;                    //consider next bit
        } //for i
      } //while
  } //build_huff_tree()

  /*
    Inputs: inputs, the encoded bitstream to be decoded
          d.huf_tree[], table containing the Huffman tree
          d.run_table[], table containing the actual run-level codewords
    Output: runs[], table containing the run-level tuples of a macroblock
  */
  short huff_decode( BitInputStream inputs,  Dtables d,  Run3D runs[] )
  {
    short n0, loc, loc0, root, k;
    int c, sign;
    boolean done = false;
    Run3D rp;
    n0 = NSymbols;               //number of symbols
    root = (short)(3 * n0 - 3);  //points to root of tree
    k = 0;
    try {
      while ( !done ) {
        loc = root;              //starts from root
        sign = inputs.readBit(); //sign-bit
        do {
          loc0 = (short) (loc - n0);
          c = inputs.readBit();             //read one bit
          if (c < 0) {done = true; break;}  //no more data, done
          if ( c == 0 )                     //a 0, go left
            loc = d.huf_tree[loc0];
          else                              //a 1, go right
            loc = d.huf_tree[loc0-1];
        } while ( loc >= n0 );              //traverse until reaches leaf
	if ( loc >= n0 ) break;             //done
        rp = d.run_table[loc];
         
 	Run3D r3d = new Run3D();
        if ( rp.run ==  ESC ) {  //ESCAPE code, read actual run-level tuple
          r3d.run  = (byte) inputs.readBits ( 6 );  //read 6 bits for run
          r3d.level  = (short) inputs.readBits ( 8 );//read 8 bits for level
          r3d.last = (byte) inputs.readBit();       //read 1 bit for last
          if ( sign == 1 )            //if sign is 1, level should be negative
            r3d.level = (short) -r3d.level;
        } else if ( rp.run == EOS ) {
          done = true;                //End of Stream, Done!
          break;
        } else {                      //not ESCAPE code 
	  r3d.run = rp.run;
	  r3d.level = rp.level;
	  r3d.last = rp.last;
          if ( sign == 1 )            //1 => negative
            r3d.level = (short) -r3d.level;
        }
        if ((r3d.run == 63)&&(r3d.level == 0))
          r3d.run = 64;               //whole block 0
        runs[k++] = r3d;              //save tuple in table runs[]
        if ( r3d.last != 0 )          //end of macroblock 
          break;
      } //while
    } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
    }

    if ( done ) return -1;            //if ( done ) => no more data
    else return 1;
  }
}



