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

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

#include "pmFPA.h"
#include "pmHDUUtils.h"
#include "pmBias.h"

static void overscanOptionsFree(pmOverscanOptions *options)
{
    psFree(options->stat);
    psFree(options->poly);
    psFree(options->spline);
}

pmOverscanOptions *pmOverscanOptionsAlloc(bool single, pmFit fitType, unsigned int order, psStats *stat)
{
    pmOverscanOptions *opts = psAlloc(sizeof(pmOverscanOptions));
    psMemSetDeallocator(opts, (psFreeFunc)overscanOptionsFree);

    // Inputs
    opts->single = single;
    opts->fitType = fitType;
    opts->order = order;
    opts->stat = psMemIncrRefCounter(stat);

    // Outputs
    opts->poly = NULL;
    opts->spline = NULL;

    return opts;
}


// subtractFrame(): this routine will take as input a readout for the input image and a readout for the bias
// image.  The bias image is subtracted in place from the input image.
static bool subtractFrame(pmReadout *in, // Input readout
                          const pmReadout *sub, // Readout to be subtracted from input
                          float scale   // Scale to apply before subtracting
                         )
{
    assert(in);
    assert(sub);

    psImage *inImage  = in->image;      // The input image
    psImage *inMask   = in->mask;       // The input mask
    psImage *subImage = sub->image;     // The image to be subtracted
    psImage *subMask  = sub->mask;      // The mask for the subtraction image

    // Check parities
    int xIpar = psMetadataLookupS32(NULL, in->parent->concepts, "CELL.XPARITY");
    int xSpar = psMetadataLookupS32(NULL, sub->parent->concepts, "CELL.XPARITY");
    if (xIpar != xSpar) {
        psError(PS_ERR_UNKNOWN, true, "images for subtraction do not have the same "
                "CELL.XPARITY (%d vs %d).\n    pmSubtractBias must be upgraded to handle this situation\n",
                xIpar, xSpar);
        return false;
    }

    int yIpar = psMetadataLookupS32(NULL, in->parent->concepts, "CELL.YPARITY");
    int ySpar = psMetadataLookupS32(NULL, sub->parent->concepts, "CELL.YPARITY");
    if (yIpar != ySpar) {
        psError(PS_ERR_UNKNOWN, true, "images for subtraction do not have the same "
                "CELL.YPARITY (%d vs %d).\n    pmSubtractBias must be upgraded to handle this situation\n",
                xIpar, xSpar);
        return false;
    }

    // Offsets of the cells
    int x0in = psMetadataLookupS32(NULL, in->parent->concepts, "CELL.X0");
    int y0in = psMetadataLookupS32(NULL, in->parent->concepts, "CELL.Y0");
    int x0sub = psMetadataLookupS32(NULL, sub->parent->concepts, "CELL.X0");
    int y0sub = psMetadataLookupS32(NULL, sub->parent->concepts, "CELL.Y0");

    if ((inImage->numCols + x0in - x0sub) > subImage->numCols) {
        psError(PS_ERR_UNKNOWN, true, "Image does not have enough columns for subtraction.\n");
        return false;
    }
    if ((inImage->numRows + y0in - y0sub) > subImage->numRows) {
        psError(PS_ERR_UNKNOWN, true, "Image does not have enough rows for subtraction.\n");
        return false;
    }

    if (scale == 1.0) {
        for (int i = 0; i < inImage->numRows; i++) {
            for (int j = 0; j < inImage->numCols; j++) {
                inImage->data.F32[i][j] -= subImage->data.F32[i+y0in-y0sub][j+x0in-x0sub];
                if (inMask && subMask) {
                    inMask->data.U8[i][j] |= subMask->data.U8[i+y0in-y0sub][j+x0in-x0sub];
                }
            }
        }
    } else {
        for (int i = 0; i < inImage->numRows; i++) {
            for (int j = 0; j < inImage->numCols; j++) {
                inImage->data.F32[i][j] -= subImage->data.F32[i+y0in-y0sub][j+x0in-x0sub] * scale;
                if (inMask && subMask) {
                    inMask->data.U8[i][j] |= subMask->data.U8[i+y0in-y0sub][j+x0in-x0sub];
                }
            }
        }
    }

    return true;
}

