/*
 * WJPGC.C 
 *
 * This is a mex interface to the Independent Jpeg Group's (IJG)
 * LIBJPEG libaray.  This can write RGB and grayscale JPEG images.
 * 
 * The syntaxes are:
 *    
 *    
 *      wjpgc(RGB, filename);
 *      wjpgc(GRAY, filename);
 *      wjpgc(... , quality);
 * 
 * RGB is either a mxnx3 uint8 array containing a 24-bit image to
 * be stored in the jpeg file filename or a mxn uint32 array with 
 * each 4 bytes being RGBA, were A is currently garbage (maybe Alpha
 * in the future).
 *       
 * GRAY is a mxn uint8 array containing a grayscale image to
 * be stored in the jpeg file filename.    
 * 
 * The quality argument specifies the compression scheme's quality
 * setting which adjusts the tradeoff between image quality and
 * file-size. 
 * 
 * 
 * KNOWN BUGS:
 * -----------
 *
 * ENHANCEMENTS UNDER CONSIDERATION:
 * ---------------------------------
 * 
 * 
 * The IJG code is available at:
 * ftp://ftp.uu.net/graphics/jpeg/jpegsrc.v6.tar.gz
 * 
 * Chris Griffin, June 1996
 * Copyright (c) 1984-98 by The MathWorks, Inc.
 * $Revision: 1.6 $  $Date: 1998/02/19 16:29:25 $
 */



#include "mex.h"
#include <stdio.h>
#include <setjmp.h>
#include "jpeglib.h"

#define RGB_IMG          0
#define GRAY_IMG         1
#define PACKED_RGB_IMG   2

static void WriteGRAYFromUint8(j_compress_ptr cinfoPtr, 
                               uint8_T *uint8Data, 
                               const int *size);
static void WriteRGBFromUint8(j_compress_ptr cinfoPtr, 
                              uint8_T *uint8Data, 
                              const int *size);
static void WriteRGBFromUint32(j_compress_ptr cinfoPtr, 
                              uint32_T *uint8Data, 
                              const int *size);

/* 
 *  These guys are to replace specific routines in the jpeg library's
 *  jerror.c module. 
 */
static void my_error_exit (j_common_ptr cinfo);
static void my_output_message (j_common_ptr cinfo);

struct my_error_mgr {
  struct jpeg_error_mgr pub;	/* "public" fields */
  jmp_buf setjmp_buffer;	/* for return to caller */
};

typedef struct my_error_mgr * my_error_ptr;


void mexFunction(int nlhs, mxArray *plhs[],
                 int nrhs, const mxArray *prhs[]) { 
    const mxArray *inputArray;
    double *doubleData;
    long i,j,k,row_stride;
    int ndims, imageType;
    const int *size;
    char *filename;
    struct jpeg_compress_struct cinfo;
    struct my_error_mgr jerr;
    FILE *outfile=NULL;
    int inClass, quality;
    long strlen;
    
    if (nrhs < 2)
    {
        mexErrMsgTxt("Not enough input arguments.");
    }      
    else if(nrhs > 2)
    {
        /* jpgwrite(rgb, 'filename.jpg', quality); */
        quality = (int) mxGetScalar(prhs[2]);
    }
    else 
    {
        quality = 75;
    }
    
    if(! mxIsChar(prhs[1]))
    {
        mexErrMsgTxt("Second argument is not a string.");
    }

    strlen = mxGetM(prhs[1]) * mxGetN(prhs[1]) + 1;
    filename = (char *) mxCalloc(strlen, sizeof(*filename));
    mxGetString(prhs[1],filename,strlen);  /* First argument is the filename */
    inputArray = prhs[0];
    
    ndims = mxGetNumberOfDimensions(inputArray);
    size = mxGetDimensions(inputArray);
    
    /*
     * Initialize the jpeg library
     */
    
    cinfo.err = jpeg_std_error(&jerr.pub);
    jerr.pub.output_message = my_output_message;
    jerr.pub.error_exit = my_error_exit;
    if(setjmp(jerr.setjmp_buffer))
    {
        /* If we get here, the JPEG code has signaled an error.
         * We need to clean up the JPEG object, close the input file, 
         * and return.
         */
        jpeg_destroy_compress(&cinfo);
        fclose(outfile);
        return;
    }

    jpeg_create_compress(&cinfo);
    
    /*
     * Open jpg file
     */
    
    if ((outfile = fopen(filename, "wb")) == NULL) {
        mexErrMsgTxt("Couldn't open file for JPEG write.");
    }
    
    /*
     * Read the jpg header to get info about size and color depth
     */
    
    jpeg_stdio_dest(&cinfo, outfile);
    
    cinfo.image_width =  size[1]; 	/* image width and height, in pixels */
    cinfo.image_height = size[0];
    
    
    inClass = mxGetClassID(inputArray); 
    
    if( (ndims == 3 && size[2] == 3) && inClass == mxUINT8_CLASS ) /* RGB image */
    {
        imageType = RGB_IMG;
        cinfo.in_color_space = JCS_RGB; 	/* colorspace of input image */
        cinfo.input_components = 3; /* # of color components per pixel */
        jpeg_set_defaults(&cinfo); /* Set default compression parameters. */
        jpeg_set_quality(&cinfo, quality, TRUE);
        WriteRGBFromUint8(&cinfo, (uint8_T *) mxGetPr(inputArray), size);
        
    }
    else if(ndims == 2 && inClass == mxUINT8_CLASS )
    {
        imageType = GRAY_IMG;
        cinfo.in_color_space = JCS_GRAYSCALE;
        cinfo.input_components = 1; /* # of color components per pixel */
        jpeg_set_defaults(&cinfo); /* Set default compression parameters. */
        jpeg_set_quality(&cinfo, quality, TRUE);
        WriteGRAYFromUint8(&cinfo, (uint8_T *) mxGetPr(inputArray), size);
        
    }
    else if(ndims == 2 && inClass == mxUINT32_CLASS ) {
        imageType = PACKED_RGB_IMG;
        cinfo.in_color_space = JCS_RGB; 	/* colorspace of input image */
        cinfo.input_components = 3; /* # of color components per pixel */
        jpeg_set_defaults(&cinfo); /* Set default compression parameters. */
        jpeg_set_quality(&cinfo, quality, TRUE);
        WriteRGBFromUint32(&cinfo, (uint32_T *) mxGetPr(inputArray), size);
    }
    else
    {
        mexErrMsgTxt("Invalid image array");
    }
    
    /*
     * Clean up
     */
    
    jpeg_finish_compress(&cinfo); fclose(outfile);
    jpeg_destroy_compress(&cinfo);
    mxFree((void *) filename);
    
    /* Wah Lah! */
    return;		
}

