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

// Add fake sources to an image so we can measure their significance
#define NUM_SOURCES 1000                // Number of sources to add
#define SIGMA 3.0                       // Gaussian width of sources
#define SOURCE_SIZE (3 * (SIGMA))       // Size of source
#define BRIGHT_PEAK_FLUX 10000          // Peak flux of bright sources
#define FAINT_PEAK_FLUX 10              // Peak flux of faint sources
#define NOISE 10                        // Gaussian noise
#define IMAGE_SIZE 4096                 // Size of image


void writeImage(const psImage *image,   // Image to write
                const char *name        // Name of FITS file
    )
{
    psFits *fits = psFitsOpen(name, "w");
    psFitsWriteImage(fits, NULL, image, 0, NULL);
    psFitsClose(fits);
}


void addStar(psImage *image, int xStar, int yStar, float peak)
{
    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++) {
            image->data.F32[y][x] += peak * expf(-(PS_SQR(u) + PS_SQR(v)) * 0.5 / PS_SQR(SIGMA));
        }
    }
    return;
}



int main(int argc, char *argv[])
{
    // Mask and weight are easy: they are constant
    psImage *mask = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_MASK);
    psImageInit(mask, 0);
    psImage *weight = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_F32);
    psImageInit(weight, PS_SQR(NOISE));

    writeImage(mask, "mask.fits");
    writeImage(weight, "weight.fits");
    psFree(mask);
    psFree(weight);

    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS, 0);

    psImage *bright = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_F32);
    psImage *faint = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_F32);

    // Start with noise
    for (int y = 0; y < IMAGE_SIZE; y++) {
        for (int x = 0; x < IMAGE_SIZE; x++) {
            bright->data.F32[y][x] = psRandomGaussian(rng) * NOISE;
            faint->data.F32[y][x] = psRandomGaussian(rng) * NOISE;
        }
    }

    // Add stars
    FILE *stars = fopen("fake_stars.dat", "w");
    for (int i = 0; i < NUM_SOURCES; i++) {
        // Position for stars
        int xStar = SOURCE_SIZE + (IMAGE_SIZE - 2 * SOURCE_SIZE - 1) * psRandomUniform(rng);
        int yStar = SOURCE_SIZE + (IMAGE_SIZE - 2 * SOURCE_SIZE - 1) * psRandomUniform(rng);
        fprintf(stars, "%d %d\n", xStar, yStar);

        addStar(bright, xStar, yStar, BRIGHT_PEAK_FLUX);
        addStar(faint, xStar, yStar, FAINT_PEAK_FLUX);
    }
    fclose(stars);

    psFree(rng);

    writeImage(bright, "bright.fits");
    writeImage(faint, "faint.fits");

    psFree(bright);
    psFree(faint);

    return PS_EXIT_SUCCESS;
}
