#ifdef HAVE_CONFIG_H
# include "config.h"
#endif

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

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

static int nFailures = 0;

void psImageBackgroundInit() {
    nFailures = 0;
    return;
}

// XXX allow the user to choose the stats method?
// (SAMPLE_MEAN, CLIPPED_MEAN, ROBUST_MEDIAN, FITTED_MEAN)
bool psImageBackground(psStats *stats, psVector **sample, const psImage *image, const psImage *mask, psImageMaskType maskValue, 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_IMAGE_MASK, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(mask, image, NULL);
    }
    if (stats->options & PS_STAT_ROBUST_QUARTILE) {
        PS_ASSERT_FLOAT_LARGER_THAN_OR_EQUAL(stats->min, 0.0, NULL);
        PS_ASSERT_FLOAT_LARGER_THAN_OR_EQUAL(stats->max, 0.0, NULL);
        PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(stats->min, 1.0, NULL);
        PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(stats->max, 1.0, NULL);
    }
    PS_ASSERT_INT_NONNEGATIVE(stats->nSubsample, NULL);
    PS_ASSERT_PTR_NON_NULL(rng, NULL);

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

    psImage *bad = psImageAlloc(nx, ny, PS_TYPE_U8); // Image with bad pixels
    psImageInit(bad, 0);

    int Npixels = 0;                    // Total number of pixels
    for (int y = 0; y < ny; y++) {
        for (int x = 0; x < nx; x++) {
            if (!isfinite(image->data.F32[y][x]) ||
                (mask && mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] & maskValue)) {
                bad->data.U8[y][x] = 0xFF;
            }
            Npixels++;
        }
    }

    const int Nsubset = (stats->nSubsample == 0) ? Npixels : PS_MIN(stats->nSubsample, Npixels); // Number of pixels in subset

    psVector *values;                   // Vector containing subsample
    if (sample) {
        *sample = psVectorRecycle(*sample, Nsubset, PS_TYPE_F32);
        values = psMemIncrRefCounter(*sample);
        values->n = 0;
    } else {
        values = psVectorAllocEmpty(Nsubset, PS_TYPE_F32);
    }

    // Minimum and maximum values
    float min = +PS_MAX_F32;
    float max = -PS_MAX_F32;

    // select a subset of the image pixels to measure the stats
    long n = 0;                         // Number of actual pixels in subset
    if (Nsubset >= Npixels) {
        // if we have an image smaller than Nsubset, just loop over the (good) image pixels
        for (int iy = 0; iy < ny; iy++) {
            for (int ix = 0; ix < nx; ix++) {
                if (bad->data.U8[iy][ix]) {
                    continue;
                }

                float value = image->data.F32[iy][ix];
                min = PS_MIN(value, min);
                max = PS_MAX(value, max);
                values->data.F32[n] = value;
                n++;
            }
        }
    } else {
        // Subsample all pixels
        // This is not optimal, since there may be a large masked fraction that leaves us with few good pixels.
        // In that case, you should have set Nsubset.......
        Npixels = nx * ny;
        for (long i = 0; i < Nsubset; i++) {
            double frnd = psRandomUniform(rng);
            int pixel = Npixels * frnd;
            int ix = pixel % nx;
            int iy = pixel / nx;

            if (bad->data.U8[iy][ix]) {
                continue;
            }

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

    psFree(bad);

    if (n < 0.01*Nsubset) {
        if ((nFailures < 3) || (nFailures % 100 == 0)) {
            psLogMsg("psLib.psImageBackground", PS_LOG_DETAIL, "Unable to measure image background: too few data points (%ld) (%d failures)", n, nFailures);
        }
        nFailures ++;
        psFree(values);
        return false;
    }

    values->n = n;

    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 = stats->min * n;
        int imax = stats->max * n;
        int npts = imax - imin + 1;

        if (!psVectorSort(values, values)) {
            if ((nFailures < 3) || (nFailures % 100 == 0)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to sort values.(%d failures)\n", nFailures);
            }
            nFailures ++;
            psFree(values);
            return false;
        }

        // 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 {
        if (!psVectorStats (stats, values, NULL, NULL, 0)) {
            if (psTraceGetLevel("psLib.imageops") >= 5) {
                FILE *f = fopen ("vector.dat", "w");
                int fd = fileno(f);
                p_psVectorPrint (fd, values, "values");
                fclose (f);
            }
            if ((nFailures < 3) || (nFailures % 100 == 0)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to measure statistics for image background "
                        "(%dx%d, (row0,col0) = (%d,%d) (%d failures)\n",
                        image->numRows, image->numCols, image->row0, image->col0, nFailures);
            }
            nFailures ++;
            psFree(values);
            return false;
        }
        if (psTraceGetLevel("psLib.imageops") >= 6) {
            FILE *f = fopen ("vector.dat", "w");
            int fd = fileno(f);
            p_psVectorPrint (fd, values, "values");
            fclose (f);
        }
    }

    psFree(values);

    return true;
}
