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

#include <unistd.h>

#include "psFits.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(psImage* output, // a psImage to recycle.
                         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);
        psFree(output);
        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);
        psFree(output);
        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);
        psFree(output);
        return NULL;
    }

    /* Validate the number of axis */
    if ((nAxis < 2) || (nAxis > 3)) {
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psFits_IMAGE_DIMENSION_UNSUPPORTED,
                nAxis);
        psFree(output);
        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);
        psFree(output);
        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);
        psFree(output);
        return NULL;
    }

    output = psImageRecycle(output,
                            lastPixel[0]-firstPixel[0]+1,
                            lastPixel[1]-firstPixel[1]+1,
                            datatype);

    if (output == NULL) {
        psError(PS_ERR_UNKNOWN, false,
                "Failed to allocate a properly sized image.");
        return false;
    }

    // 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,
                      psMetadata* header,
                      const psImage* input,
                      int numZPlanes)
{

    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;
    }
    int numCols = input->numCols;
    int numRows = input->numRows;

    int status = 0;

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

    int naxis = 3;
    long naxes[3];

    naxes[0] = numCols;
    naxes[1] = numRows;
    naxes[2] = numZPlanes;

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

    fits_create_img(fits->fd, bitPix, naxis, naxes, &status);

    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 (header != NULL) {
        psFitsWriteHeader(header, (psPtr)fits);
    }

    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,
                       psRegion region,
                       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] = 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;

    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,
                "Specified region [%d:%d,%d:%d], is not valid given the %dx%d FITS image.",
                region.y0,region.y1-1,region.x0,region.x1-1);
        return false;

    }

    int dx = lastPixel[0] - firstPixel[0];
    int dy = lastPixel[1] - firstPixel[1];
    if (dx > numCols ||
            dy > numRows) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                "The region [%d:%d,%d:%d], is not valid given the input %dx%d image.",
                firstPixel[1]-1,lastPixel[1]-1,
                firstPixel[0]-1,lastPixel[0]-1,
                numCols, numRows);
        return false;
    }

    psImage* subset;
    if (dx != numCols || dy != numRows) {
        // the input image needs to be subsetted
        subset = psImageSubset((psImage*)input, psRegionSet(0,dx+1,0,dy+1));
    } else {
        subset = psMemIncrRefCounter((psImage*)input);
    }

    fits_write_subset(fits->fd, dataType, firstPixel, lastPixel, subset->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;
}

