/***************************************************************************
 *   Copyright (C) 2012 by CADS/UV Software Team,                          *
 *                         Indian Institute of Astrophysics                *
 *                         Bangalore 560034                                *
 *                         cads_AT_iiap.res.in                             *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
/*---------------------------------------------------------------------------
  This program will write out FITS data into a jpeg file.
  Assume that data are in first extension and are two dimensional

  Libraries needed are CFITSIO, JPEG
  Original author: Jayant Murthy
  Last Modification: 11/05/02

  Reconstruction and additional features: Reks
 *---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------
  Modification History

  Reks, 5th October 2007: v0.01
  1. configure checks for jpeglib.
     optional extra arguments --with-jpeglib=<location>
                              --with-cfitsio=<location>
  2. Makes blank images! datamin and max are not set..

  Reks, 6th October 2007: v0.10
  1. Added 2 functions to fetch max and min from image
  2. user can choose various scales like log/square/square-root
  3. Increased verbosity, added banner, help message, etc
  4. Inputs through getopt. will Ask if mandatory input (fits file) is not
     provided.
  5. fitsio2

  Reks, 10th October 2007: v0.9.0 (close to 1.0)
  1. Found our niche :-) Can handle wild cards at command line,
     batch processing.
  2. Now users cannot specify the output file name though.. Its fixed to
     <input_file_root>.jpg.
  3. No getopt for input file.. getopt is limited to various image scaling
     options.

 Reks, 17th October 2007, v0.9.5
  1. Merged all image scaling options into one input param: -s
     optarg (from getopt) will decide scaling function. Now there is
     enough room for additional flags (in future).

 Reks, 22 June 2010 v1.0.0
  1. Removed datamax() and datamin(). That was required only once!
  2. option histoeq shifted to image enhancement option
  3. Added "normalize" option
     Now users can do scaling + enhance (if at all someone feels like it)
  4. All options tested. First stable release :)

  Reks 28 June 2012
  1. rebranded to CADS/UVS
 *-------------------------------------------------------------------------*/

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

/*Header Definitions*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <libgen.h>
#include <dirent.h>
#include <fitsio2.h>
#include <jpeglib.h>

#define PROGRAM "fits2jpeg"
#define r2d (90./atan2(1,0))
#define MAX_TEXT 150
int my_getopt(int, char * const *, const char *);
void signals_handler(int);
void set_signals(void);

char *optarg;
int optindex;
/*---------------------------------------------------------------------------*/
void PRINT_BANNER()
{
    printf("\n");
    printf(" %s %s\n", PROGRAM, VERSION);
    printf(" Converts fits image files to jpeg\n");
    printf("-----------------------------------------------------------\n");
}
/*---------------------------------------------------------------------------*/
void usage()
{
    printf ("\n Usage: fits2jpeg [-s <options>] <fits_file>             \n");
    printf ("  Options are:                                             \n");
    printf ("    -h help                                                \n");
    printf ("    -s <scale_type>                                        \n");
    printf ("       scale for output image, where <scale_type> can be:  \n");
    printf ("         linear         Linear scale, default            \n\n");
    printf ("         sqroot         for square root scale              \n");
    printf ("         square         for quadratic scale                \n");
    printf ("         cubic          for cubic scale                    \n");
    printf ("         log            for log scale                      \n");
    printf ("    -e <operation>                                         \n");
    printf ("       Image enhancements, where <operation> can be:       \n");
    printf ("         normalize      for linear histogram stretch       \n");
    printf ("         equalize       for histogram equalization         \n");
    printf ("  Output will be written to <fits_file_root>.jpg. Wild card\n");
    printf ("  entries allowed in <fits_file>; For eg: *.fits, m31*.fits\n");
    printf ("  ngc???.fits etc.                                         \n");
    printf ("-----------------------------------------------------------\n");
    printf ("Copyright (c) 2012 The CADS Team, IIAp. [GPL v3.0 or later]\n");
    printf ("Report Bugs to: http://cads.iiap.res.in/bugzilla           \n");
    printf ("Documentation : http://cads.iiap.res.in/software         \n\n");
    exit(1);
}
/*---------------------------------------------------------------------------*/

