#include <stdio.h>
#include <strings.h>  // for strcasecmp
#include <string.h>

#include <jpeglib.h>

#include "psMemory.h"
#include "psImage.h"
#include "psVector.h"
#include "psError.h"
#include "psAssert.h"
#include "psImageJpeg.h"

static void imageJpegColormapFree(psImageJpegColormap *map)
{

    if (!map) {
        return;
    }

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

psImageJpegColormap *psImageJpegColormapAlloc ()
{
    psImageJpegColormap *map = psAlloc(sizeof(psImageJpegColormap));
    psMemSetDeallocator(map, (psFreeFunc)imageJpegColormapFree);

    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, const char *name)
{

    if (!map) {
        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
    psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Invalid colormap name: %s --- using greyscale\n", name);
    return psImageJpegColormapSet (map, "greyscale");
}

bool psImageJpeg(const psImageJpegColormap *map, const psImage *image, const char *filename,
                 float min, float max)
{
    PS_ASSERT_PTR_NON_NULL(map, false);
    PS_ASSERT_IMAGE_NON_NULL(image, false);
    PS_ASSERT_VECTOR_NON_NULL(map->red, false);
    PS_ASSERT_VECTOR_NON_NULL(map->green, false);
    PS_ASSERT_VECTOR_NON_NULL(map->blue, false);
    PS_ASSERT_PTR_NON_NULL(filename, false);
    PS_ASSERT_INT_POSITIVE(strlen(filename), false);
    PS_ASSERT_FLOAT_REAL(min, false);
    PS_ASSERT_FLOAT_REAL(max, false);

    float zero, scale;
    struct jpeg_compress_struct cinfo;
    struct jpeg_error_mgr jerr;

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

    /* 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) {
        psError(PS_ERR_IO, true, "failed to open %s for output\n", filename);
        return false;
    }
    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;

    if (max == min) {
        zero = min - 0.1;
        scale = 256.0/0.2;
    } else {
        zero = min;
        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) {
            if (isfinite(row[i])) {
                pixel = PS_JPEG_SCALEVALUE(row[i],zero,scale);
            } else {
                pixel = 0;
            }
            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;
}
