/*
 *  MvCodec.java
 *  This class is similar to Hcodec of Chapter 8 except that
 *  it processes motion vectors ( mv ) and may be simpler as it
 *  does not need to escape-encode.  On the other hand, it 
 *  has to handle the three levels of mv's.
 */

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

class MvCodec {
  TreeMap<MvHuff, MvHuff> mvHtables[];
  public MVtables mv_decode_tables[] = new MVtables[3];

  //constructor 
  public MvCodec()
  {
    //generic array creation not allowed in java
    mvHtables = (TreeMap<MvHuff, MvHuff>[]) new TreeMap[3];
    for ( int i = 0; i < 3; i++ ) {
      mvHtables[i] = new TreeMap<MvHuff, MvHuff>();
      mv_decode_tables[i] = new MVtables();
    } 
    build_mvHtables();
    for ( int i = 0; i < 3; i++ )   //for decoding only
      build_huff_tree ( i );
  } 

  //use a map ( mvHtable ) to collect all pre-calculated motion vector ( MV ) codewords.
  //  There are totally 9 motion vectors.
  void build_mvHtables ()
  {
    short i, j, k, N = 9;         //N = # of possible motion vectors
    byte hlen[] = { 5, 3, 5, 4, 2, 3, 3, 3, 3 };  //lengths of MV codewords
    //Huffman codewords: no codeword is a prefix of another
    short hcode[] = { 0x1f, 0x02, 0x0f, 0x07, 0x00, 0x03, 0x05, 0x06, 0x01 };

    //Level 0  motion vectors (x, y) 
    int x[] = { -4, -4, -4,  0, 0, 0,  4,  4,  4 };
    int y[] = { -4,  0,  4, -4, 0, 4, -4,  0,  4 };
    Vec2 mv;
    MvHuff mvf[][] = new MvHuff[3][128]; //table containing MvHuff objects

    int scale = 1;                       //for calculating MVs of all levels
    for ( j = 0; j < 3; ++j ) {
      k = 0;
      for ( i = 0; i < N; ++i ) {
        mv = new Vec2();                 //an MV vector
        mv.set( x[i]/scale, y[i]/scale );
        mvf[j][k] = new MvHuff ( mv, hcode[i], hlen[i], i );
        k++;
      }
      scale *= 2;
    }

    //insert all N MvHuff objects into the mvtable map
    for ( k = 0; k < 3; ++k ) {
      for ( i = 0; i < N; ++i )
        mvHtables[k].put ( mvf[k][i], mvf[k][i] );
    }
  }


  /*
   * Inputs: mvtable contains all the pre-calculated Huffman codewords of VEC2 mv tuples
   *         mvs[] contains the 3 motion vectors of TSS search.
   * Outputs:bitstreams of codewords to outputs
   */
  void mv_encode (Vec2 mvs[], BitOutputStream outputs )
  {
    int nlevels = 3;
    for ( int  i = 0; i < nlevels; ++i ) {
      TreeMap<MvHuff, MvHuff> htable = mvHtables[i];
      MvHuff mvf = new MvHuff ( mvs[i], 0,(byte) 0, (short)0 ); //construct an MvHuff object; only mvs[k] is relevant here
      if ( htable.containsKey ( mvf ) ) { 
        MvHuff mvhuf = htable.get ( mvf );
        try {
          outputs.writeBits( mvhuf.codeword, mvhuf.hlen );
        } catch (IOException e) {
          e.printStackTrace();
          System.exit(0);
        }
        if ( mvhuf.mv.x == 0 && mvhuf.mv.y == 0 )
          break;                    //best match at position 0, done
      }
    }
}

  /*   
   * Input:   mvHtables[nlevel], a map containing all pre-calculated Huffman codewords of Vec2 motion vectors at three levels
   * Outputs: mv_decode_table, the Huffman tree containing pointers and indices pointing to 
   * 	    VEC2 motion vectors
   *          d.mv_table[], the table that contains the motion vectors   
   */

