/*

  FITSPNG     FITS to PNG converter
  Copyright (C) 2006-2019  Filip Hroch, Masaryk University, Brno, CZ

  Fitspng 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 3 of the License, or
  (at your option) any later version.

  Fitspng 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 Fitspng.  If not, see <http://www.gnu.org/licenses/>.

*/

#include "fitspng.h"
#include <png.h>
#include <fitsio.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <time.h>

#include <assert.h>

#define NFKEYS 11
#define NSTDKEYS 10
#define NAXES 3
#define SKIP 10

/* macros significantly speed-ups processing */
#define MAX(x,y) (x) > (y) ? (x) : (y)
#define MIN(x,y) (x) < (y) ? (x) : (y)
#define PRESCALE(flux,thresh,sense) ((flux - thresh) / sense)
#define CLIP(flux,high) MIN(MAX(((int)(flux+0.5)),0),high)


float sRGBGamma(float r)
{
  const float q = 1.0 / 2.4;
  return r < 0.0031308f ? 12.92f * r : 1.055f*powf(r,q) - 0.055f;
}

void XYZ_sRGB(float X, float Y, float Z, float *R, float *G, float *B)
{
  float x = X / 100;
  float y = Y / 100;
  float z = Z / 100;

  /* transform to RGB and apply gamma function */
  *R = sRGBGamma( 3.2406f*x - 1.5372f*y - 0.4986f*z);
  *G = sRGBGamma(-0.9689f*x + 1.8758f*y + 0.0415f*z);
  *B = sRGBGamma( 0.0557f*x - 0.2040f*y + 1.0570f*z);
}

float AdobeGamma(float r)
{
  const float q = 1.0 / 2.19921875;
  return powf(r,q);
}

void XYZ_AdobeRGB(float X, float Y, float Z, float *R, float *G,float *B)
{
  float x = X / 100;
  float y = Y / 100;
  float z = Z / 100;

  /* transform to RGB and apply gamma function */
  *R = AdobeGamma( 2.04159f*x - 0.56501f*y - 0.34473f*z);
  *G = AdobeGamma(-0.96924f*x + 1.87597f*y + 0.04156f*z);
  *B = AdobeGamma( 0.01344f*x - 0.11836f*y + 1.01517f*z);
}


float fscale(float r, int ctype, float f0, float zero)
{
  switch (ctype) {
  case 0:
    return r;
  case 5:
    return (r > 0.0f ? f0*sqrtf(r) : 0.0f) + zero;
  case 6:
    return f0*r*r + zero;
  case 7:
    return f0*(2.0f/(1.0f + expf(-2.5f*r)) - 1.0f) + zero;
  default:
    return r - trunc(r); /* iso-contours */
  }
}

float Scotopic(float X, float Y, float Z)
{
  return 0.36169f*Z + 1.18214f*Y - 0.80498f*X;
}

float farray2d(float *pic, int w, int h, int x, int y, int d)
{
  int xdim,ydim,ndim;
  int x1 = x, y1 = y, x2 = x + d, y2 = y + d;
  int i, j, l, n;
  double s, d2;

  assert(pic != NULL);
  assert(d >= 1);

  xdim = w;
  ydim = h;

  if( d == 1 ) {

    assert((0 <= x && x < xdim) && (0 <= y && y < ydim));
    return(pic[y*xdim + x]);

  }
  else if( d > 1 ) {

    d2 = d*d;

    if( (0 <= x1 && x1 < xdim) && (0 <= y1 && y1 < ydim) &&
	(0 <= x2 && x2 < xdim) && (0 <= y2 && y2 < ydim)) {

      s = 0.0;
      for(j = y1; j < y2; j++) {
	float *nrow = pic + j*xdim;
	for(i = x1; i < x2; i++)
	  s = s + *(nrow + i);
      }
      return(s/d2);
    }
    else {
      s = 0.0;
      ndim = w*h;
      n = 0;
      for(j = y1; j < y2; j++) {
        for(i = x1; i < x2; i++) {
	  l = i+j*xdim;
	  if( 0 <= l && l < ndim ) {
	    s = s + pic[l];
	    n++;
	  }
	}
      }
      return( n > 0 ? s/n : 0);
    }
  }
  return(0.0);
}


