# include <ohana.h>
# include <gfitsio.h>
# include <zlib.h>

# define ESCAPE { \
  if (ztile != NULL) free (ztile); \
  if (optname != NULL) { \
    for (j = 0; j < Noptions; j++) { \
      free (optname[j]); \
      free (optvalue[j]); \
    } \
    free (optname); \
    free (optvalue); \
  } \
  if (out != NULL) free (out); \
  if (otile != NULL) free (otile); \
  if (ntile != NULL) free (ntile); \
  return (FALSE); }

# define MOD_KEYWORD(ZNAME,NAME,TYPE,IN,OUT) { \
    if (gfits_scan (header, ZNAME, TYPE, 1, IN)) { \
      gfits_modify (header, NAME, TYPE, 1, OUT); \
    } \
    gfits_delete (header, ZNAME, 1); }

# define MOD_KEYWORD_REQUIRED(ZNAME,NAME,TYPE,IN,OUT) { \
  if (!gfits_scan (header, ZNAME, TYPE, 1, IN)) ESCAPE; \
  gfits_delete (header, ZNAME, 1); \
  gfits_modify (header, NAME, TYPE, 1, OUT); }

int gfits_uncompress_image (Header *header, Matrix *matrix, FTable *ftable) {

  int i, j, status, zimage, zcol, Nzrows, Nout, Nzdata, max_tile_size;
  char cmptype[80];
  char zaxis[10], naxis[10], key[10], word[81], exttype[81], checksum[81], datasum[81];
  int Noptions, NOPTIONS, zblank, oblank;
  VarLengthColumn zdef;
  float zscale, zzero;

  int zdata_pixsize, odata_pixsize, idata_pixsize;

  int *ztile = NULL;
  int *otile = NULL;
  int *ntile = NULL;
  char **optname = NULL;
  char **optvalue = NULL;
  char *out = NULL;
  char *zdata = NULL;

  Noptions = 0;

  // is ZIMAGE present?
  status = gfits_scan (ftable->header, "ZIMAGE", "%t", 1, &zimage);
  if (!status || !zimage) ESCAPE;

  // copy original header to output header
  gfits_copy_header (ftable->header, header);

  // delete ZIMAGE from output header
  gfits_delete (header, "ZIMAGE", 1);

  // extract compression-specific keywords, update header as needed.
  if (!gfits_scan (header, "ZCMPTYPE", "%s", 1, cmptype)) ESCAPE;
  gfits_delete (header, "ZCMPTYPE", 1);

  MOD_KEYWORD_REQUIRED ("ZBITPIX", "BITPIX", "%d", &header->bitpix, header->bitpix);
  MOD_KEYWORD_REQUIRED ("ZNAXIS",  "NAXIS",  "%d", &header->Naxes,  header->Naxes);

  for (i = 0; i < header->Naxes; i++) {
    snprintf (zaxis, 10, "ZNAXIS%d", i + 1);
    snprintf (naxis, 10, "NAXIS%d", i + 1);
    MOD_KEYWORD_REQUIRED (zaxis,  naxis,  "%d", &header->Naxis[i],  header->Naxis[i]);
  }    

  // set up the tile sizes : ztile is the default size of the tile in the Nth dimension
  // the actual tile size may be smaller at the edge of a dimension.  if the ZTILEn
  // entries are not found, default to [Nx,1,1,...]
  ALLOCATE (ztile, int, header->Naxes);
  if (!gfits_scan (header, "ZTILE1", "%d", 1, &ztile[0])) {
    ztile[0] = header->Naxis[0];
    for (i = 1; i < header->Naxes; i++) {
      ztile[i] = 1;
    }
  } else {
    gfits_delete (header, "ZTILE1", 1);
    for (i = 1; i < header->Naxes; i++) {
      snprintf (key, 10, "ZTILE%d", i + 1);
      if (!gfits_scan (header, key, "%d", 1, &ztile[i])) ESCAPE;
      gfits_delete (header, key, 1);
    }
  }

  // search for algorithm-specific keywords. these are used to control compression options
  // note the difference between the keywords (1 indexed) and the variables (0 indexed)
  NOPTIONS = 10;
  ALLOCATE (optname, char *, NOPTIONS);
  ALLOCATE (optvalue, char *, NOPTIONS);
  for (Noptions = 0; TRUE; Noptions++) {
    snprintf (key, 10, "ZNAME%d", Noptions + 1);
    if (!gfits_scan (header, key, "%s", 1, word)) break;
    gfits_delete (header, key, 1);
    optname[Noptions] = strcreate (word);

    snprintf (key, 10, "ZVAL%d", Noptions + 1);
    if (!gfits_scan (header, key, "%s", 1, word)) ESCAPE;
    gfits_delete (header, key, 1);
    optvalue[Noptions] = strcreate (word);

    if (Noptions == NOPTIONS - 1) {
      NOPTIONS += 10;
      REALLOCATE (optname, char *, NOPTIONS);
      REALLOCATE (optvalue, char *, NOPTIONS);
    }
  }

  // XXX get ZMASKCMP 
  int zsimple, have_zsimple;
  int have_ztension;

  have_zsimple  = gfits_scan (header, "ZSIMPLE", "%t", 1, &zsimple);
  have_ztension = gfits_scan (header, "ZTENSION", "%s", 1, exttype);

  // this is a very bogus case: we cannot have both keywords
  if (have_zsimple && have_ztension) ESCAPE;

  // if neither are present, we have an image that is not really following the standard:
  // assume it is a PHU
  if (!have_zsimple && !have_ztension) {
    header->simple = TRUE;
    gfits_extended_to_primary (header, header->simple, "Image data");

    MOD_KEYWORD ("ZEXTEND",  "EXTEND",   "%t", &header->extend, header->extend);
    MOD_KEYWORD ("ZBLOCKED", "BLOCKED",  "%t", &header->extend, header->extend);
  }

  // have_zsimple : image comes from a PHU
  if (have_zsimple) {
    header->simple = zsimple;
    gfits_delete (header, "ZSIMPLE", 1);
    gfits_extended_to_primary (header, header->simple, "Image data");

    MOD_KEYWORD ("ZEXTEND",  "EXTEND",   "%t", &header->extend, header->extend);
    MOD_KEYWORD ("ZBLOCKED", "BLOCKED",  "%t", &header->extend, header->extend);
  } 

  // have_ztension : image comes from an extension
  if (have_ztension) {
    gfits_delete (header, "ZTENSION", 1);
    gfits_modify_extended (header, exttype, "Image extension");

    // we may have an uncompressed PCOUNT / GCOUNT value, otherwise set to 0,1
    if (gfits_scan (header, "ZPCOUNT", "%d", 1, &header->pcount)) {
	gfits_delete (header, "ZPCOUNT", 1);
	gfits_modify (header, "PCOUNT", "%d", 1, header->pcount);
    } else {
	header->pcount = 0;
	gfits_modify (header, "PCOUNT", "%d", 1, header->pcount);
    }
    if (gfits_scan (header, "ZGCOUNT", "%d", 1, &header->gcount)) {
	gfits_delete (header, "ZGCOUNT", 1);
	gfits_modify (header, "GCOUNT", "%d", 1, header->gcount);
    } else {
	header->pcount = 1;
	gfits_modify (header, "GCOUNT", "%d", 1, header->gcount);
    }
  } else {
    header->pcount = 0;
    header->gcount = 1;
    gfits_modify (header, "PCOUNT", "%d", 1, header->pcount);
    gfits_modify (header, "GCOUNT", "%d", 1, header->gcount);
  }

  MOD_KEYWORD ("ZHECKSUM", "CHECKSUM", "%s", checksum,        checksum);
  MOD_KEYWORD ("ZDATASUM", "DATASUM",  "%s", datasum,         datasum);

  zscale = 1;
  gfits_scan (header, "ZSCALE", "%lf", 1, &zscale);

  zblank = 32767;
  gfits_scan (header, "ZBLANK", "%d", 1, &zblank);

  oblank = 32767;
  gfits_scan (header, "BLANK", "%d", 1, &oblank);

  zzero = 0;
  gfits_scan (header, "ZZERO", "%lf", 1, &zzero);

  // find the COMPRESSED_DATA column (format should be 1PB, 1PI, 1PJ)
  for (i = 1; TRUE; i++) {
    snprintf (key, 10, "TTYPE%d", i);
    if (!gfits_scan (ftable->header, key, "%s", 1, word)) ESCAPE;
    if (!strcmp (word, "COMPRESSED_DATA")) break;
  }
  zcol = i;

  if (!gfits_varlength_column_define (ftable, &zdef, zcol)) ESCAPE;
  gfits_delete (header, "TFIELDS", 1);
  snprintf (key, 10, "TTYPE%d", zcol);
  gfits_delete (header, key, 1);
  snprintf (key, 10, "TFORM%d", zcol);
  gfits_delete (header, key, 1);

  // create the output image
  gfits_create_matrix (header, matrix);

  // counters for tile numbers
  ALLOCATE (otile, int, matrix->Naxes);
  ALLOCATE (ntile, int, matrix->Naxes);
  max_tile_size = 1;
  for (i = 0; i < matrix->Naxes; i++) {
    otile[i] = 0;
    ntile[i] = (matrix->Naxis[i] % ztile[i]) ? (matrix->Naxis[i] / ztile[i] + 1) : (matrix->Naxis[i] / ztile[i]);
    max_tile_size *= ztile[i];
  }

  // this takes place in three major steps:
  // 1) read the table data : zdef.format tells the size of the varlength heap element
  // 2) uncompress the data : pixel size depends on compression method
  // 3) distribute data to the tiles : output pixel size depends on header.bitpix

  // heapdata -> zdata -> odata -> idata

  // heapdata.pixsize : zdef.format
  // zdata.pixsize : depends on compression
  // odata.pixsize : depends on compression
  // idata.pixsize : header.bitpix

  // size of an element in the vartable heap section
  zdata_pixsize = gfits_vartable_heap_pixsize (zdef.format);

  // size of a pixel in the final image
  idata_pixsize = abs(header[0].bitpix) / 8;

  // size of a pixel in the output from the decompression routine
  odata_pixsize = gfits_uncompressed_data_pixsize (cmptype, header[0].bitpix, optname, optvalue, Noptions);
  ALLOCATE (out, char, odata_pixsize*max_tile_size);

  // uncompress the data
  Nzrows = ftable->header->Naxis[1];
  for (i = 0; i < Nzrows; i++) {

    // expected output size for this tile
    Nout = gfits_tile_size (matrix, otile, ztile);

    zdata = gfits_varlength_column_pointer (ftable, &zdef, i, &Nzdata);
    if (!zdata) return (FALSE);

    // XXX not certain this is the correct place to do the swap (or if zdata_pixsize is
    // the correct size to guide the swap)

    if (strcasecmp(cmptype, "GZIP_1")) {
      if (!gfits_byteswap_zdata (zdata, Nzdata, zdata_pixsize)) return (FALSE);
    }

    // gfits_uncompress_data uncompresses from zdata to the temporary output buffer which must be allocated
    if (!gfits_uncompress_data ((char *)zdata, Nzdata, cmptype, optname, optvalue, Noptions, out, &Nout, odata_pixsize)) return (FALSE);

    // copy the uncompressed pixels into their correct locations 	    
    if (!gfits_distribute_data (matrix, odata_pixsize, out, Nout, otile, oblank, ztile, zblank, zscale, zzero)) return (FALSE);

    // update the tile counters, carrying to the next dimension if needed
    for (j = 0; j < matrix->Naxes; j++) {
      otile[j] ++;
      if (otile[j] == ntile[j]) {
	otile[j] = 0;
      } else {
	break;
      }
    }
  }
  return (TRUE);
}

