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

#include <stdio.h>
#include <assert.h>
#include "pslib.h"
#include "stac.h"

#define SQUARE(x) ((x)*(x))

#define SAMPLE 10                       // Subsample rate for images

float stacBackground(const psImage *image, // Image for which to get the background
                     int sample         // Sample in increments of this value
    )
{
    assert(image->type.type == PS_TYPE_F32);

#if 1
    int size = image->numCols * image->numRows; // Number of pixels in image
    int numSamples = size / sample;     // Number of samples in image
    psVector *values = psVectorAlloc(numSamples + 1, PS_TYPE_F32); // Vector containing sub-sample
    psVector *mask = psVectorAlloc(numSamples + 1, PS_TYPE_U8); // Mask for sample

    int offset = 0;                     // Offset from start of the row
    int index = 0;                      // Sample number
    for (int row = 0; row < image->numRows; row++) {
        // I'd cast this as a "for", but this makes it a bit easier to understand.
        int col = offset;
        while (col < image->numCols) {
            values->data.F32[index] = image->data.F32[row][col];
            if (isnan(values->data.F32[index])) {
                mask->data.U8[index] = 1;
            } else {
                mask->data.U8[index] = 0;
            }
            col += sample;
            index++;
        }
        offset = col - image->numCols;
    }

    psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEDIAN);
    psVectorStats(stats, values, NULL, mask, 1);
    float median = stats->sampleMedian;
    psFree(stats);
    psFree(values);
    psFree(mask);
#else
    // Will use robust median instead of sampling --- it's supposed to be fast.
    psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN); // Using a clipped mean because median is SLOW
    stats->clipSigma = 3.0;
    stats->clipIter = 5;
    stats = psImageStats(stats, (psImage*)image, NULL, 0);
    float median = stats->robustMedian;
    psFree(stats);
#endif

    return median;
}


bool stacScales(psVector **scalesPtr,   // Scales to return
                psVector **offsetsPtr,  // Offsets to return
                const psArray *images,  // Images on which to measure the scales and offsets
                const char *starFile,   // File containing coordinates to photometer
                const char *starMapFile, // Map for coodinates to the common output frame
                float xMapDiff, float yMapDiff, // Difference from the map to apply (due to size difference)
                float aper              // Aperture to use for photometry (radius)
    )
{
    assert(scalesPtr);
    assert(offsetsPtr);
    assert(images);
    for (int i = 0; i < images->n; i++) {
        psImage *image = images->data[i];
        assert(image->type.type == PS_TYPE_F32);
    }

    psVector *scales = NULL; // Relative scales between images
    psVector *offsets = NULL; // Offsets between images (ADU)
    if (*scalesPtr) {
        scales = *scalesPtr;
        assert(scales);
        assert(scales->n == images->n);
        assert(scales->type.type == PS_TYPE_F32);
    } else {
        scales = psVectorAlloc(images->n, PS_TYPE_F32);
        *scalesPtr = scales;
    }
    if (*offsetsPtr) {
        offsets = *offsetsPtr;
        assert(offsets);
        assert(offsets->n == images->n);
        assert(offsets->type.type == PS_TYPE_F32);
    } else {
        offsets = psVectorAlloc(images->n, PS_TYPE_F32);
        *offsetsPtr = offsets;
    }

    // Offsets are easy
    psTrace("stac.scales", 3, "Getting background levels....\n");
    double startTime = getTime();
    for (int i = 0; i < images->n; i++) {
        offsets->data.F32[i] = stacBackground(images->data[i], SAMPLE);
        psTrace("stac.scales", 5, "Background in image %d is %f\n", i, offsets->data.F32[i]);
        double time = getTime();
        psTrace("stac.scales", 10, "Took %f sec\n", time - startTime);
        startTime = time;
    }

    // Now the scales
    if (starFile == NULL || starMapFile == NULL) {
        psLogMsg("stac.scales", PS_LOG_INFO,
                 "No coordinates available to set scales --- assuming all are identical.\n");
        for (int i = 0; i < images->n; i++) {
            scales->data.F32[i] = 1.0;
            psTrace("stac.scales", 5, "Scale for image %d is %f\n", i, scales->data.F32[i]);
        }
    } else {
        // Read star coordinates and map
        psArray *starCoords = stacReadCoords(starFile); // Array of star coordinates
        psPlaneTransform *starMap = stacReadMap(starMapFile); // Transformation for star coordinates

        // Transform the stellar positions to match the transformed reference frame
        psArray *starCoordsTransformed = psArrayAlloc(starCoords->n); // Transformed positions
        // Fix up difference between map and output frame
        starMap->x->coeff[0][0] -= xMapDiff;
        starMap->y->coeff[0][0] -= yMapDiff;
        for (int i = 0; i < starCoords->n; i++) {
            starCoordsTransformed->data[i] = psPlaneTransformApply(NULL, starMap, starCoords->data[i]);
        }
        psFree(starMap);

        psArray *stars = psArrayAlloc(images->n); // Array of stellar photometry vectors
        psArray *masks = psArrayAlloc(images->n); // Array of masks for stars

        // Set scales relative to the first image
        scales->data.F32[0] = 1.0;

        // Iterate over images
        for (int i = 0; i < scales->n; i++) {
            psImage *image = images->data[i]; // The image we're working with

            // Do photometry on transformed image i
            psTrace("stac.scales", 3, "Doing photometry on image %d....\n", i);
            psVector *photometry = psVectorAlloc(starCoords->n, PS_TYPE_F32); // Photometry of the stars
            psVector *mask = psVectorAlloc(starCoords->n, PS_TYPE_U8); // Mask for the photometry
            for (int j = 0; j < starCoordsTransformed->n; j++) {
                psPlane *coords = starCoordsTransformed->data[j]; // The coordinates of the star

                if (coords->x < aper || coords->y < aper ||
                    coords->x + aper > image->numCols - 1 ||
                    coords->y + aper > image->numRows) {
                    mask->data.U8[j] = 1;
                } else {
                    // Sum flux within the aperture
                    float sum = 0.0;
                    int numPix = 0;
                    float aper2 = SQUARE(aper);
                    for (int y = (int)floorf(coords->y - aper);
                         y <= (int)ceilf(coords->y + aper); y++) {
                        for (int x = (int)floorf(coords->x - aper);
                             x <= (int)ceilf(coords->x + aper); x++) {
                            if (SQUARE((float)x + 0.5 - coords->x) + SQUARE((float)y + 0.5 - coords->y) <=
                                aper2) {
                                sum += image->data.F32[y][x];
                                numPix++;
                            }
                        }
                    }
                    // Subtract background, renormalise to account for circular aperture
                    if (numPix > 0 && sum > 0) {
                        sum -= offsets->data.F32[i] * (float)numPix;
                        photometry->data.F32[j] = sum * M_PI * aper2 / (float)numPix;
                        if (photometry->data.F32[j] > 0 && isfinite(photometry->data.F32[j])) {
                            mask->data.U8[j] = 1;
                            psTrace("stac.scales", 8, "Star at %f,%f --> %f\n", coords->x, coords->y, sum);
                        } else {
                            mask->data.U8[j] = 0;
                        }
                    } else {
                        mask->data.U8[j] = 0;
                    }
                }
            }
            stars->data[i] = photometry;
            masks->data[i] = mask;
        }
        psFree(starCoords);
        psFree(starCoordsTransformed);

        // Get the scales
        psVector *ref = stars->data[0]; // The reference photometry
        psVector *refMask = masks->data[0]; // The reference mask
#if 0
        psStats *stats = psStatsAlloc(PS_STAT_CLIPPED_MEAN); // Statistics
        stats->clipSigma = 2.5;
        stats->clipIter = 3;
#else
        psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEDIAN); // Statistics
#endif
        for (int i = 1; i < scales->n; i++) {
            psVector *input = stars->data[i];   // The comparison photometry
            psVector *inputMask = masks->data[i]; // The comparison mask

            psVector *compare = (psVector*)psBinaryOp(NULL, input, "/", ref);
            psVector *compareMask = (psVector*)psBinaryOp(NULL, inputMask, "*", refMask);
            (void)psBinaryOp(compareMask, psScalarAlloc(1, PS_TYPE_U8), "-", compareMask);

#if 0
            psTrace("stac.scales", 9, "Getting scale for image %d...\n", i);
            for (int j = 0; j < compare->n; j++) {
                if (compareMask->data.U8[j] & 1) {
                    psTrace("stac.scales", 9, "Bad star %d: %f %f --> %f %d\n", j, input->data.F32[j],
                            ref->data.F32[j], compare->data.F32[j], compareMask->data.U8[j]);
                } else {
                    psTrace("stac.scales", 9, "Good star %d: %f %f --> %f %d\n", j, input->data.F32[j],
                            ref->data.F32[j], compare->data.F32[j], compareMask->data.U8[j]);
                }
            }
#endif

            psVectorStats(stats, compare, NULL, compareMask, 1);

#if 0
            scales->data.F32[i] = stats->clippedMean;
#else
            scales->data.F32[i] = stats->sampleMedian;
#endif
            psTrace("stac.scales", 5, "Scale for image %d is %f\n", i, scales->data.F32[i]);
            psFree(compare);
            psFree(compareMask);
        }
        psFree(stats);

        psFree(stars);
        psFree(masks);
    }

    return true;
}