int fitspng(char *fitsname, char *png, int bit_depth,
	    float qblack,float rsense, int scale,
	    float f0, float thresh, float sense, float sthresh, float swidth,
	    int ctype, float satur, int dcspace, float zero, int verb,
	    int est, int scotop, int saturate)
{
  char *fkeys[NFKEYS] = {"OBJECT", "OBSERVER", "FILTER", "DATE-OBS",
			 "CAMTYPE", "EXPTIME", "SITE", "TEMPERAT",
			 "XFACTOR", "YFACTOR", "TELESCOP" };
  char *stdkeys[NSTDKEYS] = {"Title","Author","Description","Copyright",
			     "Creation Time","Software","Disclaimer",
			     "Warning","Source","Comment"};
  char *stdvalues[NSTDKEYS];
  char *fval[NFKEYS];
  long naxes[NAXES];

  /* FITS input */
  fitsfile *f;
  int naxis,bitpix,status,nullval=0,fpixel=1;
  int i,j,k,m,n,lim;
  char line[FLEN_CARD];

  /* data scaling */
  float r;

  /* png */
  png_uint_32 height, width, bytes_per_pixel, color_bytes;
  png_byte *image;
  FILE *fp;
  png_structp png_ptr;
  png_infop info_ptr;
  png_text text_ptr[NSTDKEYS];
  char buf[NFKEYS*FLEN_CARD + 100];
  png_bytep *row_pointers;
  char *tm[6], *c, *c0;

  /* intensity scaling */
  int ii,jj;
  float flux, X, Y, Z, R, G, B, L, a, b;
  int iflux;
  float *pic;
  int imgpix[3] = {0,0,0};

  /* display colour-space */
  void (*XYZ_RGB)(float,float,float,float *,float *,float *);
  float (*Gamma)(float);

  /* --------------------------------------------------------------------*/
  /* Part: Initilization */

  status = 0;

  if( dcspace == 1 ) {
    XYZ_RGB = &XYZ_AdobeRGB;
    Gamma = &AdobeGamma;
  }
  else { /* if( dcspace == 0 ) */
    XYZ_RGB = &XYZ_sRGB;
    Gamma = &sRGBGamma;
  }

  image = NULL;
  pic = NULL;

  for( i = 0; i < NFKEYS; i++)
    fval[i] = NULL;

  for( i = 0; i < NSTDKEYS; i++)
    stdvalues[i] = NULL;

  for( i = 0; i < 6; i++ )
    tm[i] = NULL;

  if( verb )
    fprintf(stderr,"Finishing initialisation..\n");

  /* --------------------------------------------------------------------*/
  /* Part: Load input FITS */

  /* check whatever filename has extension - usefull only for selecting
     of single bands in colour images
  */

  fits_open_image(&f,fitsname, READONLY, &status);
  if( status )
    goto finish;

  fits_get_img_type(f,&bitpix,&status);
  fits_get_img_dim(f,&naxis,&status);

  if( status )
    goto finish;

  if( !(naxis == 2 || naxis == 3)) {
    fprintf(stderr,"Crash: Only grey or colour FITS files are supported.\n");
    goto finish;
  }

  if( naxis == 3 ) {
    fits_read_key(f,TSTRING,"CSPACE",line,NULL,&status);
    if( status || strstr(line,"XYZ") == NULL ) {
      fprintf(stderr,"Crash: Only CIE 1931 XYZ colour-space is supported yet.\n");
      goto finish;
    }
  }

  /* keywords */
  for(i = 0; i < NFKEYS; i++) {
    fits_read_key(f,TSTRING,fkeys[i],line,NULL,&status);
    if( status == 0 ) {
      if( fval[i] == NULL )
	fval[i] = strdup(line);
      else {
	if( verb )
	  fprintf(stderr,"Ignoring keywords: %s=%s\n",fkeys[i],fval[i]);
      }
    }
    else
      status = 0;
  }

  fits_get_img_size(f,NAXES,naxes,&status);
  if( status )
    goto finish;

  n = 1;
  for( i = 0; i < naxis; i++)
    n = n*naxes[i];

  pic = malloc(n*sizeof(float));
  if( pic == NULL ) {
    fprintf(stderr,"Crash: There is no room for an input image.\n");
    goto finish;
  }

  fits_read_img(f,TFLOAT,fpixel,n,&nullval,pic,&i,&status);
  if( status )
    goto finish;

  fits_close_file(f, &status);

  if( verb )
    fprintf(stderr,"Finishing FITS load..\n");

  /* --------------------------------------------------------------------*/
  /* Part: Estimation of intensity parameters */

  if( abs(bitpix) > 8 && bit_depth < 16 && est ) {

    int res;
    int npix = naxes[0]*naxes[1];

    float *ypic = naxis == 2 ? pic : pic + npix;
    if( naxis == 2 ) /* gray images, single plane */
      ypic = pic;
    else  /* Y component for colour images */
      ypic = pic + npix;

    res = tone(npix,ypic,qblack,rsense,&thresh,&sense,verb);

    if( ! res ) {
      fprintf(stderr,"A wicked star constellation for estimation of scale.\n");
      goto finish;
    }

    if( verb )
      fprintf(stderr,"Finishing of estimate of scale parameters..\n");
  }

  /* --------------------------------------------------------------------*/
  /* Part: Intensity conversion */

  /* fill an  output array */
  height = naxes[1]/scale;
  width = naxes[0]/scale;
  bytes_per_pixel = bit_depth / 8;
  color_bytes = naxis == 2 ? 1 : 3;
  n = width*height;

  if( height < 1 || width < 1 ) {
    fprintf(stderr,"Size of scaled image is zero.\n");
    goto finish;
  }

  /* setup bit limit */
  lim = 1;
  for( i = 0; i < bit_depth; i++)
    lim = 2*lim;
  lim = lim - 1;

  if( verb )
    fprintf(stderr,"thresh=%f sense=%f lim=%d\n", thresh, sense, lim);

  if( (image = malloc(height*width*bytes_per_pixel*color_bytes)) == NULL ) {
    fprintf(stderr,"Crash: There is no room for an output image.\n");
    goto finish;
  }

  m = 0;
  int npix = naxes[0] * naxes[1];
  for( j = height-1; j >= 0; j-- ) {
    jj = j*scale;

    for( i = 0; i < width; i++ ) {
      ii = i*scale;

      /* greyscale */
      if( naxis == 2 ) {

	flux = farray2d(pic,naxes[0],naxes[1],ii,jj,scale);
	float f = MAX(PRESCALE(flux,thresh,sense),0);
	if( ctype > 0 )
	  f = fscale(f,ctype,f0,zero);
	float g = Gamma(f);
	imgpix[0] = CLIP(255.0*g,lim);

      }
      else if( naxis == 3 ) {

	/* colour */
	Z = farray2d(pic,       naxes[0],naxes[1],ii,jj,scale);
	Y = farray2d(pic + npix,naxes[0],naxes[1],ii,jj,scale);
	X = farray2d(pic+2*npix,naxes[0],naxes[1],ii,jj,scale);

	X = MAX(100*PRESCALE(X,thresh,sense),0.0);
	Y = MAX(100*PRESCALE(Y,thresh,sense),0.0);
	Z = MAX(100*PRESCALE(Z,thresh,sense),0.0);


	/* intensity scaling */
	if( ctype > 0 || saturate ) {

	  /* convert to Lab */
	  XYZ_Lab(X,Y,Z,&L,&a,&b);

	  /* ITT */
	  if( ctype > 0 )
	    L = 100 * fscale(L / 100,ctype,f0,zero);

	  /* colour saturation */
	  if( saturate ) {
	    float hue = atan2f(b,a);
	    float r2 = a*a + b*b;
	    float chroma = r2 > 0 ? satur*sqrtf(r2) : 0;
	    a = chroma*cosf(hue);
	    b = chroma*sinf(hue);
	  }

	  /* convert back to XYZ */
	  Lab_XYZ(L,a,b,&X,&Y,&Z);

	}

	/* apply night vision simulation */
	if( scotop ) {
	  float r = (Y - sthresh) / swidth;
	  float w = (1 + erf(r)) / 2;
	  float s = Scotopic(X,Y,Z);
	  float ws = s*(1.0f - w);
	  X = w*X + 0.95047f * ws;
	  Y = w*Y +            ws;
	  Z = w*Z + 1.08883f * ws;
	}

	/* convert to output colour-space */
	(*XYZ_RGB)(X,Y,Z,&R,&G,&B);

	imgpix[0] = CLIP(255.0f*R,lim);
	imgpix[1] = CLIP(255.0f*G,lim);
	imgpix[2] = CLIP(255.0f*B,lim);
      }

      /* save to the image stream */
      for( k = 0; k < color_bytes; k++) {

	if( bit_depth == 16 ) {
	  /* confirmed by Imagemagick */
	  int n = imgpix[k];
	  image[m+1]= (png_byte) (n / 256);
	  image[m] = (png_byte) (n % 256);
	}
	else { /* if ( bit_depth == 8 ) { */
	  image[m] = (png_byte) imgpix[k];
	}
	m = m + bytes_per_pixel;
      }

    }
  }

  /* intensity table */
  if( verb ) {
    FILE *itt;
    if( (itt = fopen("itt.dat","w")) ) {
      for(i = 0; i <= lim; i++) {
	flux = fscale((float)i/(float)lim,ctype,f0,zero);
	iflux = CLIP(255.0*flux,lim);
	fprintf(itt,"%d %d\n",i,iflux);
      }
      fclose(itt);
    }
  }

  if( verb )
    fprintf(stderr,"Finishing intensity transformation..\n");

  /* --------------------------------------------------------------------*/
  /* Part: Save to PNG */

  if( png )
    fp = fopen(png, "wb");
  else
    fp = stdout;

  if (!fp) {
    fprintf(stderr,"Crash: Initialising of an output file failed.\n");
    goto finish;
  }

  png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING,NULL,NULL,NULL);
  if (!png_ptr) {
    fclose(fp);
    goto finish;
  }

  info_ptr = png_create_info_struct(png_ptr);
  if (!info_ptr) {
    fclose(fp);
    png_destroy_write_struct(&png_ptr,(png_infopp)NULL);
    goto finish;
  }

  png_init_io(png_ptr, fp);
  png_set_write_status_fn(png_ptr, NULL);

  png_set_IHDR(png_ptr, info_ptr, width, height, bit_depth,
	       color_bytes == 1 ? PNG_COLOR_TYPE_GRAY : PNG_COLOR_TYPE_RGB,
	       PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE,
	       PNG_FILTER_TYPE_BASE);


  /*
  png_set_gAMA(png_ptr, info_ptr, gamma);
  */

  /*
  png_color_16 back;
  back.gray = 0;
  png_set_bKGD(png_ptr, info_ptr, &back);

  png_set_pHYs(png_ptr, info_ptr, res_x, res_y,unit_type);
  png_set_sCAL(png_ptr, info_ptr, unit, width, height)
  */

  if( fval[0] )
    stdvalues[0] = strdup(fval[0]);
  else
    stdvalues[0] = strdup("");

  if( fval[1] )
    stdvalues[1] = strdup(fval[1]);
  else
    stdvalues[1] = strdup("");

  /* numerical constant in declaraton of buf must be greater than max. length
     of sum of following string(s) */
  strcpy(buf,"An image");
  if( fval[0] )
    sprintf(buf+strlen(buf)," of the %s",fval[0]);
  if( fval[6] )
    sprintf(buf+strlen(buf)," taken at %s observatory",fval[6]);
  if( fval[1] )
    sprintf(buf+strlen(buf)," by %s",fval[1]);
  if( fval[4] )
    sprintf(buf+strlen(buf)," by the %s instrument",fval[4]);
  if( fval[10] )
    sprintf(buf+strlen(buf)," of the %s telescope",fval[10]);
  if( fval[3] )
    sprintf(buf+strlen(buf)," at %s UT (start time)",fval[3]);
  if( fval[5] )
    sprintf(buf+strlen(buf)," of exposure %s sec",fval[5]);
  if( fval[2] && color_bytes == 1 )
    sprintf(buf+strlen(buf)," with the %s filter",fval[2]);
  strcat(buf,".");
  if( fval[7] )
    sprintf(buf+strlen(buf)," The instrument temperature: %s.",fval[7]);
  if( fval[8] )
    sprintf(buf+strlen(buf)," XBinnig: %s.",fval[8]);
  if( fval[9] )
    sprintf(buf+strlen(buf)," YBinnig: %s.",fval[9]);

  stdvalues[2] = strdup(buf);
  stdvalues[3] = strdup("");

  /* decode time (round fractional seconds and adds timezone),
     Standard FITS headers must contains DATE-OBS as
     YYYY-MM-DDTHH:MM:SSS.SSS, however the rule is violated
     by many software writing obsolete YYYY-MM-DD keyword.
  */
  i = 0;
  if( fval[3] ) {
    for( i = 0, c = fval[3], c0 = fval[3]; *c != '\0'; c++) {
      if( *c == '-' || *c == 'T' || *c == ':' || *c == ' ' || *(c+1) == '\0') {
	n = c - c0;
	if( *(c+1) == '\0' )
	  n++;
	tm[i++] = strndup(c0,n);
	c0 = c + 1;
      }
    }
    /* This loop is my own bugy implementation of bugy strtok
       (https://sourceware.org/bugzilla/show_bug.cgi?id=16640) */
  }
  if( i == 6 ) {
    if( sscanf(tm[5],"%f",&r) == 1 )
      i = rint(r);
    else
      i = 0;
    sprintf(buf,"%s-%s-%s %s:%s:%02d GMT",tm[0],tm[1],tm[2],tm[3],tm[4],i);
    stdvalues[4] = strdup(buf);
  } else if( i == 3 ) {
    sprintf(buf,"%s-%s-%s",tm[0],tm[1],tm[2]);
    stdvalues[4] = strdup(buf);
  }
  else
    stdvalues[4] = strdup("");

  stdvalues[5] = strdup("Created by FITSPNG.");
  stdvalues[6] = strdup("");
  stdvalues[7] = strdup("");

  strcpy(buf,"");
  if( fval[4] )
    strcat(buf,fval[4]);
  if( fval[10] ) {
    strcat(buf,", ");
    strcat(buf,fval[10]);
  }
  stdvalues[8] = strdup(buf);

  strcpy(buf,"Converted from the original FITS image:");
  sprintf(buf+strlen(buf)," %s",fitsname);
  stdvalues[9] = strdup(buf);

  for(i = 0; i < NSTDKEYS; i++ ) {

    text_ptr[i].key = stdkeys[i];
    text_ptr[i].text = stdvalues[i];
    text_ptr[i].compression = PNG_TEXT_COMPRESSION_NONE;

  }
  png_set_text(png_ptr, info_ptr, text_ptr, NSTDKEYS);

  png_write_info(png_ptr, info_ptr);

#ifdef WORDS_BIGENDIAN
  ;
#else
  if( bit_depth == 16 )
    png_set_swap(png_ptr);
#endif

  if( (row_pointers = malloc(height*sizeof(row_pointers))) == NULL ) {
    fprintf(stderr,"There is no room for all rows of image.\n");
    png_destroy_write_struct(&png_ptr,(png_infopp)NULL);
    goto finish;
  }
  for (i = 0; i < height; i++)
    row_pointers[i] = image + i*width*bytes_per_pixel*color_bytes;
  png_write_image(png_ptr,row_pointers);

  png_write_end(png_ptr, NULL);

  png_free(png_ptr,row_pointers);

  png_destroy_write_struct(&png_ptr, &info_ptr);
  fclose(fp);

  if( verb )
    fprintf(stderr,"Finishing save of PNG.\n");

 finish:

  fits_report_error(stderr, status);

  for( i = 0; i < NFKEYS; i++)
    free(fval[i]);
  for( i = 0; i < NSTDKEYS; i++)
    free(stdvalues[i]);
  for( i = 0; i < 6; i++ )
    free(tm[i]);

  free(image);
  free(pic);

  return(status);
}