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

#include <unistd.h>

#include "psFits.h"
#include "psFitsImage.h"
#include "string.h"
#include "psError.h"
#include "psErrorText.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


static bool convertPsTypeToFits(int type, int* bitPix, double* bZero, int* dataType)
{

    int bitpix;
    int datatype;
    double bzero = 0.0;

    switch (type) {

    case PS_TYPE_U8:
        bitpix = BYTE_IMG;
        datatype = TBYTE;
        break;

    case PS_TYPE_S8:
        bitpix = BYTE_IMG;
        bzero = INT8_MIN;
        datatype = TSBYTE;
        break;

    case PS_TYPE_U16:
        bitpix = SHORT_IMG;
        bzero = -1.0 * INT16_MIN;
        datatype = TUSHORT;
        break;

    case PS_TYPE_S16:
        bitpix = SHORT_IMG;
        datatype = TSHORT;
        break;

    case PS_TYPE_U32:
        bitpix = LONG_IMG;
        bzero = -1.0 * INT32_MIN;
        datatype = TUINT;
        break;

    case PS_TYPE_S32:
        bitpix = LONG_IMG;
        datatype = TINT;
        break;

    case PS_TYPE_F32:
        bitpix = FLOAT_IMG;
        datatype = TFLOAT;
        break;

    case PS_TYPE_F64:
        bitpix = DOUBLE_IMG;
        datatype = TDOUBLE;
        break;

    case PS_DATA_STRING:
        bitpix = BYTE_IMG;
        datatype = TSTRING;
        break;

    case PS_DATA_BOOL:
        bitpix = BYTE_IMG;
        datatype = TLOGICAL;
        break;

    default: {
            char* typeStr;
            PS_TYPE_NAME(typeStr,type);
            psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                    PS_ERRORTEXT_psFits_TYPE_UNSUPPORTED,
                    typeStr);
            return false;
        }
    }

    // pass the requested parameters  (NULL parameters are not set, of course).
    if (bitPix != NULL) {
        *bitPix = bitpix;
    }

    if (dataType != NULL) {
        *dataType = datatype;
    }

    if (bZero != NULL) {
        *bZero = bzero;
    }

    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
{
    psS32 status = 0;           /* CFITSIO file vars */
    psS32 nAxis = 0;
    psS32 anynull = 0;
    psS32 bitPix = 0;           /* Pixel type */
    long nAxes[3];
    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 */
    char fitsErr[80] = "";      /* CFITSIO error message string */
    psS32 fitsDatatype = 0;
    psS32 datatype = 0;

    if (fits == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psFits_NULL);
        return NULL;
    }

    // check to see if we even 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,
                PS_ERRORTEXT_psFits_GET_HDU_TYPE_FAILED,
                fitsErr);
        return NULL;
    }
    if (hdutype != IMAGE_HDU) {
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_NOT_IMAGE_TYPE);
        return NULL;
    }

    /* Get the data type 'bitPix' from the FITS image */
    if (fits_get_img_equivtype(fits->fd, &bitPix, &status) != 0) {
        fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_DATATYPE_UNKNOWN,
                fitsErr);
        return NULL;
    }

    /* Get the dimensions 'nAxis' from the FITS image */
    if (fits_get_img_dim(fits->fd, &nAxis, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_IMAGE_DIM_UNKNOWN,
                fitsErr);
        return NULL;
    }

    /* Validate the number of axis */
    if ((nAxis < 2) || (nAxis > 3)) {
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_IMAGE_DIMENSION_UNSUPPORTED,
                nAxis);
        return NULL;
    }

    /* Get the Image size from the FITS file */
    if (fits_get_img_size(fits->fd, nAxis, nAxes, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_IMAGE_SIZE_UNKNOWN,
                fitsErr);
        return NULL;
    }

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

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

    increment[0] = 1;
    increment[1] = 1;
    increment[2] = 1;

    switch (bitPix) {
    case BYTE_IMG:
        datatype = PS_TYPE_U8;
        fitsDatatype = TBYTE;
        break;
    case SBYTE_IMG:
        datatype = PS_TYPE_S8;
        fitsDatatype = TSBYTE;
        break;
    case USHORT_IMG:
        datatype = PS_TYPE_U16;
        fitsDatatype = TUSHORT;
        break;
    case SHORT_IMG:
        datatype = PS_TYPE_S16;
        fitsDatatype = TSHORT;
        break;
    case ULONG_IMG:
        datatype = PS_TYPE_U32;
        fitsDatatype = TUINT;
        break;
    case LONG_IMG:
        datatype = PS_TYPE_S32;
        fitsDatatype = TINT;
        break;
    case LONGLONG_IMG:
        datatype = PS_TYPE_S64;
        fitsDatatype = TLONGLONG;
        break;
    case FLOAT_IMG:
        datatype = PS_TYPE_F32;
        fitsDatatype = TFLOAT;
        break;
    case DOUBLE_IMG:
        datatype = PS_TYPE_F64;
        fitsDatatype = TDOUBLE;
        break;
    default:
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_FITS_TYPE_UNSUPPORTED,
                bitPix);
        return NULL;
    }

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

    // n.b., this assumes contiguous image buffer
    if (fits_read_subset(fits->fd, fitsDatatype, firstPixel, lastPixel, increment,
                         NULL, output->data.V[0], &anynull, &status) != 0) {
        psFree(output);
        (void)fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_READ_FAILED,
                fitsErr);
        return NULL;
    }

    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, PS_ERRORTEXT_psFits_NULL);
        return false;
    }

    if (!input) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true, PS_ERRORTEXT_psFits_IMAGE_NULL);
        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 (! convertPsTypeToFits(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,
                PS_ERRORTEXT_psFits_WRITE_FAILED,
                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,
                PS_ERRORTEXT_psFits_NULL);
        return false;
    }

    if (input == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psFits_IMAGE_NULL);
        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,
                PS_ERRORTEXT_psFits_GET_HDU_TYPE_FAILED,
                fitsErr);
        return NULL;
    }
    if (hdutype != IMAGE_HDU) {
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_NOT_IMAGE_TYPE);
        return NULL;
    }

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

    // determine the FITS-equivalent parameters
    int bitPix;
    double bZero;
    int dataType;
    if (! convertPsTypeToFits(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,
                PS_ERRORTEXT_psFits_IMAGE_UPDATE_TYPE_MISMATCH,
                fitsTypeStr, imageTypeStr);
        return false;
    }

    //check if the HDU has the z-plane requested
    if (z >= nAxes[2]) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true,
                PS_ERRORTEXT_psFits_FITS_Z_SMALL,
                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 %ix%i 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,
                PS_ERRORTEXT_psFits_WRITE_FAILED,
                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,
                PS_ERRORTEXT_psFits_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,
                PS_ERRORTEXT_psFits_IMAGE_DIM_UNKNOWN,
                fitsErr);
        return NULL;
    }

    if (nAxis == 2) {
        psArray *images = psArrayAlloc(1); // Single image plane
        images->n = 1;
        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,
                    PS_ERRORTEXT_psFits_IMAGE_SIZE_UNKNOWN,
                    fitsErr);
            return NULL;
        }

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

        return images;
    }

    // Bad dimensionality
    psError(PS_ERR_IO, true, PS_ERRORTEXT_psFits_IMAGE_DIMENSION_UNSUPPORTED, 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,
                PS_ERRORTEXT_psFits_NULL);
        return false;
    }

    if (input == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psFits_IMAGE_NULL);
        return false;
    }

    if (input->n == 0) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, PS_ERRORTEXT_psFits_ARRAY_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, PS_ERRORTEXT_psFits_ARRAY_SIZE_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, PS_ERRORTEXT_psFits_METADATA_ADD_FAILED,
                "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, PS_ERRORTEXT_psFits_WRITE_PLANE_FAILED, 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, PS_ERRORTEXT_psFits_WRITE_PLANE_FAILED, 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,
                PS_ERRORTEXT_psFits_NULL);
        return false;
    }

    if (input == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psFits_IMAGE_NULL);
        return false;
    }

    if (input->n == 0) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, PS_ERRORTEXT_psFits_ARRAY_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, PS_ERRORTEXT_psFits_UPDATE_PLANE_FAILED, i);
            return false;
        }
    }

    return true;
}

