/*
  Encode.java
  Contains functions that convert an RGB frame to YCbCr, then to DCT 
  coefficients which are saved in a file.  The program reads the data 
  back from the file, convert them back to YCbCr and to RGB, which
  will be saved in another file.
*/

import java.io.*;

class Encode {

  //save one YCbCr macroblock.
  public void save_yccblocks( YCbCr_MACRO ycbcr_macro, DataOutputStream out )
  {
    int b, i, j, k, r;
  
    //one macroblock has 256 Y samples
    byte [] Yblock = new byte[256];

     r = 0;
    //save four 8x8 Y sample blocks
    for ( b = 0; b < 4; b++ ) {
      if ( b < 2 )
        k = 8 * b;                //points to beginning of block
      else
        k = 128 + 8 * ( b - 2 );  //points to beginning of block
      for ( i = 0; i < 8; i++ ) { //one sample-block
        if ( i > 0 ) k += 16;     //advance k by 16 ( length of one row )
        for ( j = 0; j < 8; j++ ) {
          Yblock[r] = ( byte ) ycbcr_macro.Y[k+j]; 
          r++;
        } 
      }
    }

    //save one 8x8 Cb block
    byte [] Cb_block = new byte[64]; 
    k = 0;
    for ( i = 0; i < 8; ++i ) {
      for ( j = 0; j < 8; ++j ) {
        Cb_block[k] = ( byte ) ycbcr_macro.Cb[k];
        k++;
      }
    }

    //save one 8x8 Cr block
    byte [] Cr_block = new byte[64]; 
    k = 0; 
    for ( i = 0; i < 8; ++i ) {
      for ( j = 0; j < 8; ++j ) {
        Cr_block[k] = ( byte ) ycbcr_macro.Cr[k];
        k++;
      }
    }

    //save four Y sample blocks, one Cb block, one Cr block in file
    try {
      out.write ( Yblock );
      out.write ( Cb_block );
      out.write ( Cr_block );
    } catch (IOException e) {
      e.printStackTrace();
      System.exit(0);
    }
  }

  /*
    Convert one frome of RGB to YCbCr and save the converted data.
  */
  public void encode ( RGBImage image, DataOutputStream out )
  {
    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
    RgbYcc rgbycc = new RgbYcc ();
    for ( row = 0; row < image.height; row += 16 ) {
      for ( col = 0; col < image.width; col += 16 ) {
        k = row * image.width + col;
        r = 0;				
        for ( i = 0; i < 16; ++i ) {
	  for ( j = 0; j < 16; ++j ) 
	    rgb_macro.rgb[r++] = image.ibuf[k++];
	  k += ( image.width - 16 );  //points to next row within macroblock
        }
        rgbycc.macroblock2ycbcr( rgb_macro, ycbcr_macro );//convert from RGB to YCbCr
        save_yccblocks( ycbcr_macro, out );		  //save one YCbCr macroblock
      } //for col
    } //for row
  }

  static void print_block( int [][] X )
  {
    for ( int i = 0; i < 8; ++i ){
      System.out.printf("\n");
      for ( int j = 0; j < 8; ++j ) {
        System.out.printf("%4d, ", X[i][j] );
      }
    }
    System.out.printf("\n-------------------------------------");
  }

  //The following functions are added in Chapter 6
  void save_one_dctblock ( int [][] Y, DataOutputStream out )
  {
    for ( int i = 0; i < 8; ++i )
      for ( int j = 0; j < 8; ++j ) 
        try {  
          out.writeShort( Y[i][j] );   //save DCT coefficients of the block
       } catch (IOException e) {
          e.printStackTrace();
          System.exit(0);
       }
  }

  /*
   *  Apply DCT to the six 8x8 sample blocks of a 4:2:0 YCbCr macroblock
   *  and save the coefficients in a file pointed by out  
   */
  void save_dct_yccblocks( YCbCr_MACRO ycbcr_macro, DataOutputStream out )
  {
    int b, i, j, k, r;
    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
      save_one_dctblock(Y, out);  //save DCT coefficients of the 8x8 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
    save_one_dctblock( Y, out );
    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
    save_one_dctblock ( Y, out );
  }


  /*
   * Convert RGB  data to YCbCr, then to DCT coeffs and save DCT coeffs 
   */
  void encode_dct ( RGBImage image, DataOutputStream out )
  {
    int row, col, i, j, k, r;
    
    RGB_MACRO rgbmacro = new RGB_MACRO();
    YCbCr_MACRO ycbcr_macro = new YCbCr_MACRO(); //macroblock for YCbCr samples
    RgbYcc rgbycc = new RgbYcc();

    for ( row = 0; row < image.height; row += 16 ) {
      for ( col = 0; col < image.width; col += 16 ) {
        k = row * image.width + col;   //points to beginning of macroblock
        r = 0;                   
        for ( i = 0; i < 16; ++i ) {
          for ( j = 0; j < 16; ++j ) {
            rgbmacro.rgb[r++] = image.ibuf[k++];
          }
          k += image.width - 16;       //points to next row within macroblock 
        }
        rgbycc.macroblock2ycbcr ( rgbmacro, ycbcr_macro );  
        save_dct_yccblocks( ycbcr_macro, out );
      } //for col
    } //for row
  }
}