// gfits_varlength returns a pointer to a chunk of Nzdata of data elements starting at zdata.

// XXX need to have an API to send the user data

// XXX need to byte-swap the table column; this needs to be worked out more clearly
// in the APIs: do the table read / column extract functions swap or not?.  we have
// put the byte-swapping in gifts_varlength_column_pointer, but this is fairly weak.

// XXX inconsistency between data element size between compressed data, uncompressed data, 
// and output pixel data.....

// true sizes of this tile (in pixels)
int gfits_tile_size (Matrix *matrix, int *otile, int *ztile) {

  int i, Npixels, Ndimen;

  Npixels = 1;
  for (i = 0; i < matrix->Naxes; i++) {
    Ndimen = MIN ((matrix->Naxis[i] - otile[i]*ztile[i]), ztile[i]);
    Npixels *= Ndimen;
  }
  
  return (Npixels);
}

// bitpix is the input data size/type
int gfits_distribute_data (Matrix *matrix, int bitpix, char *data, int Ndata, int *otile, int oblank, int *ztile, int zblank, float zscale, float zzero) {

  int i, j, start, offset, coord, Nline;
  int *counter = NULL;
  int *Ztile = NULL;

  ALLOCATE (counter, int, matrix->Naxes);
  ALLOCATE (Ztile, int, matrix->Naxes);

  // counter for current row in tile to copy
  // true sizes of this tile (in pixels)
  for (i = 0; i < matrix->Naxes; i++) {
    counter[i] = 0;
    Ztile[i] = MIN ((matrix->Naxis[i] - otile[i]*ztile[i]), ztile[i]);
  }

  // number of lines in the tile (in pixels)
  Nline = 1;
  for (i = 1; i < matrix->Naxes; i++) {
    Nline *= Ztile[i];
  }

  // double check reported size (needs the Isize)
  assert (Ndata == Ztile[0]*Nline);

  // set the starting point of the tile:
  // start = otile[0]*ztile[0] + otile[1]*ztile[1]*Naxis[0] + otile[2]*ztile[2]*Naxis[0]*Naxis[1] + ...;
  // start = otile[0]*ztile[0] + Naxis[0]*(otile[1]*ztile[1] + Naxis[1]*(otile[2]*ztile[2] + ...));
  start = otile[matrix->Naxes-1]*ztile[matrix->Naxes-1];
  for (i = matrix->Naxes - 2; i >= 0; i--) {
    coord = otile[i]*ztile[i];
    start = start*matrix->Naxis[i] + coord;
  }
  
  // pixel offset in output array relative to tile start
  offset = 0;

  // we need to set up switches for all of the possible combinations:
  // this macro is used at the inner switch to run the actual loop
# define SCALE_AND_DIST(TYPE, SIZE) { \
    TYPE *Optr = (TYPE *) &matrix->buffer[SIZE*(offset + start)]; \
    for (j = 0; j < Ztile[0]; j++, Iptr++, Optr++) { \
      if (*Iptr == zblank) { \
	  *Optr = oblank;    \
      } else { \
	  *Optr = *Iptr * zscale + zzero;	\
      } } }

  // this macro sets up the outer switch and calls above macro with all output bitpix options