#if 0
static void 
WriteRGBFromDouble(j_compress_ptr cinfoPtr, 
                   double *doubleData, 
                   const int *size)
{
    int row_stride,i,j;
  JSAMPARRAY buffer;
    
  jpeg_start_compress(cinfoPtr, TRUE);
  row_stride = size[1] * 3;	
  buffer = (*cinfoPtr->mem->alloc_sarray)
                ((j_common_ptr) cinfoPtr, JPOOL_IMAGE, row_stride, 1);

  /* jpeg_write_scanlines expects an array of pointers to scanlines.
   * Here the array is only one element long, but you could pass
   * more than one scanline at a time if that's more convenient.
   */
  
  while (cinfoPtr->next_scanline < cinfoPtr->image_height) {
      /* construct a buffer which contains the next scanline with data in
       * RGBRGBRGB... format */
      i = cinfoPtr->next_scanline;
      for(j=0; j<cinfoPtr->image_width; j++)
      {
          buffer[0][3*j]   = (uint8_T) doubleData[i+(j*size[0])]; /* Red */
          buffer[0][3*j+1] = (uint8_T) doubleData[i+(j*size[0])+(size[0]*size[1])]; 
          buffer[0][3*j+2] = (uint8_T) doubleData[i+(j*size[0])+(2*size[0]*size[1])];
          
      }
      (void) jpeg_write_scanlines(cinfoPtr, buffer, 1);
  }
}

#endif /* 0 */



static void 
WriteRGBFromUint8(j_compress_ptr cinfoPtr, 
                  uint8_T *uint8Data, 
                  const int *size)
{
    int row_stride,i,j;
  JSAMPARRAY buffer;

  jpeg_start_compress(cinfoPtr, TRUE);
  row_stride = size[1] * 3;	
  buffer = (*cinfoPtr->mem->alloc_sarray)
                ((j_common_ptr) cinfoPtr, JPOOL_IMAGE, row_stride, 1);

  /* jpeg_write_scanlines expects an array of pointers to scanlines.
   * Here the array is only one element long, but you could pass
   * more than one scanline at a time if that's more convenient.
   */

  while (cinfoPtr->next_scanline < cinfoPtr->image_height) {
      /* construct a buffer which contains the next scanline with data in
       * RGBRGBRGB... format */
      i = cinfoPtr->next_scanline;
      for(j=0; j<cinfoPtr->image_width; j++)
      {
          buffer[0][3*j]   = uint8Data[i+(j*size[0])]; /* Red */
          buffer[0][3*j+1] = uint8Data[i+(j*size[0])+(size[0]*size[1])]; /* Green */
          buffer[0][3*j+2] = uint8Data[i+(j*size[0])+(2*size[0]*size[1])]; /* Blue */
      }
      (void) jpeg_write_scanlines(cinfoPtr, buffer, 1);
  } 
}  
  
