/** @file  psFitsImage.c
 *
 *  @brief Contains Fits I/O routines
 *
 *  @ingroup FileIO
 *
 *  @author Robert DeSonia, MHPCC
 *
 *  @version $Revision: 1.16 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-10-24 22:52:55 $
 *
 *  Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
 */

#include <unistd.h>
#include <assert.h>
#include <string.h>

#include "psFits.h"
#include "psAssert.h"
#include "psFitsImage.h"
#include "psError.h"

#include "psImageStructManip.h"
#include "psMemory.h"
#include "psString.h"
#include "psLogMsg.h"
#include "psTrace.h"
#include "psVector.h"

#include "psFitsHeader.h"

#define MAX_STRING_LENGTH 256  // maximum length string for FITS routines

// Information required to read a FITS file
typedef struct
{
    int nAxis;                          // Number of axes
    int bitPix;                         // Bits per pixel
    long nAxes[3];                      // Length of each axis
    long firstPixel[3];                 // lower-left corner of image subset
    long lastPixel[3];                  // upper-right corner of image subset
    long increment[3];                  // increment for image subset
    int fitsDatatype;                   // cfitsio data type
    int psDatatype;                     // psLib data type
}
p_psFitsReadInfo;

static p_psFitsReadInfo *p_psFitsReadInfoAlloc(const psFits *fits, // The FITS file handle
        psRegion region, // Region to read
        int z // z-plane to read in cube
                                              )
{
    PS_ASSERT_PTR_NON_NULL(fits, NULL);
    PS_ASSERT_INT_NONNEGATIVE(z, NULL);

    p_psFitsReadInfo *info = psAlloc(sizeof(p_psFitsReadInfo));
    memset(info, 0, sizeof(p_psFitsReadInfo));

    int status = 0;                     // CFITSIO status
    char fitsErr[80];                   // CFITSIO error message string

    // check to see if we even are positioned on an image HDU
    int hduType;                        // Type of HDU
    if (fits_get_hdu_type(fits->fd, &hduType, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not determine the HDU type. CFITSIO Error: %s"),
                fitsErr);
        goto bad;
    }
    if (hduType != IMAGE_HDU) {
        psError(PS_ERR_IO, true,
                _("Current FITS HDU type must be an image."));
        goto bad;
    }

    // Get the data type 'bitPix' from the FITS image
    if (fits_get_img_equivtype(fits->fd, &info->bitPix, &status) != 0) {
        fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not determine image data type. CFITSIO Error: %s"),
                fitsErr);
        goto bad;
    }

    /* Get the dimensions 'nAxis' from the FITS image */
    if (fits_get_img_dim(fits->fd, &info->nAxis, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not determine image dimensions. CFITSIO Error: %s"),
                fitsErr);
        goto bad;
    }

    /* Validate the number of axis */
    if ((info->nAxis < 2) || (info->nAxis > 3)) {
        psError(PS_ERR_IO, true,
                _("Image number of dimensions, %d, is not valid.  Only two or three dimensions supported for FITS I/O."),
                info->nAxis);
        goto bad;
    }

    /* Get the Image size from the FITS file */
    if (fits_get_img_size(fits->fd, info->nAxis, info->nAxes, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not determine image size. CFITSIO Error: %s"),
                fitsErr);
        goto bad;
    }

    info->firstPixel[0] = region.x0 + 1;
    info->firstPixel[1] = region.y0 + 1;
    info->firstPixel[2] = z + 1;

    if (region.x1 > 0) {
        info->lastPixel[0] = region.x1;
    } else {
        info->lastPixel[0] = info->nAxes[0] + region.x1; // n.b., region.x1 < 0
    }
    if (region.y1 > 0) {
        info->lastPixel[1] = region.y1;
    } else {
        info->lastPixel[1] = info->nAxes[1] + region.y1; // n.b., region.y1 < 0
    }
    info->lastPixel[2] = z + 1;

    info->increment[0] = 1;
    info->increment[1] = 1;
    info->increment[2] = 1;

    switch (info->bitPix) {
    case BYTE_IMG:
        info->psDatatype = PS_TYPE_U8;
        info->fitsDatatype = TBYTE;
        break;
    case SBYTE_IMG:
        info->psDatatype = PS_TYPE_S8;
        info->fitsDatatype = TSBYTE;
        break;
    case USHORT_IMG:
        info->psDatatype = PS_TYPE_U16;
        info->fitsDatatype = TUSHORT;
        break;
    case SHORT_IMG:
        info->psDatatype = PS_TYPE_S16;
        info->fitsDatatype = TSHORT;
        break;
    case ULONG_IMG:
        info->psDatatype = PS_TYPE_U32;
        info->fitsDatatype = TUINT;
        break;
    case LONG_IMG:
        info->psDatatype = PS_TYPE_S32;
        info->fitsDatatype = TINT;
        break;
    case LONGLONG_IMG:
        info->psDatatype = PS_TYPE_S64;
        info->fitsDatatype = TLONGLONG;
        break;
    case FLOAT_IMG:
        info->psDatatype = PS_TYPE_F32;
        info->fitsDatatype = TFLOAT;
        break;
    case DOUBLE_IMG:
        info->psDatatype = PS_TYPE_F64;
        info->fitsDatatype = TDOUBLE;
        break;
    default:
        psError(PS_ERR_IO, true,
                _("FITS image type, BITPIX=%d, is not supported."),
                info->bitPix);
        goto bad;
    }

    return info;

    // Common error cleanup
