# include <stdio.h>
# include <strings.h>  // for strcasecmp
# include <unistd.h>   // for unlink
# include <pslib.h>
# include "jpeglib.h"
# include "psImageJpeg.h"

static void psImageJpegColormapFree (psImageJpegColormap *map)
{

    if (map == NULL)
        return;

    psFree (map->red);
    psFree (map->green);
    psFree (map->blue);
    return;
}

psImageJpegColormap *psImageJpegColormapAlloc ()
{

    psImageJpegColormap *map;
    map = psAlloc (sizeof(psImageJpegColormap));
    psMemSetDeallocator(map, (psFreeFunc) psImageJpegColormapFree);

    map->red   = psVectorAlloc (256, PS_TYPE_U8);
    map->blue  = psVectorAlloc (256, PS_TYPE_U8);
    map->green = psVectorAlloc (256, PS_TYPE_U8);

    return (map);
}

psImageJpegColormap *psImageJpegColormapSet (psImageJpegColormap *map, char *name)
{

    if (map == NULL) {
        map = psImageJpegColormapAlloc ();
    }

    /* grayscale */
    if ((!strcasecmp (name, "grayscale")) || (!strcasecmp (name, "greyscale"))) {
        for (int i = 0; i < map->red->n; i++) {
            map->red->data.U8[i]   = PS_JPEG_RANGELIM(i);
            map->green->data.U8[i] = PS_JPEG_RANGELIM(i);
            map->blue->data.U8[i]  = PS_JPEG_RANGELIM(i);
        }
        return map;
    }

    /* -grayscale */
    if ((!strcasecmp (name, "-grayscale")) || (!strcasecmp (name, "-greyscale"))) {
        for (int i = 0; i < map->red->n; i++) {
            map->red->data.U8[i]   = PS_JPEG_RANGELIM(256 - i);
            map->green->data.U8[i] = PS_JPEG_RANGELIM(256 - i);
            map->blue->data.U8[i]  = PS_JPEG_RANGELIM(256 - i);
        }
        return map;
    }

    /* rainbow */
    if (!strcasecmp (name, "rainbow")) {
        int I1 = 0.25*map->red->n;
        int I2 = 0.50*map->red->n;
        int I3 = 0.75*map->red->n;
        for (int i = 0; i < I1; i++) {
            map->red->data.U8[i]   = 0;
            map->green->data.U8[i] = 0;
            map->blue->data.U8[i]  = PS_JPEG_RANGELIM(4*i);
        }
        for (int i = I1; i < I2; i++) {
            map->red->data.U8[i]   = PS_JPEG_RANGELIM(4*(i - I1));
            map->green->data.U8[i] = 0;
            map->blue->data.U8[i]  = PS_JPEG_RANGELIM(4*(I2 - i));
        }
        for (int i = I2; i < I3; i++) {
            map->red->data.U8[i]   = 255;
            map->green->data.U8[i] = 4*(i - I2);
            map->blue->data.U8[i]  = 0;
        }
        for (int i = I3; i < map->red->n; i++) {
            map->red->data.U8[i]   = 255;
            map->green->data.U8[i] = 255;
            map->blue->data.U8[i]  = PS_JPEG_RANGELIM(4*(i - I3));
        }
        return map;
    }

    /* heat */
    if (!strcasecmp (name, "heat")) {
        int I1 = 0.25*map->red->n;
        int I2 = 0.50*map->red->n;
        int I3 = 0.75*map->red->n;
        for (int i = 0; i < I1; i++) {
            map->red->data.U8[i]   = PS_JPEG_RANGELIM(2*i);
            map->green->data.U8[i] = 0;
            map->blue->data.U8[i]  = 0;
        }
        for (int i = I1; i < I2; i++) {
            map->red->data.U8[i]   = PS_JPEG_RANGELIM(2*i);
            map->green->data.U8[i] = PS_JPEG_RANGELIM(2*(i - I1));
            map->blue->data.U8[i]  = 0;
        }
        for (int i = I2; i < I3; i++) {
            map->red->data.U8[i]   = 255;
            map->green->data.U8[i] = PS_JPEG_RANGELIM(2*(i - I1));
            map->blue->data.U8[i]  = PS_JPEG_RANGELIM(2*(i - I2));
        }
        for (int i = I3; i < map->red->n; i++) {
            map->red->data.U8[i]   = 255;
            map->green->data.U8[i] = 255;
            map->blue->data.U8[i]  = PS_JPEG_RANGELIM(2*(i - I2));
        }
        return map;
    }

    // invalid colormap : warn user
    map = psImageJpegColormapSet (map, "greyscale");
    return map;
}

bool psImageJpeg (psImageJpegColormap *map, psImage *image, char *filename, float min, float max)
{

    struct jpeg_compress_struct cinfo;
    struct jpeg_error_mgr jerr;

    int pixel;
    JSAMPLE *jpegLine;   // Points to data for current line
    JSAMPROW jpegLineList[1];  // pointer to JSAMPLE row[s]
    JSAMPLE *outPix;

    if (map->red == NULL)
        return false;
    if (map->green == NULL)
        return false;
    if (map->blue == NULL)
        return false;

    /* JPEG init calls */
    cinfo.err = jpeg_std_error (&jerr);
    jpeg_create_compress (&cinfo);

    /* open file, prep for jpeg */
    FILE *f = fopen (filename, "w");
    if (f == NULL) {
        fprintf (stderr, "failed to open %s for output\n", filename);
        return (TRUE);
    }
    jpeg_stdio_dest(&cinfo, f);

    /* set up color jpeg buffers */
    int quality = 75;
    cinfo.image_width = image->numCols; // image width and height, in pixels
    cinfo.image_height = image->numRows;
    cinfo.input_components = 3;
    cinfo.in_color_space = JCS_RGB;
    jpeg_set_defaults (&cinfo);
    jpeg_set_quality (&cinfo, quality, true); // limit to baseline-JPEG values
    jpeg_start_compress (&cinfo, true);

    jpegLine = psAlloc (3*image->numCols*sizeof(JSAMPLE));
    jpegLineList[0] = jpegLine;

    psU8 *Rpix = map->red->data.U8;
    psU8 *Gpix = map->green->data.U8;
    psU8 *Bpix = map->blue->data.U8;

    float zero = min;
    float scale = 256.0/(max - min);

    for (int j = 0; j < image->numRows; j++) {
        psF32 *row = image->data.F32[j];

        outPix = jpegLine;
        for (int i = 0; i < image->numCols; i++, outPix += 3) {
            pixel = PS_JPEG_SCALEVALUE(row[i],zero,scale);
            outPix[0] = Rpix[pixel];
            outPix[1] = Gpix[pixel];
            outPix[2] = Bpix[pixel];
        }
        jpeg_write_scanlines (&cinfo, jpegLineList, 1);
    }

    jpeg_finish_compress (&cinfo);
    fclose (f);
    jpeg_destroy_compress (&cinfo);

    psFree (jpegLine);
    return true;
}