bool stacRescale(psArray *images,       // Images to rescale
                 psArray *errImages,    // Variance images to rescale
                 const psImage *mask,   // Mask for pixels to scale
                 const psVector *scales,// Scales for images
                 const psVector *offsets // Offsets for images
    )
{
    assert(images);
    assert(scales);
    assert(offsets);
    assert(images->n == scales->n);
    assert(images->n == offsets->n);
    assert(!errImages || errImages->n == images->n);
    assert(scales->type.type == PS_TYPE_F32);
    assert(offsets->type.type == PS_TYPE_F32);
    for (int i = 0; i < images->n; i++) {
        psImage *image = images->data[i]; // Image of interest
        assert(image->type.type == PS_TYPE_F32);
        assert(scales->data.F32[i] != 0);
        if (mask) {
            assert(mask->type.type == PS_TYPE_U8);
            assert(image->numCols == mask->numCols && image->numRows == mask->numRows);
        }
        if (errImages) {
            psImage *errImage = errImages->data[i];
            assert(errImage->type.type == PS_TYPE_F32);
            assert(errImage->numCols == image->numCols && errImage->numRows == image->numRows);
        }
    }

    for (int i = 0; i < images->n; i++) {
        psImage *image = images->data[i]; // Image to rescale
        psImage *errImage = NULL;       // Variance image to rescale
        if (errImages) {
            errImage = errImages->data[i];
        }
        float scale = scales->data.F32[i]; // Scale to use
        float offset = offsets->data.F32[i]; // Offset to use
        for (int y = 0; y < image->numRows; y++) {
            for (int x = 0; x < image->numCols; x++) {
                if (!mask || mask->data.F32[y][x]) {
                    image->data.F32[y][x] = (image->data.F32[y][x] - offset) / scale;
                    if (errImage) {
                        errImage->data.F32[y][x] /= SQUARE(scale);
                    }
                }
            }
        }
    }

    return true;
}