// Produce an overscan vector from an array of pixels
static psVector *overscanVector(float *chi2, // chi^2 from fit
                                pmOverscanOptions *overscanOpts, // Overscan options
                                const psArray *pixels, // Array of vectors containing the pixel values
                                psStats *myStats // Statistic to use in reducing the overscan
                               )
{
    assert(overscanOpts);
    assert(pixels);
    assert(myStats);

    psStatsOptions statistic = psStatsSingleOption(myStats->options); // Statistic to use
    assert(statistic != 0);

    // Reduce the overscans
    psVector *reduced = psVectorAlloc(pixels->n, PS_TYPE_F32); // Overscan for each row
    psVector *ordinate = psVectorAlloc(pixels->n, PS_TYPE_F32); // Ordinate
    psVector *mask = psVectorAlloc(pixels->n, PS_TYPE_U8); // Mask for fitting

    for (int i = 0; i < pixels->n; i++) {
        psVector *values = pixels->data[i]; // Vector with overscan values
        if (values->n > 0) {
            mask->data.U8[i] = 0;
            ordinate->data.F32[i] = 2.0*(float)i/(float)pixels->n - 1.0; // Scale to [-1,1]
            psVectorStats(myStats, values, NULL, NULL, 0);
            reduced->data.F32[i] = psStatsGetValue(myStats, statistic);
        } else if (overscanOpts->fitType == PM_FIT_NONE) {
            psError(PS_ERR_UNKNOWN, true, "The overscan is not supplied for all points on the "
                    "image, and no fit is requested.\n");
            return NULL;
        } else {
            // We'll fit this one out
            mask->data.U8[i] = 1;
        }
    }

    // Fit the overscan, if required
    psVector *fitted;                   // Fitted overscan values
    switch (overscanOpts->fitType) {
    case PM_FIT_NONE:
        // No fitting --- that's easy.
        fitted = psMemIncrRefCounter(reduced);
        break;
    case PM_FIT_POLY_ORD:
        if (overscanOpts->poly && (overscanOpts->poly->nX != overscanOpts->order ||
                                   overscanOpts->poly->type != PS_POLYNOMIAL_ORD)) {
            psFree(overscanOpts->poly);
            overscanOpts->poly = NULL;
        }
        if (! overscanOpts->poly) {
            overscanOpts->poly = psPolynomial1DAlloc(PS_POLYNOMIAL_ORD, overscanOpts->order);
        }
        psVectorFitPolynomial1D(overscanOpts->poly, mask, 1, reduced, NULL, ordinate);
        fitted = psPolynomial1DEvalVector(overscanOpts->poly, ordinate);
        break;
    case PM_FIT_POLY_CHEBY:
        if (overscanOpts->poly && (overscanOpts->poly->nX != overscanOpts->order ||
                                   overscanOpts->poly->type != PS_POLYNOMIAL_CHEB)) {
            psFree(overscanOpts->poly);
            overscanOpts->poly = NULL;
        }
        if (! overscanOpts->poly) {
            overscanOpts->poly = psPolynomial1DAlloc(PS_POLYNOMIAL_CHEB, overscanOpts->order);
        }
        psVectorFitPolynomial1D(overscanOpts->poly, mask, 1, reduced, NULL, ordinate);
        fitted = psPolynomial1DEvalVector(overscanOpts->poly, ordinate);
        break;
    case PM_FIT_SPLINE:
        // XXX I don't think psSpline1D is up to scratch yet --- it has no mask, and requires an
        // input spline
        overscanOpts->spline = psVectorFitSpline1D(reduced, ordinate);
        fitted = psSpline1DEvalVector(overscanOpts->spline, ordinate);
        break;
    default:
        psError(PS_ERR_UNKNOWN, true, "Unknown value for the fitting type: %d\n", overscanOpts->fitType);
        return NULL;
        break;
    }

    if (chi2) {
        *chi2 = 0.0;                    // chi^2 (sort of)
        for (int i = 0; i < reduced->n; i++) {
            *chi2 += PS_SQR(fitted->data.F32[i] - reduced->data.F32[i]);
        }
    }

    psFree(reduced);
    psFree(ordinate);
    psFree(mask);

    return fitted;
}


