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

#define NUM_BG 1000000                  // Number of random background measurements to make
#define SIGMA 3.0                       // Gaussian width for photometry
#define SOURCE_SIZE (3 * SIGMA)         // Size of source for photometry

static psImage *readImage(const char *name)
{
    psFits *fits = psFitsOpen(name, "r");
    psImage *image = psFitsReadImage(fits, psRegionSet(0,0,0,0), 0);
    psFitsClose(fits);
    return image;
}


double measure(const psImage *image, const psImage *mask, psMaskType maskVal, double bg, int xStar, int yStar)
{
    double sumGauss = 0.0;
    double sumFluxGauss = 0.0;
    for (int y = yStar - SOURCE_SIZE, v = -SOURCE_SIZE; v <= SOURCE_SIZE; y++, v++) {
        for (int x = xStar - SOURCE_SIZE, u = -SOURCE_SIZE; u <= SOURCE_SIZE; x++, u++) {
            if (mask->data.PS_TYPE_MASK_DATA[y][x] & maskVal) {
                continue;
            }
            if (!isfinite(image->data.F32[y][x])) {
                continue;
            }
            float gauss = expf(-(PS_SQR(u) + PS_SQR(v)) * 0.5 / PS_SQR(SIGMA));

            sumGauss += gauss;
            sumFluxGauss += gauss * (image->data.F32[y][x] - bg);
        }
    }

    return sumFluxGauss / sumGauss;
}



int main(int argc, char *argv[])
{
    psImage *image = readImage("faint.conv.fits");
    psImage *mask = readImage("mask.conv.fits");
    psImage *significance = readImage("faint.sig.fits");

    assert(image->type.type == PS_TYPE_F32);
    assert(mask->type.type == PS_TYPE_MASK);

    int numCols = image->numCols, numRows = image->numRows;
    psMaskType maskVal = 0;

    psStats *bgStats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN);
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS, 0);
    psImageBackground(bgStats, NULL, image, mask, maskVal, rng);
    double bg = bgStats->robustMedian;
    assert(isfinite(bg));

    psArray *stars = psVectorsReadFromFile("fake_stars.dat", "%f %f");
    psVector *xStars = stars->data[0], *yStars = stars->data[1];
    int number = xStars->n;
    psVector *fluxStars = psVectorAlloc(number, PS_TYPE_F64);
    psVector *sigStars = psVectorAlloc(number, PS_TYPE_F64);

    int numStars = 0;
    for (int i = 0; i < number; i++) {
        int xStar = xStars->data.F32[i] + 0.5;
        int yStar = yStars->data.F32[i] + 0.5;

        if (xStar < SOURCE_SIZE || xStar > numCols - SOURCE_SIZE - 1 ||
            yStar < SOURCE_SIZE || yStar > numRows - SOURCE_SIZE - 1) {
            continue;
        }

        if (significance->data.F32[yStar][xStar] <= 0.0 || !isfinite(significance->data.F32[yStar][xStar])) {
            continue;
        }

        float flux = measure(image, mask, maskVal, bg, xStar, yStar);
        if (!isfinite(flux)) {
            continue;
        }

        fluxStars->data.F64[numStars] = flux;
        sigStars->data.F64[numStars] = significance->data.F32[yStar][xStar];
        numStars++;
    }

    psVector *fluxBG = psVectorAlloc(NUM_BG, PS_TYPE_F64);
    int numBG = 0;
    double meanBG = 0.0;
    for (int i = 0; i < NUM_BG; i++) {
        int xBG = SOURCE_SIZE + (numCols - 2 * SOURCE_SIZE - 1) * psRandomUniform(rng);
        int yBG = SOURCE_SIZE + (numRows - 2 * SOURCE_SIZE - 1) * psRandomUniform(rng);

        float flux = measure(image, mask, maskVal, bg, xBG, yBG);

        if (!isfinite(flux)) {
            continue;
        }

        fluxBG->data.F64[numBG] = flux;
        meanBG += flux;
        numBG++;
    }
    meanBG /= numBG;

    psFree(image);
    psFree(mask);
    psFree(significance);

    double stdevBG = 0.0;
    for (int i = 0; i < numBG; i++) {
        stdevBG += PS_SQR(fluxBG->data.F64[i] - meanBG);
    }
    stdevBG = sqrt(stdevBG / numBG);

    double sum = 0.0;
    for (int i = 0; i < numStars; i++) {
        double actual = fabs(fluxStars->data.F64[i]) / stdevBG;
        double expected = sqrt(sigStars->data.F64[i]);
        sum += actual/expected;
    }
    double meanRatio = sum / numStars;

    printf("%d %d %lf %lf %lf %lf\n", numBG, numStars, bg, meanBG, stdevBG, meanRatio);

    return PS_EXIT_SUCCESS;
}
