/** @file  psImageIO.c
 *
 *  @brief Contains image I/O routines.
 *
 *  This file defines the file input/output functions for the psImage structure.
 *
 *  @author Robert DeSonia, MHPCC
 *
 *  @version $Revision: 1.15 $ $Name:  $
 *  @date $Date: 2004/11/04 01:05:00 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 */

#include <fitsio.h>
#include <unistd.h>

#include "psImageIO.h"
#include "psError.h"
#include "psMemory.h"

#include "psImageErrors.h"

psImage* psImageReadSection(psImage* output,
                            psS32 col,
                            psS32 row,
                            psS32 numCols,
                            psS32 numRows,
                            psS32 z,
                            char *extname,
                            psS32 extnum,
                            char *filename)
{
    fitsfile *fptr = NULL;      /* Pointer to the FITS file */
    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 hduType = IMAGE_HDU;
    psS32 fitsDatatype = 0;
    psS32 datatype = 0;

    if (filename == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psImageIO_FILENAME_NULL);
        psFree(output);
        return NULL;
    }

    /* Open the FITS file */
    (void)fits_open_file(&fptr, filename, READONLY, &status);
    if (fptr == NULL || status != 0) {
        fits_get_errstatus(status, fitsErr);
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                PS_ERRORTEXT_psImageIO_FILENAME_INVALID,
                filename, fitsErr);
        psFree(output);
        return NULL;
    }

    /* find the specified extension */
    if (extname != NULL) {
        if (fits_movnam_hdu(fptr, hduType, extname, 0, &status) != 0) {
            fits_get_errstatus(status, fitsErr);
            status = 0;
            (void)fits_close_file(fptr, &status);
            psError(PS_ERR_LOCATION_INVALID, true,
                    PS_ERRORTEXT_psImageIO_EXTNAME_INVALID,
                    extname, filename, fitsErr);
            psFree(output);
            return NULL;
        }
    } else {
        if (fits_movabs_hdu(fptr, extnum + 1, &hduType, &status) != 0) {
            fits_get_errstatus(status, fitsErr);
            status = 0;
            (void)fits_close_file(fptr, &status);
            psError(PS_ERR_LOCATION_INVALID, true,
                    PS_ERRORTEXT_psImageIO_EXTNUM_INVALID,
                    extnum, filename, fitsErr);
            psFree(output);
            return NULL;
        }
    }

    /* Get the data type 'bitPix' from the FITS image */
    if (fits_get_img_equivtype(fptr, &bitPix, &status) != 0) {
        fits_get_errstatus(status, fitsErr);
        status = 0;
        (void)fits_close_file(fptr, &status);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psImageIO_DATATYPE_UNKNOWN,
                filename, fitsErr);
        psFree(output);
        return NULL;
    }

    /* Get the dimensions 'nAxis' from the FITS image */
    if (fits_get_img_dim(fptr, &nAxis, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        status = 0;
        (void)fits_close_file(fptr, &status);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psImageIO_IMAGE_DIM_UNKNOWN,
                filename, fitsErr);
        psFree(output);
        return NULL;
    }

    /* Validate the number of axis */
    if ((nAxis < 2) || (nAxis > 3)) {
        status = 0;
        (void)fits_close_file(fptr, &status);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psImageIO_IMAGE_DIMENSION_UNSUPPORTED,
                nAxis);
        psFree(output);
        return NULL;
    }

    /* Get the Image size from the FITS file */
    if (fits_get_img_size(fptr, nAxis, nAxes, &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        status = 0;
        (void)fits_close_file(fptr, &status);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psImageIO_IMAGE_SIZE_UNKNOWN,
                filename, fitsErr);
        psFree(output);
        return NULL;
    }

    if (numCols < 1) {
        numCols += nAxes[0] - col;
    }
    if (numRows < 1) {
        numRows += nAxes[1] - row;
    }

    firstPixel[0] = col + 1;
    firstPixel[1] = row + 1;
    firstPixel[2] = z + 1;

    lastPixel[0] = firstPixel[0] + numCols - 1;
    lastPixel[1] = firstPixel[1] + numRows - 1;
    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_psImageIO_FITS_TYPE_UNSUPPORTED,
                bitPix, filename);
        psFree(output);
        return NULL;
    }
    output = psImageRecycle(output, numCols, numRows, datatype);
    if (fits_read_subset(fptr, fitsDatatype, firstPixel, lastPixel, increment,
                         NULL, output->data.V[0], &anynull, &status) != 0) {
        psFree(output);
        (void)fits_get_errstatus(status, fitsErr);
        status = 0;
        (void)fits_close_file(fptr, &status);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psImageIO_READ_FAILED,
                filename, fitsErr);
        return NULL;
    }

    (void)fits_close_file(fptr, &status);

    return output;
}