bool pmBiasSubtract(pmReadout *in, pmOverscanOptions *overscanOpts,
                    const pmReadout *bias, const pmReadout *dark)
{
    psTrace("psModules.detrend", 4,
            "---- pmBiasSubtract() begin ----\n");
    PS_ASSERT_PTR_NON_NULL(in, NULL);
    PS_ASSERT_IMAGE_NON_NULL(in->image, NULL);
    PS_ASSERT_IMAGE_TYPE(in->image, PS_TYPE_F32, NULL);
    if (bias) {
        PS_ASSERT_IMAGE_NON_NULL(bias->image, NULL);
        PS_ASSERT_IMAGE_TYPE(bias->image, PS_TYPE_F32, NULL);
    }
    if (dark) {
        PS_ASSERT_IMAGE_NON_NULL(dark->image, NULL);
        PS_ASSERT_IMAGE_TYPE(dark->image, PS_TYPE_F32, NULL);
    }

    pmHDU *hdu = pmHDUFromReadout(in);  // HDU of interest

    psImage *image = in->image;         // The input image

    // Overscan processing
    if (overscanOpts) {
        // Check for an unallowable pmFit.
        if (overscanOpts->fitType != PM_FIT_NONE && overscanOpts->fitType != PM_FIT_POLY_ORD &&
                overscanOpts->fitType != PM_FIT_POLY_CHEBY && overscanOpts->fitType != PM_FIT_SPLINE) {
            psError(PS_ERR_UNKNOWN, true, "Invalid fit type (%d).  Returning original image.\n",
                    overscanOpts->fitType);
            return false;
        }

        psList *overscans = in->bias; // List of the overscan images

        psStatsOptions statistic = psStatsSingleOption(overscanOpts->stat->options); // Statistic to use
        if (statistic == 0) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, false, "Multiple or no statistics options set: %p\n",
                    overscanOpts->stat);
            return false;
        }
        psStats *stats = psStatsAlloc(statistic); // A new psStats, to avoid clobbering original

        psString comment = NULL;    // Comment to add
        psStringAppend(&comment, "Subtracting overscan (stat %x; type %x; order %d)",
                       statistic, overscanOpts->fitType, overscanOpts->order);
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                         comment, "");
        psFree(comment);

        // Reduce all overscan pixels to a single value
        if (overscanOpts->single) {
            psVector *pixels = psVectorAlloc(0, PS_TYPE_F32);
            psListIterator *iter = psListIteratorAlloc(overscans, PS_LIST_HEAD, false); // Iterator
            psImage *overscan = NULL;   // Overscan image from iterator
            while ((overscan = psListGetAndIncrement(iter))) {
                int index = pixels->n;  // Index
                pixels = psVectorRealloc(pixels, pixels->n + overscan->numRows * overscan->numCols);
                pixels->n += overscan->numRows * overscan->numCols;
                for (int i = 0; i < overscan->numRows; i++) {
                    memcpy(&pixels->data.F32[index], overscan->data.F32[i],
                           overscan->numCols * sizeof(psF32));
                    index += overscan->numCols;
                }
            }
            psFree(iter);

            (void)psVectorStats(stats, pixels, NULL, NULL, 0);
            psFree(pixels);
            double reduced = psStatsGetValue(stats, statistic); // Result of statistics

            psString comment = NULL;    // Comment to add
            psStringAppend(&comment, "Overscan value: %f", reduced);
            psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK, comment, "");
            psFree(comment);

            (void)psBinaryOp(image, image, "-", psScalarAlloc((float)reduced, PS_TYPE_F32));
        } else {
            // We do the regular overscan subtraction
            bool readRows = psMetadataLookupBool(NULL, in->parent->concepts, "CELL.READDIR");// Read direction
            float chi2 = NAN;           // chi^2 from fit

            if (readRows) {
                // The read direction is rows
                psArray *pixels = psArrayAlloc(image->numRows); // Array of vectors containing pixels
                for (int i = 0; i < pixels->n; i++) {
                    pixels->data[i] = psVectorAlloc(0, PS_TYPE_F32);
                }

                // Pull the pixels out into the vectors
                psListIterator *iter = psListIteratorAlloc(overscans, PS_LIST_HEAD, false); // Iterator
                psImage *overscan = NULL; // Overscan image from iterator
                while ((overscan = psListGetAndIncrement(iter))) {
                    int diff = image->row0 - overscan->row0; // Offset between the two regions
                    for (int i = PS_MAX(0,diff); i < PS_MIN(image->numRows, overscan->numRows + diff); i++) {
                        // i is row on overscan
                        psVector *values = pixels->data[i];
                        int index = values->n; // Index in the vector
                        values = psVectorRealloc(values, values->n + overscan->numCols);
                        values->n += overscan->numCols;
                        memcpy(&values->data.F32[index], overscan->data.F32[i],
                               overscan->numCols * PSELEMTYPE_SIZEOF(PS_TYPE_F32));
                        index += overscan->numCols;
                        pixels->data[i] = values; // Update the pointer in case it's moved
                    }
                }
                psFree(iter);

                // Reduce the overscans
                psVector *reduced = overscanVector(&chi2, overscanOpts, pixels, stats);
                psFree(pixels);
                if (! reduced) {
                    psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to generate overscan vector.\n");
                    return false;
                }

                // Subtract row by row
                for (int i = 0; i < image->numRows; i++) {
                    for (int j = 0; j < image->numCols; j++) {
                        image->data.F32[i][j] -= reduced->data.F32[i];
                    }
                }
                psFree(reduced);

            } else {
                // The read direction is columns
                psArray *pixels = psArrayAlloc(image->numCols); // Array of vectors containing pixels
                for (int i = 0; i < pixels->n; i++) {
                    psVector *values = psVectorAlloc(0, PS_TYPE_F32);
                    pixels->data[i] = values;
                }

                // Pull the pixels out into the vectors
                psListIterator *iter = psListIteratorAlloc(overscans, PS_LIST_HEAD, false); // Iterator
                psImage *overscan = NULL; // Overscan image from iterator
                while ((overscan = psListGetAndIncrement(iter))) {
                    int diff = image->col0 - overscan->col0; // Offset between the two regions
                    for (int i = PS_MAX(0,diff); i < PS_MIN(image->numCols, overscan->numCols + diff); i++) {
                        // i is column on overscan
                        psVector *values = pixels->data[i];
                        int index = values->n; // Index in the vector
                        values = psVectorRealloc(values, values->n + overscan->numRows);
                        for (int j = 0; j < overscan->numRows; j++) {
                            values->data.F32[index++] = overscan->data.F32[j][i];
                        }
                        values->n += overscan->numRows;
                        pixels->data[i] = values; // Update the pointer in case it's moved
                    }
                }
                psFree(iter);

                // Reduce the overscans
                psVector *reduced = overscanVector(&chi2, overscanOpts, pixels, stats);
                psFree(pixels);
                if (! reduced) {
                    psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to generate overscan vector.\n");
                    return false;
                }

                // Subtract column by column
                for (int i = 0; i < image->numCols; i++) {
                    for (int j = 0; j < image->numRows; j++) {
                        image->data.F32[j][i] -= reduced->data.F32[i];
                    }
                }
                psFree(reduced);
            }

            switch (overscanOpts->fitType) {
            case PM_FIT_POLY_ORD:
            case PM_FIT_POLY_CHEBY: {
                    psString comment = NULL;    // Comment to add
                    psStringAppend(&comment, "Overscan fit (chi2: %.2f): ", chi2);
                    psPolynomial1D *poly = overscanOpts->poly; // The polynomial
                    for (int i = 0; i < poly->nX; i++) {
                        psStringAppend(&comment, "%.1f ", poly->coeff[i]);
                    }
                    psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK, comment, "");
                    psFree(comment);
                    break;
                }
            case PM_FIT_SPLINE: {
                    psSpline1D *spline = overscanOpts->spline; // The spline
                    for (int i = 0; i < spline->n; i++) {
                        psStringAppend(&comment, "Overscan fit (chi2: %.2f) %d:", chi2, i);
                        psPolynomial1D *poly = spline->spline[i]; // i-th polynomial
                        for (int j = 0; j < poly->nX; j++) {
                            psStringAppend(&comment, "%.1f ", poly->coeff[i]);
                        }
                        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                                         comment, "");
                        psFree(comment);
                    }
                    break;
                }
            case PM_FIT_NONE:
                break;
            default:
                psAbort(__func__, "Should never get here!!!\n");
            }


        }
        psFree(stats);
    } // End of overscan subtraction

    // Bias frame subtraction
    if (bias) {
        psVector *md5 = psImageMD5(bias->image); // md5 hash
        psString md5string = psMD5toString(md5); // String
        psFree(md5);
        psStringPrepend(&md5string, "BIAS image MD5: ");
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                         md5string, "");
        psFree(md5string);

        if (!subtractFrame(in, bias, 1.0)) {
            return false;
        }
    }

    if (dark) {
        // Get the scaling
        float inTime = psMetadataLookupF32(NULL, in->parent->concepts, "CELL.DARKTIME");
        float darkTime = psMetadataLookupF32(NULL, dark->parent->concepts, "CELL.DARKTIME");

        psVector *md5 = psImageMD5(dark->image); // md5 hash
        psString md5string = psMD5toString(md5); // String
        psFree(md5);
        psStringPrepend(&md5string, "DARK image (scale %.3f) MD5: ", inTime/darkTime);
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                         md5string, "");
        psFree(md5string);

        if (!subtractFrame(in, dark, inTime/darkTime)) {
            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, "Overscan/bias/dark processing completed at ");
    psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                     timeString, "");
    psFree(timeString);


    return true;
}