# define SETUP_INSIZE(TYPE, SIZE) { \
    TYPE *Iptr = (TYPE *) &data[i*SIZE*Ztile[0]]; \
    switch (matrix->bitpix) { \
      case 8: \
	SCALE_AND_DIST (char, 1); \
	break; \
      case 16: \
	SCALE_AND_DIST (short, 2); \
	break; \
      case 32: \
	SCALE_AND_DIST (int, 4); \
	break; \
      case -32: \
	SCALE_AND_DIST (float, 4); \
	break; \
      case -64: \
	SCALE_AND_DIST (double, 8); \
	break; \
      default: \
	abort(); \
    } }

  // loop over lines in the tile
  for (i = 0; i < Nline; i++) {
    switch (bitpix) {
      case 1:
	SETUP_INSIZE (char, 1);
	break;
      case 2: 
	SETUP_INSIZE (short, 2);
	break; 
      case 4:
	SETUP_INSIZE (int, 4);
	break;
      default:
	abort();
    }

    // update the counters, carrying to the next dimension if needed
    for (j = 1; j < matrix->Naxes; j++) {
      counter[j] ++;
      if (counter[j] == Ztile[j]) {
	counter[j] = 0;
      } else {
	break;
      }
    }
    if (j == matrix->Naxes) assert (i == Nline - 1); // we should be done here...

    // Naxes = 3
    // offset = counter[1]*matrix->Naxis[0] + counter[2]*matrix->Naxis[0]*matrix->Naxis[1] + 
    // offset = matrix->Naxis[0]*(counter[1] + matrix->Naxis[1]*(counter[2] + matrix->Naxis[2]*...))

    // determine the offset of the next line relative to the start position
    offset = counter[matrix->Naxes - 1];
    for (j = matrix->Naxes - 2; j >= 0; j--) {
      offset = offset*matrix->Naxis[j] + counter[j];
    }      
  }

  free (counter);
  free (Ztile);
  return (TRUE);
}