/*--------------------------- Begin Main Progam -----------------------------*/
int main(int argc, char *argv[], char *envp[])
{
    fitsfile *fptr;

    int scale = 0, process = 0, status = 0, nfound, anynull;
    unsigned int i = 0, j = 0, npixels = 0;
    long naxes[2], row_stride;
    float datamin = 0.0, datamax = 0.0;
    float nullval = 0.0, scl_data;
    float *data;
    float tmp = 0.0;
    float hist[256] = {0.0}, cumhist[256] = {0.0};
    char jpeg_file_name[MAX_TEXT] = "", fits_file_name[MAX_TEXT] = "";

    int opt;                                                   /* For getopt */
    extern char *optarg;
    extern int optindex;

    FILE *jpeg_file;

    struct jpeg_error_mgr jerr;
    struct jpeg_compress_struct cinfo;
    int JMAXVAL = 255;                    /* Max value for greyscale in jpeg */
    JSAMPROW row_pointer[1];
    JSAMPLE *image_buffer;
    /*---------------- End of variable initialization -----------------------*/

    PRINT_BANNER();        /* Header: Talks about program, version & purpose */
    set_signals();                                  /* Trap un-natural exits */

    /*-------------------- Parse command line inputs ------------------------*/

    if (argc < 2) usage();  /* People who don't give any arguments need help */
    while ( ( opt = my_getopt ( argc, argv, "s:e:h" ) ) != EOF )
        switch ( opt )
        {
        case 's':
            if (strcmp(optarg, "sqroot")   ==0) scale = 1;
            if (strcmp(optarg, "square")   ==0) scale = 2;
            if (strcmp(optarg, "cube")     ==0) scale = 3;
            if (strcmp(optarg, "log")      ==0) scale = 4;
            break;

        case 'e':
            if (strcmp(optarg, "normalize")==0) process = 1;
            if (strcmp(optarg, "equalize") ==0) process = 2;
            break;

        case 'h':
            usage();
            break;

        case '*':
            usage();
            break;
        }

    if ((argc - optindex) == 0) usage();       /* Somebody gave only options */

    /*---------------------- Begin processing files -------------------------*/

    /* Getopt would have processed all argument with a "-" in front. optindex
       is the counter, which keeps the number of options supplied. Whatever
       remains in the list of command line arguments are the fits files.
       NOTE: POSIX systems will do the regex matching and expand wild card
             entries for us. Even windoze XP Pro is known to do that...      */
    for (j = optindex; j < argc; j++)
    {

        strcpy(fits_file_name, argv[j]);

        strcpy(jpeg_file_name, fits_file_name);
        strtok(jpeg_file_name, ".");
        strcat(jpeg_file_name, ".jpg");

        fits_open_file(&fptr, fits_file_name, READONLY, &status);
        fits_read_keys_lng(fptr, "NAXIS", 1, 2, naxes, &nfound, &status);
        if (status)
        {
            fprintf(stderr, "ERROR  : Could not open file: %s\n",
                    fits_file_name);
            exit(EXIT_FAILURE);
        }/*Endif*/

        /* Read in data */
        npixels = naxes[0] * naxes[1];
        data = (float *) malloc(sizeof(float) * npixels);

        nullval = 0;
        if (fits_read_img(fptr, TFLOAT, 1, npixels, &nullval, data, &anynull,
                          &status))
        {
            fprintf(stderr, "ERROR  : No valid fits image data in %s\n",
                    fits_file_name);
            exit(EXIT_FAILURE);
        }/*Endif*/

        fits_close_file(fptr, &status);
        /*Close data file*/
        printf("INFO   : data input from %s\n", fits_file_name);

        /* find min & max in image */
        datamax = -1.0e9;
        datamin = +1.0e9;
        for (i = 0; i < npixels; ++i)
        {
            if (data[i] > datamax) datamax = data[i];
            if (data[i] < datamin) datamin = data[i];
        } /*endfor*/

        /* Convert data into bytscaled values for jpeg file                 */
        /* the dynamic range is reduced to 255 for jpeg                     */
        scl_data = (datamax - datamin)/(float)JMAXVAL;
        for (i = 0; i < npixels; ++i)
            data[i] = (data[i] - datamin)/scl_data;

        /* All data is now squeezed into the range 0 - 255                   */
        /* NOTE: At this point onwards min & max is 0 and 255 respectively   */


        /* Now we need to look at user options..                             */
        /* 1. scale it as per user's requirement.                            */
        /* 2. image enhancing operations                                     */
                                                    /* Allocate image buffer */
        image_buffer = (unsigned char *) malloc(sizeof(char) * npixels);
        scl_data = 1.0;
        switch (scale)
        {
        case 1 :                                              /* Square root */
            printf("INFO   : Using square-root scale\n");
            scl_data = sqrt((float)JMAXVAL)/(float)JMAXVAL;
            for (i = 0; i < npixels; ++i)
                image_buffer[i] = (int)(sqrt(data[i])/scl_data);
            break;

        case 2 :                                                   /* Square */
            printf("INFO   : Using quadratic scale\n");
            scl_data = pow((float)JMAXVAL,2)/(float)JMAXVAL;
            for (i = 0; i < npixels; ++i)
                image_buffer[i] = (int)abs((pow(data[i],2) - 1.0)/scl_data);
            break;

        case 3 :                                                    /* Cubic */
            printf("INFO   : Using cubic scale\n");
            scl_data = pow((float)JMAXVAL,3)/(float)JMAXVAL;
            for (i = 0; i < npixels; ++i)
                image_buffer[i] = (int)abs((pow(data[i],3) - 1.0)/scl_data);
            break;

        case 4 :                                                      /* log */
            printf("INFO   : Using log scale\n");
            scl_data = log(1.0 + (float)JMAXVAL)/(float)JMAXVAL;
            for (i = 0; i < npixels; ++i)
                image_buffer[i] = (int)((log(abs(data[i]) + 1.0))/scl_data);
            break;

        default:                                   /* Linear scale (min-max) */
            for (i = 0; i < npixels; ++i)
                image_buffer[i] = (int)(data[i]);
            break;
        }

        /* initialize image histogram. ensure all are zeroes in hist[]       */
        for (i = 0; i <= JMAXVAL; ++i) hist[i] = 0;

        /* construct the image histogram */
        tmp = 1.0/(float)npixels;
        for (i = 0; i <= npixels; ++i)
            hist[(int)floor(data[i])] += tmp;

        /* And the cumulative histogram */
        cumhist[0] = hist[0];
        for (i = 1; i <= JMAXVAL; ++i)
            cumhist[i] += cumhist[i - 1] + hist[i];

        switch (process)
        {
        case 1 :                                         /* contrast stretch */
            printf("INFO   : Performing contrast stretch (normalization) \n");

            datamax = (float)JMAXVAL;
            datamin = 0.0;

            /* We need to go through the cumulative histogram to pick the
               appropriate values for datamin and datamax                    */
            i = 0;
            while (i < JMAXVAL)
            {
                if (cumhist[i] >= 0.01)
                {
                    datamin = (float) i;
                    break;
                }
                i++;
            }
            i = JMAXVAL;
            while (i > 0)
            {
                if (cumhist[i] <= 0.99)
                {
                    datamax = (float) i;
                    break;
                }
                i--;
            }
            scl_data = (datamax - datamin)/(float)JMAXVAL;
            for (i = 0; i < npixels; ++i)
            {
                if (image_buffer[i] >= datamax)
                    image_buffer[i] = JMAXVAL;
                else if (image_buffer[i] <= datamin)
                    image_buffer[i] = 0;
                else
                    image_buffer[i]=(int)abs((image_buffer[i]-datamin)/scl_data);
            }
            break;


        case 2 :                                   /* histogram equalization */
            printf("INFO   : Performing Histogram Equalization\n");
            for (i = 0; i <  npixels; ++i)
                image_buffer[i] = cumhist[image_buffer[i]]*255;
            break;
        }

        /*Write out data into JPEG file*/
        cinfo.err = jpeg_std_error(&jerr);             /* JPEG error handler */
        jpeg_create_compress(&cinfo);                 /* JPEG initialization */
        jpeg_file = fopen(jpeg_file_name, "wb");/* Open JPEG file for writing*/
        jpeg_stdio_dest(&cinfo, jpeg_file); /* Send compressed data to stdio */
        cinfo.image_width = naxes[0];                         /* Image width */
        cinfo.image_height = naxes[1];                       /* Image height */
        cinfo.input_components = 1;            /* Number of colors per pixel */
        cinfo.in_color_space = JCS_GRAYSCALE;   /* colorspace of input image */
        jpeg_set_defaults(&cinfo);                      /* Set JPEG defaults */
        jpeg_start_compress(&cinfo, TRUE);       /* default data compression */
        row_stride = naxes[0];            /* JSAMPLEs per row inimage buffer */

        /* Now we have to turn the data upside down                          */
        while (cinfo.next_scanline < cinfo.image_height)
        {
            row_pointer[0] = &image_buffer[(cinfo.image_height -
                                            cinfo.next_scanline) *
                                            row_stride];
            (void) jpeg_write_scanlines(&cinfo, row_pointer, 1);
        }/*Loop through file writing one line at a time*/

        jpeg_finish_compress(&cinfo);                  /* Finish compression */
        fclose(jpeg_file);                                     /* Close file */
        jpeg_destroy_compress(&cinfo);                     /* Release memory */

        free(data);
        free(image_buffer);
        printf("INFO   : wrote output to %s\n", jpeg_file_name);
    }
    exit(EXIT_SUCCESS);
} /*End of program*/