/*
 * Encoder.java
 * Encoder with ME, and MC.
 * The Encoder acts as a consumer.  It takes frames produced
 * by AviProduceFrames and uses TSS to search for best-matched
 * block from the reconstructed frame. It encodes the motion vectors
 * and 3d run-level tuples using precalculated Huffman codes. It then 
 * saves the encoded data in a file.
 * Need classes in ../11/, ../12/ and some other chapters.
 * Run: java Vcodec AVI_input_filename
 *      ( Vcodec class is in ../11/ )
 */

import java.io.*;
import java.util.*;

class Encoder extends Thread {
  private CircularQueue buf;         //shared buffer, class defined in ../11/
  private BitOutputStream bitout;    //class defined in ../8/
  private Hcodec hcodec;             //in ../8/
  private MvCodec mvCodec;           //for encoding MVs
  private Display display;           //in ../11/
  private RefVectorMV refVectorMV;   //extends RefVector of ../12/
  private Tss tss;                   //Three step search class 

  //constructor
  public Encoder ( CircularQueue q, BitOutputStream outs, Display d )
  {
    buf = q;
    bitout = outs;
    hcodec = new Hcodec();
    hcodec.build_htable();
    hcodec.print_htable();
    mvCodec = new MvCodec();
    display = d;
    refVectorMV = new RefVectorMV ( buf );
    tss = new Tss( buf.width, buf.height );
  }

  //save one DCT block in an array
  private void put_one_dct_block( int Y[][], short dctcoefs[][][], int bn )
  {
    for ( int i = 0; i < 8; ++i )
      for ( int j = 0; j < 8; ++j )
        dctcoefs[bn][i][j] = ( short ) Y[i][j];  //convert to short  ( 2 bytes );
  }
 
  //make DCT coeffs for one  macroblock.
  private void get_dctcoefs ( YCbCr_MACRO ycbcr_macro, short dctcoefs[][][] ) 
  {
    int b, i, j, k, r;
    int bn = 0;
    int [][] X = new int[8][8], Y = new int[8][8];
    DctVideo dct_idct = new DctVideo();

    //save DCT of Y
    for ( b = 0; b < 4; b++){   //Y has four 8x8 sample blocks
       if ( b < 2 )
        r =  8 * b;            //points to beginning of block
      else
        r = 128 + 8*(b-2);    //points to beginning of block
      k = 0;
      for ( i = 0; i < 8; i++ ) {         //one sample-block
        if ( i > 0 ) r += 16;            //multiply i by 16 ( length of one row )
        for ( j = 0; j < 8; j++ ) {
          X[i][j] = ycbcr_macro.Y[r+j];
        }
      }
           
      dct_idct.dct (  X, Y );              //DCT tranform of 8x8 block 
      put_one_dct_block (Y, dctcoefs, bn++); //save DCT coefs for Y block in array
    }

    //save one 8x8 Cb block
    k = 0;
    for ( i = 0; i < 8; ++i ) {
      for ( j = 0; j < 8; ++j ) {
        X[i][j] = ycbcr_macro.Cb[k];
        k++;
      }
    }
    dct_idct.dct (  X, Y );           //DCT of Cb 8x8 sample block
    put_one_dct_block( Y, dctcoefs, bn++ );

    //save one 8x8 Cr block
    k = 0;
    for ( i = 0; i < 8; ++i ) {
      for ( j = 0; j < 8; ++j ) {
        X[i][j] = ycbcr_macro.Cr[k];
        k++;
      }
    }
    dct_idct.dct ( X, Y );                        //DCT of Cr 8x8 sample block
    put_one_dct_block ( Y, dctcoefs, bn );
  }

  public void print_dctcoefs (short dctcoefs[][][] )  
  {
    for ( int k = 0; k < 6; k++ ){
      System.out.printf("Block %d:\n", k);
      for ( int i = 0; i < 8; i++ ){
        for ( int j = 0; j < 8; j++ ){
          System.out.printf("%4d, ", dctcoefs[k][i][j] );
        }
        System.out.println();
      }
    } 
  }

