/* 
 * Test_dct_ppm.java
 * Program to test integer implementations of DCT, IDCT, and RGB-YCbCr 
 * conversions using macroblocks. PPM files are used for testing. 
 * It reads from the file specified by args[0] the RGB data, converts 
 * them to 4:2:0 YCbCr macroblocks, and then to 16-bit DCT coefficients
 * which will be saved in the file specified by args[1].
 * The program then reads back the DCT coefficients from file args[1], 
 * performs IDCT, converts them to YCbCr macroblocks and then to RGB data.
 * The recovered RGB data are saved in the file speicified by args[2].  
 * PPM files can be viewed using "xview".
 *
 * Compile:  javac Test_dct_ppm.java
 * Execute:  e.g. java Test_dct_ppm ../data/beach.ppm t.dct t.ppm
 * Need to set classpath by: "export CLASSPATH=$CLASSPATH:../5/"
 */

import java.io.*;
import java.awt.Frame;
import java.awt.image.*;
import javax.media.jai.JAI;
import javax.media.jai.RenderedOp;
import com.sun.media.jai.codec.FileSeekableStream;
import javax.media.jai.widget.ScrollingImagePanel;
import com.sun.media.jai.codec.PNMEncodeParam;

public class Test_dct_ppm {
   
  //DCT file header 
  private static byte []  header = { 'D', 'C', 'T', '4', ':', '2', ':', '0' };

  public static void write_dct_header( int width, int height, DataOutputStream out )
  {

    try {
       out.write ( header );
       out.writeInt ( width );
       out.writeInt ( height );
    } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
    }
  }

  public static int read_dct_header(RGBImage rgbimage,  DataInputStream in)
  {
    byte [] bytes = new byte[header.length];
    try {
      in.read ( bytes );
      rgbimage.width  = in.readInt ();
      rgbimage.height = in.readInt ();
    } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
    }
    for ( int i = 0; i < header.length; ++i )
      if ( bytes[i] != header[i] )
        return -1;      //wrong header

    return 1;
  }

  //save PPM header and RGB data
  public static int save_ppm ( RGBImage rgbimage, DataOutputStream out )
  {
    byte [] P6 = { 'P', '6', '\n' };
    byte [] colorLevels = { '2', '5', '5', '\n' };
    String sw = Integer.toString ( rgbimage.width ) + " ";
    String sh = Integer.toString ( rgbimage.height ) + "\n";
    
    int isize = rgbimage.width * rgbimage.height;
    byte [] bytes = new byte[3*isize];
    for ( int i = 0, k = 0; i < isize; i++, k+=3 ){
        bytes[k] = (byte) (rgbimage.ibuf[i].R);
        bytes[k+1] = (byte) (rgbimage.ibuf[i].G);
        bytes[k+2] = (byte) (rgbimage.ibuf[i].B);
    }

    try {
      out.write ( P6 );
      out.writeBytes ( sw );
      out.writeBytes ( sh );
      out.write ( colorLevels );
      out.write ( bytes );
    } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
    }

    return 1;
  }

  public static void main(String[] args) throws InterruptedException
  {
    if (args.length < 3) {
      System.out.println("Usage: java " + "Test_dct_ppm" +
        " input_ppm_filename output_dct_filename recoverd_ppm_filename\n" +
        "e.g. java Test_dct_ppm ../data/beach.ppm t.dct t.ppm");
      System.exit(-1);
    }

    /*
     * Create an input stream from the specified file name
     * to be used with the file decoding operator.
     */

     FileSeekableStream stream = null;
     try {
       stream = new FileSeekableStream(args[0]);
     } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
     }
    /* Create an operator to decode the image file. */
    RenderedOp image = JAI.create("stream", stream);

    /* Get the width and height of image. */
    int width = image.getWidth();
    int height = image.getHeight();

    if ( width % 16 != 0 || height % 16 != 0 ) {
      System.out.println("Program only works for image dimensions");
      System.out.println("divisible by 16. Use 'convert' to change");
      System.out.println("image dimensions.");
      System.exit(1);
    }
    int [] samples = new int[3*width*height];

    Raster ras = image.getData();
    //save pixel RGB data in samples[]
    ras.getPixels( 0, 0, width, height, samples );
    RGBImage rgbimage = new RGBImage( width, height );
    //copy image data to RGBImage object buffer
    int isize = width * height;
    for ( int i = 0, k = 0; i < isize; ++i, k+=3 ) {
      rgbimage.ibuf[i].R =  samples[k];
      rgbimage.ibuf[i].G = samples[k+1];
      rgbimage.ibuf[i].B = samples[k+2];
    }
    //Convert data to DCT coefficients and save them in file args[1]
    try {
      File f = new File ( args[1] );
      OutputStream o = new FileOutputStream( f );
      DataOutputStream out = new DataOutputStream ( o );
      write_dct_header ( width, height, out );
      Encode enc = new Encode ();
      enc.encode_dct ( rgbimage, out );
      out.close();
    } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
    }

    System.out.printf("\nEncoding done, DCT coeff saved in %s\n",args[1]);
  
    //read the DCT data back from args[1] and convert to RGB
    DataInputStream in;
    Decode dct_decoder = new Decode ();
    try {
      File f = new File ( args[1] );
      InputStream ins = new FileInputStream( f );
      in = new DataInputStream ( ins );
      if ( read_dct_header ( rgbimage,  in ) == -1 ){
        System.out.println("Not DCT File");
        return;
      }
      //To be fair in the demo, we use a new image buffer.
      //Java collects any memory garbage.
      //Apply IDCT
      rgbimage = new RGBImage ( rgbimage.width, rgbimage.height );
      dct_decoder.decode_dct (rgbimage, in);
      in.close();
    } catch (IOException e) {
       e.printStackTrace();
       System.exit(0);
    }

    //Save recovered RGB data in file args[2] in PPM format.
    try {
      File f = new File ( args[2] );
      OutputStream o = new FileOutputStream( f );
      DataOutputStream out = new DataOutputStream ( o );
      save_ppm ( rgbimage, out ); //save PPM header and RGB data
      out.close();
    } catch (IOException e) {
      e.printStackTrace();
      System.exit(0);
    }
    System.out.printf("Decoded data saved in %s \n", args[2] );
  }
}