psBool psImageWriteSection(psImage* input,
                           psS32 col0,
                           psS32 row0,
                           psS32 z,
                           char *extname,
                           psS32 extnum,
                           char *filename)
{
    psS32 numCols = 0;
    psS32 numRows = 0;

    psS32 status = 0;             /* CFITSIO status */
    fitsfile *fptr = NULL;      /* pointer to the FITS file */
    long nAxes[3];              /* Image axis vars */
    long firstPixel[3];         /* First Pixel to read */
    long lastPixel[3];          /* Last Pixel to read */
    char fitsErr[80];           /* FITSIO message string */
    psS32 datatype = 0;           /* the datatype of the image */
    psS32 bitPix = 0;             /* FITS bitPix value */
    psS32 hduType = IMAGE_HDU;    /* the HDU type (image,table, etc.) */
    double bscale = 1.0;
    double bzero = 0.0;
    psBool createNewHDU = false;

    /* need a valid image to write */
    if (input == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psImage_IMAGE_NULL);
        return false;
    }

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

    switch (input->type.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.0f * INT16_MIN;
        datatype = TUSHORT;
        break;
    case PS_TYPE_S16:
        bitPix = SHORT_IMG;
        datatype = TSHORT;
        break;
    case PS_TYPE_U32:
        bitPix = LONG_IMG;
        bzero = -1.0f * 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;
    default: {
            char* typeStr;
            PS_TYPE_NAME(typeStr,input->type.type);
            psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                    PS_ERRORTEXT_psImageIO_TYPE_UNSUPPORTED,
                    typeStr);
            return false;
        }
    }

    /* Open the FITS file */
    if (access(filename, F_OK) == 0) {     // file
        // exists
        (void)fits_open_file(&fptr, filename, READWRITE, &status);
        if (fptr == NULL || status != 0) {
            fits_get_errstatus(status, fitsErr);
            psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                    PS_ERRORTEXT_psImageIO_FILENAME_INVALID,
                    filename, fitsErr);
            return false;
        }

        /* find the specified extension */
        if (extname != NULL) {
            if (fits_movnam_hdu(fptr, hduType, extname, 0, &status) != 0) {
                fits_get_errstatus(status, fitsErr);
                status = 0;
                (void)fits_close_file(fptr, &status);
                psError(PS_ERR_LOCATION_INVALID, true,
                        PS_ERRORTEXT_psImageIO_EXTNAME_INVALID,
                        extname, filename, fitsErr);
                return false;
            }
        } else {
            psS32 numHDUs = 0;

            fits_get_num_hdus(fptr, &numHDUs, &status);
            if (numHDUs < extnum) {
                status = 0;
                (void)fits_close_file(fptr, &status);
                psError(PS_ERR_LOCATION_INVALID, true,
                        PS_ERRORTEXT_psImageIO_WRITE_EXTNUM_INVALID,
                        extnum, numHDUs);
                return false;
            } else if (numHDUs == extnum) {
                createNewHDU = true;
            } else if (fits_movabs_hdu(fptr, extnum + 1, &hduType, &status) != 0) {
                fits_get_errstatus(status, fitsErr);
                status = 0;
                (void)fits_close_file(fptr, &status);
                psError(PS_ERR_LOCATION_INVALID, true,
                        PS_ERRORTEXT_psImageIO_EXTNUM_INVALID,
                        extnum, filename, fitsErr);
                return false;
            }
        }

    } else {
        // file does not exist

        (void)fits_create_file(&fptr, filename, &status);
        if (fptr == NULL || status != 0) {
            fits_get_errstatus(status, fitsErr);
            psError(PS_ERR_IO, true,
                    PS_ERRORTEXT_psImageIO_FILENAME_CREATE_FAILED,
                    filename, fitsErr);
            return false;
        }
        createNewHDU = true;
    }

    if (createNewHDU) {
        /* create the mandatory image keywords */
        nAxes[0] = col0 + numCols;
        nAxes[1] = row0 + numRows;
        nAxes[2] = z + 1;
        if (fits_create_img(fptr, bitPix, 3, nAxes, &status) != 0) {
            (void)fits_get_errstatus(status, fitsErr);
            status = 0;
            (void)fits_close_file(fptr, &status);
            psError(PS_ERR_IO, true,
                    PS_ERRORTEXT_psImageIO_CREATE_HDU_FAILED,
                    filename, fitsErr);
            return false;
        }
        // set the bscale/bzero
        fits_write_key_dbl(fptr, "BZERO", bzero, 12, "Pixel Value Offset", &status);
        fits_write_key_dbl(fptr, "BSCALE", bscale, 12, "Pixel Value Scale", &status);
        fits_set_bscale(fptr, bscale, bzero, &status);

        if (extname != NULL) {
            /* create the extension for the Primary HDU */
            if (fits_write_key_str(fptr, "EXTNAME", extname, "Extension Name", &status) != 0) {
                (void)fits_get_errstatus(status, fitsErr);
                status = 0;
                (void)fits_close_file(fptr, &status);
                psError(PS_ERR_IO, true,
                        PS_ERRORTEXT_psImageIO_CREATE_EXTENSION_FAILED,
                        filename, fitsErr);
                return false;
            }
        }
    }

    firstPixel[0] = col0 + 1;
    firstPixel[1] = row0 + 1;
    firstPixel[2] = z + 1;

    lastPixel[0] = firstPixel[0] + numCols - 1;
    lastPixel[1] = firstPixel[1] + numRows - 1;
    lastPixel[2] = z + 1;

    if (fits_write_subset(fptr, datatype, firstPixel, lastPixel, input->data.V[0], &status) != 0) {
        (void)fits_get_errstatus(status, fitsErr);
        status = 0;
        (void)fits_close_file(fptr, &status);
        psError(PS_ERR_IO, true,
                PS_ERRORTEXT_psImageIO_WRITE_FAILED,
                filename, fitsErr);
        return false;
    }

    status = 0;
    (void)fits_close_file(fptr, &status);

    return true;
}