static void 
WriteRGBFromUint32(j_compress_ptr cinfoPtr, 
                  uint32_T *uint32Data, 
                  const int *size)
{
    int row_stride,i,j;
    JSAMPARRAY buffer;

    jpeg_start_compress(cinfoPtr, TRUE);
    row_stride = size[1] * 3;	
    buffer = (*cinfoPtr->mem->alloc_sarray)
        ((j_common_ptr) cinfoPtr, JPOOL_IMAGE, row_stride, 1);

    /* jpeg_write_scanlines expects an array of pointers to scanlines.
     * Here the array is only one element long, but you could pass
     * more than one scanline at a time if that's more convenient.
     */

    while (cinfoPtr->next_scanline < cinfoPtr->image_height) {
        /* construct a buffer which contains the next scanline with data in
         * RGBRGBRGB... format */
        i = cinfoPtr->next_scanline;
        for(j=0; j<cinfoPtr->image_width; j++)
            {
                /* Shifts are hardcoded TCPI offsets of IMSAVE driver */
                buffer[0][3*j]   = (*uint32Data >> 16) & (uint32_T)255; /* Red Pixel */
                buffer[0][3*j+1] = (*uint32Data >>  8) & (uint32_T)255; /* Green Pixel */
                buffer[0][3*j+2] =  *uint32Data        & (uint32_T)255; /* Blue Pixel */
                uint32Data++;
            }
        (void) jpeg_write_scanlines(cinfoPtr, buffer, 1);
    } 
}  
  


#if 0
static void 
WriteGRAYFromDouble(j_compress_ptr cinfoPtr, 
                   double *doubleData, 
                   const int *size)
{
    int row_stride,i,j;
  JSAMPARRAY buffer;

  jpeg_start_compress(cinfoPtr, TRUE);
  row_stride = size[1];	
  buffer = (*cinfoPtr->mem->alloc_sarray)
                ((j_common_ptr) cinfoPtr, JPOOL_IMAGE, row_stride, 1);

  /* jpeg_write_scanlines expects an array of pointers to scanlines.
   * Here the array is only one element long, but you could pass
   * more than one scanline at a time if that's more convenient.
   */
  
  while (cinfoPtr->next_scanline < cinfoPtr->image_height) {
      /* construct a buffer which contains the next scanline with data in
       * RGBRGBRGB... format */
      i = cinfoPtr->next_scanline;
      for(j=0; j<cinfoPtr->image_width; j++)
      {
          buffer[0][j]   = (uint8_T) doubleData[i+(j*size[0])]; 
      }
      (void) jpeg_write_scanlines(cinfoPtr, buffer, 1);
  }
}
#endif /* 0 */



static void 
WriteGRAYFromUint8(j_compress_ptr cinfoPtr, 
                  uint8_T *uint8Data, 
                  const int *size)
{
    int row_stride,i,j;
  JSAMPARRAY buffer;

  jpeg_start_compress(cinfoPtr, TRUE);
  row_stride = size[1];	
  buffer = (*cinfoPtr->mem->alloc_sarray)
                ((j_common_ptr) cinfoPtr, JPOOL_IMAGE, row_stride, 1);

  /* jpeg_write_scanlines expects an array of pointers to scanlines.
   * Here the array is only one element long, but you could pass
   * more than one scanline at a time if that's more convenient.
   */

  while (cinfoPtr->next_scanline < cinfoPtr->image_height) {
      /* construct a buffer which contains the next scanline with data in
       * RGBRGBRGB... format */
      i = cinfoPtr->next_scanline;
      for(j=0; j<cinfoPtr->image_width; j++)
      {
          buffer[0][j]   = uint8Data[i+(j*size[0])]; 
      }
      (void) jpeg_write_scanlines(cinfoPtr, buffer, 1);
  } 
}  
  



/*
 * Here's the routine that will replace the standard error_exit method:
 */

static void
my_error_exit (j_common_ptr cinfo)
{
  /* cinfo->err really points to a my_error_mgr struct, so coerce pointer */
  my_error_ptr myerr = (my_error_ptr) cinfo->err;

  /* Always display the message. */
  /* We coulD postpone this until after returning, if we chose. */
  (*cinfo->err->output_message) (cinfo);

  /* Return control to the setjmp point */
  longjmp(myerr->setjmp_buffer, 1);
}


/* 
 *  Here's the routine to replace the standard output_message method:
 */


static void
my_output_message (j_common_ptr cinfo)
{
  char buffer[JMSG_LENGTH_MAX];

  /* Create the message */
  (*cinfo->err->format_message) (cinfo, buffer);

  /* Send it to stderr, adding a newline */
  mexWarnMsgTxt(buffer);
}