bad:
    psFree(info);
    return NULL;
}

// Read into an extant image of just the right size
static bool fitsReadImage(psImage *output,   // Output image
                          const psFits *fits, // FITS file handle
                          p_psFitsReadInfo *info // Info on how to read
                         )
{
    // n.b., this assumes contiguous image buffer
    assert(output);
    assert(fits);
    assert(info);
    assert(output->numCols == info->lastPixel[0] - info->firstPixel[0] + 1);
    assert(output->numRows == info->lastPixel[1] - info->firstPixel[1] + 1); // Right size
    assert(!output->parent);            // No parents means the buffer is contiguous

    int anynull = 0;                    // Are there any NULLs in the data?
    int status = 0;                     // cfitsio status
    if (fits_read_subset(fits->fd, info->fitsDatatype, info->firstPixel, info->lastPixel,
                         info->increment, NULL, output->data.V[0], &anynull, &status) != 0) {
        char fitsErr[80];               // CFITSIO error message string
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Reading FITS file failed. CFITSIO Error: %s"),
                fitsErr);
        return false;
    }

    return true;
}

psImage* psFitsReadImage(const psFits *fits, // the psFits object
                         psRegion region, // the region in the FITS image to read
                         int z          // the z-plane in the FITS image cube to read
                        )
{
    p_psFitsReadInfo *info = p_psFitsReadInfoAlloc(fits, region, z);

    psImage *output = psImageAlloc(info->lastPixel[0] - info->firstPixel[0] + 1,
                                   info->lastPixel[1] - info->firstPixel[1] + 1,
                                   info->psDatatype);

    if (!fitsReadImage(output, fits, info)) {
        psFree(info);
        psFree(output);
        return NULL;
    }

    psFree(info);
    return output;
}

psImage* psFitsReadImageBuffer(psImage *output, // Output image buffer
                               const psFits *fits,    // the psFits object
                               psRegion region, // the region in the FITS image to read
                               int z           // the z-plane in the FITS image cube to read
                              )
{
    if (output && output->parent) {
        psError(PS_ERR_IO, true, "Unable to read into a buffer for a child image.\n");
        return NULL;
    }

    p_psFitsReadInfo *info = p_psFitsReadInfoAlloc(fits, region, z);

    output = psImageRecycle(output, info->lastPixel[0] - info->firstPixel[0] + 1,
                            info->lastPixel[1] - info->firstPixel[1] + 1,
                            info->psDatatype);

    if (!fitsReadImage(output, fits, info)) {
        psFree(info);
        psFree(output);
        return NULL;
    }

    psFree(info);
    return output;
}

