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

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

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

# 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); }

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

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

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

  // 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);

  // 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
  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 {
    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
  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 is Noptions set to the correct value?
  // XXX get ZMASKCMP 
  // XXX can primary be defined on basis of the header?

  if (primary) {
    // XXX SIMPLE & XTENSION need to be set on the first line
    if (!gfits_scan (header, "ZSIMPLE", "%t", 1, &header->simple)) ESCAPE;
    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);
  } else {
    // XXX XTENSION needs to be set on the first line
    if (!gfits_scan (header, "ZTENSION", "%s", 1, exttype)) ESCAPE;
    gfits_delete (header, "ZTENSION", 1);
    gfits_modify_extended (header, exttype, "Image extension");

    MOD_KEYWORD ("ZPCOUNT",  "PCOUNT",   "%d", &header->pcount, header->pcount);
    MOD_KEYWORD ("ZGCOUNT",  "GCOUNT",   "%d", &header->gcount, header->gcount);
  }

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

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

  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;

  // 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];
  }

  // max_tile_size is in pixels
  bytes_per_pixel = abs(header[0].bitpix) / 8;
  ALLOCATE (out, char, bytes_per_pixel*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);

    switch (zdef.format) {
      // gfits_varlength returns a pointer to a chunk of Nzdata of data elements starting at zdata.
      // uncompress_data uncompresses from zdata to the temporary output buffer which must be allocated

      // 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.

      case 'B': {
	char *zdata;
	zdata = gfits_varlength_column_pointer (ftable, &zdef, i, &Nzdata);
	if (!zdata) return (FALSE);
	if (!gfits_byteswap_zdata ((char *)zdata, Nzdata, 8)) return (FALSE);
	if (!gfits_uncompress_data ((char *)zdata, Nzdata, cmptype, optname, optvalue, Noptions, out, &Nout, bytes_per_pixel)) return (FALSE);
	if (!gfits_distribute_data (matrix, 8, out, Nout, otile, ztile, zscale, zzero)) return (FALSE);
 	// copy the uncompressed pixels into their correct locations 	    
	break;
      }

      case 'I': {
	short *zdata;
	zdata = gfits_varlength_column_pointer (ftable, &zdef, i, &Nzdata);
	if (!zdata) return (FALSE);
	if (!gfits_byteswap_zdata ((char *)zdata, Nzdata, 16)) return (FALSE);
	if (!gfits_uncompress_data ((char *)zdata, Nzdata, cmptype, optname, optvalue, Noptions, out, &Nout, bytes_per_pixel)) return (FALSE);
	if (!gfits_distribute_data (matrix, 16, out, Nout, otile, ztile, zscale, zzero)) return (FALSE);
 	// copy the uncompressed pixels into their correct locations 	    
	break;
      }

      case 'J': {
	int *zdata;
	if (!zdata) return (FALSE);
	if (!gfits_byteswap_zdata ((char *)zdata, Nzdata, 32)) return (FALSE);
	if (!gfits_uncompress_data ((char *)zdata, Nzdata, cmptype, optname, optvalue, Noptions, out, &Nout, bytes_per_pixel)) return (FALSE);
	if (!gfits_distribute_data (matrix, 32, out, Nout, otile, ztile, zscale, zzero)) return (FALSE);
 	// copy the uncompressed pixels into their correct locations 	    
	break;
      }

      default:
	fprintf (stderr, "invalid size for compressed data: %c\n", zdef.format);  
	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);
}

int gfits_tile_size (Matrix *matrix, int *otile, int *ztile) {

  int i, Npixels, Ndimen;

  // true sizes of this tile (in pixels)
  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 *ztile, float zscale, float zzero) {

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

  // counter for current row in tile to copy
  ALLOCATE (counter, int, matrix->Naxes);
  memset (counter, 0, matrix->Naxes*sizeof(int));

  // true sizes of this tile (in pixels)
  ALLOCATE (Ztile, int, matrix->Naxes);
  for (i = 0; i < matrix->Naxes; i++) {
    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++) { \
      *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 8:
	SETUP_INSIZE (char, 1);
	break;
      case 16: 
	SETUP_INSIZE (short, 2);
	break; 
      case 32:
	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 = 1; j < matrix->Naxes; j++) {
      offset = offset*matrix->Naxis[j] + counter[j];
    }      
  }

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

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

# ifdef BYTE_SWAP

  int i;
  char tmp;

  switch (bitpix) {
    case +8:
    case -8:
      break;

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

    case +32:
    case -32:
      for (i = 0; i < 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 +64:
    case -64:
      for (i = 0; i < 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_extname, has_zimage, zimage;
    char extname[80];

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

    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);
}
