#if HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <pslib.h>

#include "pmFPA.h"
#include "pmHDUUtils.h"
#include "pmFPAMaskWeight.h"
#include "pmFlatField.h"


bool pmFlatField(pmReadout *in, const pmReadout *flat)
{
    PS_ASSERT_PTR_NON_NULL(in, false);
    PS_ASSERT_PTR_NON_NULL(in->image, false);
    PS_ASSERT_IMAGE_NON_EMPTY(in->image, false);
    PS_ASSERT_PTR_NON_NULL(flat, false);
    PS_ASSERT_PTR_NON_NULL(flat->image, false);
    PS_ASSERT_IMAGE_NON_EMPTY(flat->image, false);
    if (in->mask) {
        PS_ASSERT_IMAGE_TYPE(in->mask, PS_TYPE_MASK, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(in->mask, in->image, false);
    }
    PS_ASSERT_IMAGE_TYPE(flat->image, in->image->type.type, false);
    if (flat->mask) {
        PS_ASSERT_IMAGE_TYPE(flat->mask, PS_TYPE_MASK, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(flat->mask, flat->image, false);
    }

    psImage *inImage   = in->image;     // Input image
    psImage *inMask    = in->mask;      // Mask for input image
    psImage *flatImage = flat->image;   // Flat-field image
    psImage *flatMask  = flat->mask;    // Mask for flat-field image

    // Add flat-field MD5 to header
    pmHDU *hdu = pmHDUFromReadout(in);  // HDU of interest
    psVector *md5 = psImageMD5(flat->image); // md5 hash
    psString md5string = psMD5toString(md5); // String
    psFree(md5);
    psStringPrepend(&md5string, "FLAT image MD5: ");
    psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                     md5string, "");
    psFree(md5string);

    // Check input image is not larger than flat image; mask is the same size as the input
    if (inImage->numRows > flatImage->numRows || inImage->numCols > flatImage->numCols) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, "Input image (%dx%d) is larger than flat-field image "
                "(%dx%d).\n", inImage->numCols, inImage->numRows, flatImage->numCols, flatImage->numRows);
        return false;
    }

    // Offsets on the chip
    int x0in = psMetadataLookupS32(NULL, in->parent->concepts, "CELL.X0");
    int y0in = psMetadataLookupS32(NULL, in->parent->concepts, "CELL.Y0");
    int x0flat = psMetadataLookupS32(NULL, flat->parent->concepts, "CELL.X0");
    int y0flat = psMetadataLookupS32(NULL, flat->parent->concepts, "CELL.Y0");

    // Determine offset based on image offset with chip offset: input frame to flat frame
    int yOffset = in->row0 + y0in - flat->row0 - y0flat;
    int xOffset = in->col0 + x0in - flat->col0 - x0flat;

    // Check that offsets are within image limits
    if (inImage->numRows + yOffset > flatImage->numRows ||
            inImage->numCols + xOffset > flatImage->numCols) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true, "Input image (%dx%d) with offsets (%d,%d) is larger than "
                "flat-field image (%dx%d).\n", inImage->numCols, inImage->numRows, xOffset, yOffset,
                flatImage->numCols, flatImage->numRows);
        return false;
    }

    // Macro for all PS types
    #define FLAT_DIVISION_CASE(TYPE, SPECIAL) \
case PS_TYPE_##TYPE: \
    for (int j = 0; j < inImage->numRows; j++) { \
        for (int i = 0; i < inImage->numCols; i++) { \
            ps##TYPE flatValue = flatImage->data.TYPE[j + yOffset][i + xOffset]; \
            if (flatValue <= 0.0 || (flatMask && flatMask->data.U8[j + yOffset][i + xOffset])) { \
                if (inMask) { \
                    inMask->data.U8[j][i] |= PM_MASK_FLAT; \
                } \
                inImage->data.TYPE[j][i] = SPECIAL; \
            } else { \
                inImage->data.TYPE[j][i] /= flatValue; \
            } \
        } \
    } \
    break;

    switch (inImage->type.type) {
        FLAT_DIVISION_CASE(U8,  0);
        FLAT_DIVISION_CASE(U16, 0);
        FLAT_DIVISION_CASE(U32, 0);
        FLAT_DIVISION_CASE(U64, 0);
        FLAT_DIVISION_CASE(S8,  0);
        FLAT_DIVISION_CASE(S16, 0);
        FLAT_DIVISION_CASE(S32, 0);
        FLAT_DIVISION_CASE(S64, 0);
        FLAT_DIVISION_CASE(F32, NAN);
        FLAT_DIVISION_CASE(F64, NAN);
    default:
        psError(PS_ERR_BAD_PARAMETER_TYPE, true, "Unsupported type for input image: %x\n",
                inImage->type.type);
        return false;
    }

    psTime *time = psTimeGetNow(PS_TIME_TAI); // The time now, used for reporting
    psString timeString = psTimeToISO(time); // String with time
    psFree(time);
    psStringPrepend(&timeString, "Flat-field processing completed at ");
    psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                     timeString, "");
    psFree(timeString);


    return true;
}