bool psFitsWriteImage(psFits* fits,
                      const psMetadata* header,
                      const psImage* input,
                      int numZPlanes,
                      const char* extname)
{
    // this is equivalent to insert after the last HDU

    psFitsMoveLast(fits);
    return psFitsInsertImage(fits,header,input,numZPlanes,extname,true);
}

bool psFitsInsertImage(psFits* fits,
                       const psMetadata* header,
                       const psImage* input,
                       int numZPlanes,
                       const char* extname,
                       bool after)
{

    if (!fits) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true, _("The input psFits object can not NULL."));
        return false;
    }

    if (!input) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true, _("The input psImage was NULL.  Need a non-NULL psImage for operation to be performed."));
        return false;
    }

    int numCols = input->numCols;       // Number of columns for image
    int numRows = input->numRows;       // Number of rows for image
    int status = 0;                     // Status from cfitsio

    // determine the FITS-equivalent parameters
    int bitPix;                         // Bits per pixel
    double bZero;                       // Zero offset
    int dataType;                       // cfitsio data type
    if (! p_psFitsTypeToCfitsio(input->type.type, &bitPix, &bZero, &dataType) ) {
        return false;
    }

    int naxis = 3;                      // Number of axes
    long naxes[3];                      // Length of each axis
    naxes[0] = numCols;
    naxes[1] = numRows;
    naxes[2] = numZPlanes;

    if (numZPlanes < 2) {
        naxis = 2;
    }

    // Create the image HDU
    int hdus = psFitsGetSize(fits);     // Number of HDUs in file
    if (hdus == 0) {
        // We're creating the first image
        fits_create_img(fits->fd, bitPix, naxis, naxes, &status);
    } else {
        if (!after) {
            if (psFitsGetExtNum(fits) == 0) {
                // We're creating a replacement primary HDU.
                // Set status to signal fits_insert_img to insert a new primary HDU
                status = PREPEND_PRIMARY;
            } else {
                // Move back one to perform an insert after the previous HDU
                psFitsMoveExtNum(fits, -1, true);
            }
        }
        // Insert after the current position
        fits_insert_img(fits->fd, bitPix, naxis, naxes, &status);
    }

    if (header && !psFitsWriteHeader(fits, header)) {
        psError(PS_ERR_IO, false, "Unable to write FITS header.\n");
        return false;
    }

    if (bZero != 0) {        // set the bscale/bzero
        fits_write_key_dbl(fits->fd, "BZERO", bZero, 12, "Pixel Value Offset", &status);
        fits_write_key_dbl(fits->fd, "BSCALE", 1.0, 12, "Pixel Value Scale", &status);
        fits_set_bscale(fits->fd, 1.0, bZero, &status);
    }

    // write the header, if any.
    if (extname && strlen(extname) > 0) {
        psFitsSetExtName(fits,extname);
    }

    if (input->parent == NULL) { // if no parent, assume that the image data is contiguous
        fits_write_img(fits->fd,
                       dataType,              // datatype
                       1,                     // writing to the first z-plane
                       numCols*numRows,       // number of elements to write, i.e., the whole image
                       input->data.V[0],      // the data
                       &status);
    } else { // image data may not be contiguous; write one row at a time
        int firstPixel = 1;
        for (int row = 0; row < numRows; row++) {
            fits_write_img(fits->fd,
                           dataType,          // datatype
                           firstPixel,
                           numCols,           // number of elements to write, i.e., one row's worth
                           input->data.V[row],// the raw row data
                           &status);
            firstPixel += numCols;  // move to next row
        }
    }

    if (status != 0) {
        char fitsErr[MAX_STRING_LENGTH];
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not write data to file. CFITSIO Error: %s"),
                fitsErr);
        return false;
    }

    return true;

}

