#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"

// XXX allow the user to choose the stats method?
// (SAMPLE_MEAN, CLIPPED_MEAN, ROBUST_MEDIAN, FITTED_MEAN)
psStats *psImageBackground(const psImage *image, const psImage *mask, psMaskType maskValue,
                           double fmin, double fmax, long nMax, psStatsOptions option, 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_NONNEGATIVE(nMax, NULL);
    PS_ASSERT_PTR_NON_NULL(rng, NULL);

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

    const int Npixels = nx*ny;                // Total number of pixels
    const int Nsubset = (nMax == 0) ? Npixels : PS_MIN(nMax, Npixels); // Number of pixels in subset

    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];

    // select a subset of the image pixels to measure the stats
    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;

    psStats *stats = psStatsAlloc(option); // Statistics, for return

    if (stats->options & PS_STAT_ROBUST_QUARTILE) {
        // use simple stats code (old verions)
        // XXX this hack is just for testing, drop when I am happy with the psStats version of the values

        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];
    } else {
        // XXX leave this as a psphot user option (passed in as part of stats?)
        if (stats->options & PS_STAT_FITTED_MEAN) {
            stats->clipSigma = 1.0;
        }
        if (psVectorStats (stats, values, NULL, NULL, 0) == NULL) {
            psError(PS_ERR_UNKNOWN, false, "Unable to measure statistics for image background "
                    "(%dx%d, (row0,col0) = (%d,%d)",
                    image->numRows, image->numCols, image->row0, image->col0);
            psFree(stats);
            psFree(values);
            return NULL;
        }
    }

    psFree(values);
    return stats;
}