  void build_huff_tree ( int nlevel )
  {
    int i, j, n0, free_slot, loc, loc0, root, ntotal;
    int mask, hcode;
    MVtables d; 

    n0 = 9;                //number of symbols ( = number of leaves in tree, hypothetical here )
    ntotal = 2 * n0 - 2;   //Huffman tree has  n0 - 1  internal nodes 
    root = 3 * n0 - 3;     //location of root ( see Chapter 8 ), offset n0 has been added
    free_slot = root - 2;  //next free table entry for filing in with a pointer or an index
    d = mv_decode_tables[nlevel];

                           //  ( note: root has root_left, root_right )
    for ( i = 0; i < ntotal; ++i ) //initialize the table
      d.mv_tree[i] = -1;           //all entries empty
   
    TreeMap<MvHuff, MvHuff> htable = mvHtables[nlevel];
    Iterator itr = htable.entrySet().iterator();
    MvHuff mvfs[] = new MvHuff[htable.size()];
    while ( itr.hasNext() ) {
      Map.Entry entry = (Map.Entry) itr.next();
      MvHuff mvhuf = (MvHuff) entry.getValue();
      mvfs[mvhuf.index] = mvhuf;
    }

    for ( int ii = 0; ii < htable.size(); ii++ ) {
      MvHuff mvhuf = mvfs[ii];
      d.mv_table[mvhuf.index] = mvhuf.mv;
      loc = root;	         //always start from root
      mask = 0x01;               //for examining bits of Huffman codeword
      hcode = mvhuf.codeword;
      for ( i = 0; i < mvhuf.hlen; ++i ) {  //traverse the codeword
        loc0 = loc - n0;                    //everything offset by n0
        if ( i == ( mvhuf.hlen - 1 ) ) {    //last bit; points to leaf
	  if ( ( mask & hcode ) == 0 )	    //a 0, save it in 'left' leaf
	    d.mv_tree[loc0] = mvhuf.index;
	  else				    //a 1, save it in 'right' leaf
	    d.mv_tree[loc0-1] = mvhuf.index;
	  continue;    //get out of i-for-loop; consider next codeword
        }
        if ( ( mask & hcode ) == 0 ) {	    //a 0, go left
	  if ( d.mv_tree[loc0] == -1 ) {    //slot empty
            d.mv_tree[loc0] = (short)free_slot;//point to left new child
	    free_slot -= 2;		    //next free table entry
          }				    //else: already has left child
          loc = d.mv_tree[loc0];            //follow the left child
        } else {                            //a 1, go right
 	  if (d.mv_tree[loc0-1]== -1){ 	    //slot empty
            d.mv_tree[loc0-1]=(short)free_slot;//point to right new child
            free_slot -= 2;
          }                            	    //else: already has right child
          loc = d.mv_tree[loc0-1];     	    //follow the right child
        }
        mask <<= 1;		    	    //consider next bit
      } //for i
    } //for ii
  }  //build_huff_tree()


  /*
   *   Inputs: bf, the encoded bitstream to be decoded
   *           d.mv_tree[], table containing the Huffman tree
   *           d.mv_table[], table containing the actual motion vectors
   *   Output: mvs[], table containing the motion vectos
   *
   **/
  short mv_decode( BitInputStream bf,  Vec2 mvs[] )
  {
    short n0, loc, loc0, root, k;
    int c;
    boolean done = false;
    Vec2 mv;

    n0 = 9;                                 //number of symbols
    root = (short)( 3 * n0 - 3);            //points to root of tree
    k = 0;
    MVtables [] d =  mv_decode_tables;
    try {
      while ( !done ) {
        loc = root;
        do {
          loc0 = (short)(loc - n0);
          c = bf.readBit();                 //read one bit
          if (c < 0) {done = true; break;}  //no more data, done    
          if ( c == 0 )                     //a 0, go left
            loc = d[k].mv_tree[loc0];
          else                              //a 1, go right
            loc = d[k].mv_tree[loc0-1];
        } while ( loc >= n0 );              //traverse until reaches leaf
        mvs[k] = d[k].mv_table[loc];
        if ( d[k].mv_table[loc].x == 0 && d[k].mv_table[loc].y == 0 ) break;
        if ( k >= 2 ) break;
        k++;
      }
    } catch (IOException e) {
      e.printStackTrace();
      System.exit(0);
    }
    if ( done ) return -1;                  //if ( done ) => no more data
    else return 1;
  }
}