bool psFitsUpdateImage(psFits* fits,
                       const psImage* input,
                       int x0,
                       int y0,
                       int z)
{
    int status = 0;

    if (fits == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                _("The input psFits object can not NULL."));
        return false;
    }

    if (input == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                _("The input psImage was NULL.  Need a non-NULL psImage for operation to be performed."));
        return false;
    }

    // check to see if we are positioned on an image HDU
    int hduType;
    if ( fits_get_hdu_type(fits->fd, &hduType, &status) != 0) {
        char fitsErr[MAX_STRING_LENGTH];
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not determine the HDU type. CFITSIO Error: %s"),
                fitsErr);
        return NULL;
    }
    if (hduType != IMAGE_HDU) {
        psError(PS_ERR_IO, true,
                _("Current FITS HDU type must be an image."));
        return NULL;
    }

    int numCols = input->numCols;
    int numRows = input->numRows;

    // determine the FITS-equivalent parameters
    int bitPix;
    double bZero;
    int dataType;
    if (! p_psFitsTypeToCfitsio(input->type.type, &bitPix, &bZero, &dataType) ) {
        return false;
    }

    //check to see if the HDU has the same datatype
    int fileBitpix;
    int naxis;
    long nAxes[3];
    nAxes[2] = 1;
    fits_get_img_param(fits->fd, 3, &fileBitpix, &naxis, nAxes, &status);

    //check to see if the HDU has the same datatype
    if (bitPix != fileBitpix) {
        char* fitsTypeStr;
        char* imageTypeStr;
        PS_TYPE_NAME(fitsTypeStr,fileBitpix);
        PS_TYPE_NAME(imageTypeStr,input->type.type);
        psError(PS_ERR_IO, true,
                _("Can not update a %s image given a %s image."),
                fitsTypeStr, imageTypeStr);
        return false;
    }

    //check if the HDU has the z-plane requested
    if (z >= nAxes[2]) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true,
                _("Current FITS HDU has %ld z-planes, but z-plane %d was specified."),
                nAxes[2], z);
        return false;
    }

    // determine the region in the FITS file domain
    long firstPixel[3];
    long lastPixel[3];

    firstPixel[0] = x0 + 1;
    firstPixel[1] = y0 + 1;
    firstPixel[2] = z + 1;

    lastPixel[0] = x0 + numCols;
    lastPixel[1] = y0 + numRows;
    lastPixel[2] = z + 1;

    if (firstPixel[0] < 1 || firstPixel[0] > nAxes[0] ||
            firstPixel[1] < 1 || firstPixel[1] > nAxes[1] ||
            lastPixel[0] < 1 || lastPixel[0] > nAxes[0] ||
            lastPixel[1] < 1 || lastPixel[1] > nAxes[1]) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                "Input image [size of %ix%i] at position (%i,%i) does not all lay in the %lix%li FITS image.",
                numCols, numRows,
                x0, y0,
                nAxes[0], nAxes[1]);
        return false;
    }

    fits_write_subset(fits->fd, dataType, firstPixel, lastPixel, input->data.V[0], &status);

    if ( status != 0) {
        char fitsErr[MAX_STRING_LENGTH];
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not write data to file. CFITSIO Error: %s"),
                fitsErr);
        return false;
    }

    return true;
}

psArray *psFitsReadImageCube(const psFits *fits, psRegion region)
{
    int nAxis = 0;                      // Number of axes
    long nAxes[3];                      // Number of pixels on each axis
    int status = 0;                     // cfitsio status value
    char fitsErr[80] = "";              // CFITSIO error message string

    if (fits == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                _("The input psFits object can not NULL."));
        return NULL;
    }

    // Some of this replicates what is in psFitsReadImage, so it's a little inefficient.  But it saves
    // code replication, and should be sufficient for our needs.

    if (fits_get_img_dim(fits->fd, &nAxis, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                _("Could not determine image dimensions. CFITSIO Error: %s"),
                fitsErr);
        return NULL;
    }

    if (nAxis == 2) {
        psArray *images = psArrayAlloc(1); // Single image plane
        images->data[0] = psFitsReadImage(fits, region, 0);
        return images;
    }
    if (nAxis == 3) {
        if (fits_get_img_size(fits->fd, nAxis, nAxes, &status) != 0) {
            (void)fits_get_errstatus(status, fitsErr);
            psError(PS_ERR_IO, true,
                    _("Could not determine image size. CFITSIO Error: %s"),
                    fitsErr);
            return NULL;
        }

        psArray *images = psArrayAlloc(nAxes[2]); // Array of image planes
        for (int i = 0; i < nAxes[2]; i++) {
            images->data[i] = psFitsReadImage(fits, region, i);
        }

        return images;
    }

    // Bad dimensionality
    psError(PS_ERR_IO, true, _("Image number of dimensions, %d, is not valid.  Only two or three dimensions supported for FITS I/O."), nAxis);
    return NULL;
}

