#include <iostream>
#include <sstream>
#include <fstream>

#include <jpeg-c++.h>

#define _DEBUG

unsigned char* read_image(char *image_file, int *width, int *height, int *depth);
void write_image(char *image_file, unsigned char *image_buffer, int width, int height, int depth);
unsigned char * grayscale_conversion(unsigned char* color_image_buffer, int width, int height, int depth);

int main(int argc, char **argv)
{
  int                width, height, depth;
  unsigned char*     color_image_buffer;
  unsigned char*     greyscale_image_buffer;
  
  if (argc != 3) {
    std::cerr << "Usage : convert_image input_image_file grayscale_image_file\n";
    exit(EXIT_FAILURE);
  }

  color_image_buffer = read_image(argv[1], &width, &height, &depth);
  greyscale_image_buffer = grayscale_conversion(color_image_buffer, width, height, depth);
  write_image(argv[2], greyscale_image_buffer, width, height, 1);
}

/*************************************************************************************/
/* Read an image from a jpeg file                                                    */
/*                                                                                   */
/* Input parameters :                                                                */
/*    image_file  : the name of the jpeg file                                        */
/*                                                                                   */
/* Output paramters :                                                                */
/*   width        : image width                                                      */
/*   height       : image height                                                     */
/*   depth        : image depth (3 for color image,1 for greyscale image             */
/*                                                                                   */
/* Return value function :                                                           */
/*   image_buffer : pixel values of the image                                        */
/*************************************************************************************/
unsigned char* read_image(char *image_file, int *width, int *height, int *depth) 
{  
  unsigned char*     image_buffer;
  std::ostringstream jpeg_in_memory;
  jpeg::Decompress   jpg_to_bytes(jpeg_in_memory);
  std::ifstream      input_image;    
  int                size;
 
  // Open the input stream
  input_image.open(image_file, std::ifstream::binary);
  if(!input_image) {
    std::cerr << "Cannot read" << "\"" << image_file << "\"" << "file. Aborting." << std::endl;
    exit(EXIT_FAILURE);
  }
  jpg_to_bytes.setInputStream(input_image); 
  
  // Read the image 
  jpg_to_bytes.readHeader(*width, *height, *depth);
#ifdef _DEBUG
  std::cout << "Image " << image_file << " : " << "width = " << *width << " height = " << *height << " depth = " << *depth << std::endl; 
#endif
  size = *width * *height * *depth;
  image_buffer = new unsigned char[size];
  // Now, image is decompressed in rgb buffer (rgbrgbrgbrgb... or graygraygray... line by line).
  jpg_to_bytes.readImage(image_buffer); 

  //Close the input stream
  input_image.close();

  return image_buffer;
}

/*************************************************************************************/
/* Write the image in a jpeg file                                                    */
/*                                                                                   */
/* Input parameters :                                                                */
/*    image_file   : the name of the jpeg file                                       */
/*    image_buffer : pixel values of the image                                       */
/*    width        : image width                                                     */
/*    height       : image height                                                    */
/*    depth        : image depth                                                     */
/*************************************************************************************/
void write_image(char *image_file, unsigned char *image_buffer,  int width, int height, int depth) 
{
  jpeg::Compress     bytes_to_jpg;
  std::ofstream      output_image;

  // Open an output stream
  output_image.open(image_file, std::ofstream::binary);
  if(!output_image) {
    std::cerr << "Cannot read" << "\"" << image_file << "\"" << "file. Aborting." << std::endl;
    exit(EXIT_FAILURE);;
  }
  bytes_to_jpg.setOutputStream(output_image);
  // Write the image in the jpeg file
  bytes_to_jpg.writeImage(width, height, depth, image_buffer);

  // Close the ouput stream
  output_image.close();
}

/*************************************************************************************/
/* Convert the image in grayscale                                                    */
/*                                                                                   */
/* Input parameters :                                                                */
/*    image_buffer : pixels of the image                                             */
/*    width        : image width                                                     */
/*    height       : image height                                                    */
/*    depth        : image depth                                                     */
/*                                                                                   */
/* Return value function :                                                           */
/*   grayscale_image : a grayscale image                                             */
/*************************************************************************************/
unsigned char * grayscale_conversion(unsigned char* image_buffer, int width, int height, int depth)
{
  unsigned char *gray_scale_image;
  int k = 0;
  int size = width * height * depth;

  gray_scale_image = new unsigned char[width + height];
  for (int i = 0; i < size; i += depth) {
    gray_scale_image[k] = (unsigned char) (0.2 * image_buffer[i] +
					   0.7 * image_buffer[i + 1] + 
					   0.1 * image_buffer[i + 2]);
    k++;
  }
  return gray_scale_image;
}