#include <stdio.h>
#include <stdlib.h>

#include "psMemory.h"
#include "psImage.h"
#include "psVector.h"
#include "psStats.h"
#include "psType.h"
#include "psAssert.h"
#include "psRandom.h"
#include "psError.h"


psStats *psImageBackground(const psImage *image, const psImage *mask, psMaskType maskValue,
                           double fmin, double fmax, long nMax, psRandom *rng)
{
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    if (mask) {
        PS_ASSERT_IMAGE_NON_NULL(mask, NULL);
        PS_ASSERT_IMAGE_TYPE(mask, PS_TYPE_U8, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(mask, image, NULL);
    }
    PS_ASSERT_FLOAT_LARGER_THAN_OR_EQUAL(fmin, 0.0, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN_OR_EQUAL(fmax, 0.0, NULL);
    PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(fmin, 1.0, NULL);
    PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(fmax, 1.0, NULL);
    PS_ASSERT_INT_POSITIVE(nMax, NULL);
    PS_ASSERT_PTR_NON_NULL(rng, NULL);

    // Size of image
    long nx = image->numCols;
    long ny = image->numRows;

    psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_QUARTILE); // Statistics, for return

    int Nsubset = PS_MIN(nMax, nx*ny);  // Number of pixels in nubset
    int Npixels = nx*ny;                // Total number of pixels

    psVector *values = psVectorAllocEmpty(Nsubset, PS_TYPE_F32); // Vector containing subsample

    // Minimum and maximum values
    float min = values->data.F32[0];
    float max = values->data.F32[0];

    long n = 0;                         // Number of actual pixels in subset
    for (long i = 0; i < Nsubset; i++) {
        double frnd = psRandomUniform(rng);
        int pixel = Npixels * frnd;
        int ix = pixel % nx;
        int iy = pixel / nx;

        if (!isfinite(image->data.F32[iy][ix]) || (mask && mask->data.U8[iy][ix] & maskValue)) {
            continue;
        }

        float value = image->data.F32[iy][ix];
        min = PS_MIN(value, min);
        max = PS_MIN(value, max);
        values->data.F32[n] = value;
        n++;
    }
    values->n = n;

    int imin = fmin * n;
    int imax = fmax * n;
    int npts = imax - imin + 1;

    if (!psVectorSort(values, values)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to sort values.\n");
        psFree(stats);
        psFree(values);
        return NULL;
    }

    // Subtract the median when we add the numbers, so we don't get numerical problems
    float median = npts % 2 ? 0.5 * (values->data.F32[npts/2 - 1] + values->data.F32[npts/2]) :
                   values->data.F32[npts/2];
    double value = 0;
    for (long i = imin; (i <= imax) && (i < n); i++) {
        value += values->data.F32[i] - median;
    }
    value = value / npts + median;

    stats->robustMedian = value;
    stats->robustUQ = values->data.F32[imax];
    stats->robustLQ = values->data.F32[imin];

    psFree(values);
    return stats;
    // XXX correct for selection bias??
}