bool psFitsWriteImageCube(psFits *fits, psMetadata *header, const psArray *input, const char *extname)
{
    if (fits == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                _("The input psFits object can not NULL."));
        return false;
    }

    if (input == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                _("The input psImage was NULL.  Need a non-NULL psImage for operation to be performed."));
        return false;
    }

    if (input->n == 0) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, _("The input array was empty."));
        return false;
    }

    if (input->n == 1) {
        // The problem reduces to one already solved
        return psFitsWriteImage(fits, header, input->data[0], 1, extname);
    }

    // Check that all images are of the same size
    psImage *testImage = input->data[0];// First image off the array
    int numCols = testImage->numCols;   // Number of columns
    int numRows = testImage->numRows;   // Number of rows
    for (int i = 1; i < input->n; i++) {
        testImage = input->data[i];
        if (testImage->numCols != numCols || testImage->numRows != numRows) {
            psError(PS_ERR_BAD_PARAMETER_SIZE, true, _("The sizes of images in the array differ."));
            return false;
        }
    }

    // Need to check the header to make sure NAXIS and NAXIS[1-3] are correct
    psMetadata *headerCopy = NULL;      // Copy of header
    if (header) {
        headerCopy = psMemIncrRefCounter(header);
    } else {
        headerCopy = psMetadataAlloc();
    }
    bool update = psMetadataAddS32(headerCopy, PS_LIST_HEAD, "NAXIS", PS_META_REPLACE, "Dimensionality", 3) &&
                  psMetadataAddS32(headerCopy, PS_LIST_HEAD, "NAXIS1", PS_META_REPLACE, "Number of columns", numCols) &&
                  psMetadataAddS32(headerCopy, PS_LIST_HEAD, "NAXIS2", PS_META_REPLACE, "Number of rows", numRows) &&
                  psMetadataAddS32(headerCopy, PS_LIST_HEAD, "NAXIS3", PS_META_REPLACE, "Number of image planes",
                                   input->n);
    if (! update) {
        psError(PS_ERR_UNKNOWN, false, _("Failed to add metadata item, %s."),
                "NAXIS, NAXIS1, NAXIS2, NAXIS3");
        psFree(headerCopy);
        return false;
    }

    // Now we can safely write the images out.
    // The first is an psFitsImageWrite to create the extension.
    // The next are psFitsImageUpdate to write into the extension.
    if (! psFitsWriteImage(fits, headerCopy, input->data[0], input->n, extname)) {
        psError(PS_ERR_UNKNOWN, false, _("Could not write image plane %d."), 0);
        psFree(headerCopy);
        return false;
    }
    psFree(headerCopy);                 // Free, or drop reference

    for (int i = 1; i < input->n; i++) {
        if (! psFitsUpdateImage(fits, input->data[i], 0, 0, i)) {
            psError(PS_ERR_UNKNOWN, false, _("Could not write image plane %d."), i);
            return false;
        }
    }

    return true;
}

bool psFitsUpdateImageCube(psFits *fits, const psArray *input, int x0, int y0)
{
    if (fits == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                _("The input psFits object can not NULL."));
        return false;
    }

    if (input == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                _("The input psImage was NULL.  Need a non-NULL psImage for operation to be performed."));
        return false;
    }

    if (input->n == 0) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, _("The input array was empty."));
        return false;
    }

    for (int i = 0; i < input->n; i++) {
        if (! psFitsUpdateImage(fits, input->data[i], x0, y0, i)) {
            psError(PS_ERR_UNKNOWN, false, _("Could not update image plane %d."), i);
            return false;
        }
    }

    return true;
}