int gfits_byteswap_zdata (char *zdata, int Nzdata, int pixsize) {

# ifdef BYTE_SWAP

  int i;
  char tmp;

  switch (pixsize) {
    case 1:
      break;

    case 2:
      for (i = 0; i < 2*Nzdata; i+=2) {
	tmp = zdata[i];
	zdata[i] = zdata[i+1];
	zdata[i+1] = tmp;
      }
      break;

    case 4:
      for (i = 0; i < 4*Nzdata; i+=4) {
	tmp = zdata[i+1];
	zdata[i+1] = zdata[i+2];
	zdata[i+2] = tmp;
	tmp = zdata[i];
	zdata[i] = zdata[i+3];
	zdata[i+3] = tmp;
      }
      break;

    case 8:
      for (i = 0; i < 8*Nzdata; i+=8) {
	tmp = zdata[i+0];
	zdata[i+0] = zdata[i+7];
	zdata[i+7] = tmp;
	tmp = zdata[i+1];
	zdata[i+1] = zdata[i+6];
	zdata[i+6] = tmp;
	tmp = zdata[i+2];
	zdata[i+2] = zdata[i+5];
	zdata[i+5] = tmp;
	tmp = zdata[i+3];
	zdata[i+3] = zdata[i+4];
	zdata[i+4] = tmp;
      }
      break;
  }
# endif
  return (TRUE);
}

