// Program to fix bug in FITS images written with psLib prior to trunk@r25002.  For these images, funpack
// won't put a single compressed image back in the PHU because they are missing ZSIMPLE.  This program simply
// reads the FITS file and writes it back out the way it came in, but since psLib is fixed, the output should
// behave well with funpack.
#include <stdio.h>
#include <string.h>
#include <pslib.h>

int main(int argc, char *argv[])
{
    if (argc != 3) {
        fprintf(stderr, "Insufficient arguments.  Please supply input FITS file and output FITS file.");
        return PS_EXIT_CONFIG_ERROR;
    }
    const char *inName = argv[1];
    const char *outName = argv[2];

    if (strcmp(inName, outName) == 0) {
        fprintf(stderr, "Can't overwrite input image.");
        return PS_EXIT_CONFIG_ERROR;
    }

    psFits *inFits = psFitsOpen(inName, "r");
    psFits *outFits = psFitsOpen(outName, "w");
    outFits->options = psFitsOptionsAlloc();

    int numExt = psFitsGetSize(inFits);
    for (int i = 0; i < numExt; i++) {
        if (!psFitsMoveExtNum(inFits, i, false)) {
            psErrorStackPrint(stderr, "Can't move to extension %d\n", i);
            return PS_EXIT_SYS_ERROR;
        }

        if (i == 0 && psFitsCheckCompressedImagePHU(inFits, NULL)) {
            // We're now positioned on the image
            i++;
        }

        psMetadata *header = psFitsReadHeader(NULL, inFits);
        psString extname = psFitsGetExtName(inFits);
        psFitsType type = psFitsGetExtType(inFits);
        switch (type) {
          case PS_FITS_TYPE_IMAGE: {
              int naxis = psMetadataLookupS32(NULL, header, "NAXIS");
              if (naxis == 0) {
                  if (!psFitsWriteBlank(outFits, header, extname)) {
                      psErrorStackPrint(stderr, "Can't write blank extension %d\n", i);
                      return PS_EXIT_SYS_ERROR;
                  }
              } else {
                  psString compType = psMetadataLookupStr(NULL, header, "ZCMPTYPE");
                  if (compType) {
                      psVector *tile = psVectorAlloc(3, PS_TYPE_S32);
                      psVectorInit(tile, 0);
                      for (int i = 0; i < naxis; i++) {
                          psString name = NULL;
                          psStringAppend(&name, "ZTILE%d", i + 1);
                          tile->data.S32[i] = psMetadataLookupS32(NULL, header, name);
                          psFree(name);
                      }
                      psFitsCompressionType type = psFitsCompressionTypeFromString(compType);
                      psFitsSetCompression(outFits, type, tile, 8, 0, 0);
                      psFree(tile);
                  } else {
                      psFitsSetCompression(outFits, PS_FITS_COMPRESS_NONE, NULL, 8, 0, 0);
                  }

                  outFits->options->bitpix = psMetadataLookupS32(NULL, header, "BITPIX");
                  outFits->options->bscale = psMetadataLookupF64(NULL, header, "BSCALE");
                  outFits->options->bzero = psMetadataLookupF64(NULL, header, "BZERO");
                  if (isfinite(outFits->options->bscale) && isfinite(outFits->options->bzero)) {
                      outFits->options->scaling = PS_FITS_SCALE_MANUAL;
                  } else {
                      outFits->options->scaling = PS_FITS_SCALE_NONE;
                  }
                  outFits->options->fuzz = false;

                  psArray *cube = psFitsReadImageCube(inFits, psRegionSet(0,0,0,0));
                  if (!cube) {
                      psErrorStackPrint(stderr, "Can't read image(s) from extension %d\n", i);
                      return PS_EXIT_SYS_ERROR;
                  }
                  if (!psFitsWriteImageCube(outFits, header, cube, extname)) {
                      psErrorStackPrint(stderr, "Can't write image(s) to extension %d\n", i);
                      return PS_EXIT_SYS_ERROR;
                  }
                  psFree(cube);
              }
              break;
          }
          case PS_FITS_TYPE_BINARY_TABLE:
          case PS_FITS_TYPE_ASCII_TABLE: {
              psArray *table = psFitsReadTable(inFits);
              if (!table) {
                  psErrorStackPrint(stderr, "Can't read table from extension %d\n", i);
                  return PS_EXIT_SYS_ERROR;
              }
              if (!psFitsWriteTable(outFits, header, table, extname)) {
                  psErrorStackPrint(stderr, "Can't write table to extension %d\n", i);
                  return PS_EXIT_SYS_ERROR;
              }
              psFree(table);
              break;
          }
          default:
            psAbort("Unrecognised FITS extension type: %x\n", type);
        }

        psFree(header);
        psFree(extname);
    }

    psFitsClose(inFits);
    psFitsClose(outFits);

    return PS_EXIT_SUCCESS;
}