  //Encode One Frame
  private void encode_one_frame ()
  {
    //encode the frame at the head of buf
    int h = buf.getHead();
    int row, col, i, j, k, r;

    RGB_MACRO rgb_macro = new RGB_MACRO();
    YCbCr_MACRO ycbcr_macro = new YCbCr_MACRO(); //macroblock for YCbCr samples
    YCbCr_MACRO ycc_diff = new YCbCr_MACRO();    //difference between current and reference YCbCr macroblocks

    RgbYcc rgbycc = new RgbYcc ();
    short dctcoefs[][][] = new short[6][8][8];
    short Yr[][] = new short[8][8];
    short Y[][] = new short[8][8];
    int nm = 0;                            //for indexing macroblock
    Point2 cp = new Point2();              //current point
    Point2 ref_point = new Point2();       //reference point
    Vec2 [] mvs = new Vec2[3];    

    Quantizer quantizer = new Quantizer();
    Reorder reorder = new Reorder();
    Run3D runs[] = new Run3D[64];
    for ( i = 0; i < 64; i++ )
      runs[i] = new Run3D();
    Run run = new Run();
    for ( i = 0; i < 3; i++ )
      mvs[i] = new Vec2(); 

    for ( row = 0; row < buf.height; row += 16 ){  //scan all rows of image
      for ( col = 0; col < buf.width; col += 16 ){ //scan all columns of image
        k = (row * buf.width + col) * 3;   //x3 for RGB
        cp.set ( col, row );               //set x, y values of current position
        r = 0;
        for ( i = 0; i < 16; i++ ) {
          for ( j = 0; j < 16; j++ ) {
            rgb_macro.rgb[r].B = 0x000000ff & (int) buf.buffer[h][k];
            rgb_macro.rgb[r].G = 0x000000ff & (int) buf.buffer[h][k+1];
            rgb_macro.rgb[r].R = 0x000000ff & (int) buf.buffer[h][k+2];
            ++r;
            k += 3;
          }
          k += ( buf.width - 16 )*3;       //points to next row within macroblock
        }
        rgbycc.macroblock2ycbcr( rgb_macro, ycbcr_macro );//convert from RGB to YCbCr
        //motion estimation and motion compensation 
        tss.search ( refVectorMV.ref_frame, cp, ycbcr_macro, mvs );
        mvCodec.mv_encode ( mvs, bitout ); 
        //obtain the point at which the reference block is subtracted
        refVectorMV.get_ref_point ( cp, mvs, ref_point );
        //obtain difference between current and reference; save the reference YCC macro in  ycc_refv[nm]
        refVectorMV.current_ref_diff ( ycbcr_macro, ref_point, ycc_diff, nm );
        get_dctcoefs ( ycc_diff, dctcoefs );
        for ( int bn = 0; bn < 6; bn++ ) {
          quantizer.quantize_block ( dctcoefs[bn] );  //quantize one dct sample block
          reorder.reorder ( dctcoefs[bn], Yr );       //reorder the quantized sample block
          run.run_block ( Yr, runs );                 //encode reordered DCT coefs using 3D run-level code
          /*
            Encode the 3D runs with precalculated Huffman codes using the provided
            Huffman table htable. Save the encoded bit stream in file pointed by
            bitout.
          */
          hcodec.huff_encode ( runs, bitout );        //encode and save
        } //for bn
        //reconstruct YCbCr macroblock from quantized DCT coefficients; save block in ycc_refv   
        refVectorMV.reconstruct_macro ( dctcoefs, nm );
        nm++;                                         //next macroblock
      } //for col
    } //for row
    //reconstruct frame from reconstructed macros
    refVectorMV.reconstruct_frame ( nm );
  }

  public void run()
  {
     System.out.printf("Running Encoder with ME, MC ...\n");
     while ( !buf.quit || !buf.bufferEmpty() ) {
       buf.waitIfBufferEmpty();
       encode_one_frame();
       display.display_frame();
       buf.headInc();
       try {
         Thread.sleep ( 10 );   //See ThreadedPlayer.java of ../10/ for playing at desired fps
       } catch (InterruptedException e )  {}
     }
  } 
}