int gfits_extension_is_compressed (Header *header) {

    int has_extension, has_extname, has_zimage, zimage;
    char extname[80], extension[80];

    has_extension = gfits_scan (header, "XTENSION", "%s", 1, extension);
    has_extname = gfits_scan (header, "EXTNAME", "%s", 1, extname);
    has_zimage  = gfits_scan (header, "ZIMAGE", "%t", 1, &zimage);

    if (has_extension && !strcmp (extension, "IMAGE")) return (FALSE);
    if (has_zimage && zimage) return (TRUE);
    if (has_extname) {
	if (!strcmp (extname, "COMPRESSED_IMAGE")) return (TRUE);
    }

    return (FALSE);
}

int gfits_compressed_is_primary (Header *header) {

    int has_ztension, has_zimage;
    int ztension, zimage;

    has_zimage   = gfits_scan (header, "ZIMAGE",   "%t", 1, &zimage);
    has_ztension = gfits_scan (header, "ZTENSION", "%t", 1, &ztension);

    if (has_zimage && zimage) return (TRUE);

    return (FALSE);
}

int gfits_uncompressed_data_pixsize (char *cmptype, int out_bitpix, char **optname, char **optvalue, int Noptions) {

  int i, Nbyte;

  if (!strcasecmp(cmptype, "GZIP_1")) {
    if (out_bitpix == 8)  return (1);
    if (out_bitpix == 16) return (2);
    if (out_bitpix == 32) return (4);
    if (out_bitpix == -32) return (4);
    if (out_bitpix == -64) return (8);
    return (1);
  }
  if (!strcasecmp(cmptype, "RICE_1")) {

    // if BYTEPIX option is specified, use that for Nbyte
    for (i = 0; i < Noptions; i++) {
      if (!strcmp(optname[i], "BYTEPIX")) {
	Nbyte = atoi (optvalue[i]);
	return (Nbyte);
      }
    }

    return (4);
  }
  if (!strcasecmp(cmptype, "PLIO_1")) {
    return (4);
  }
  if (!strcasecmp(cmptype, "HCOMPRESS_1")) {
    if (out_bitpix == 8)  return (4);
    if (out_bitpix == 16) return (4);
    return (8);
  }
  return (0);
}

int gfits_vartable_heap_pixsize (char format) {

  if (format == 'B') {
      return (1);
  }
  if (format == 'I') {
      return (2);
  }
  if (format == 'J') {
      return (4);
  }
  fprintf (stderr, "invalid size for compressed data: %c\n", format);  
  abort ();
}
