#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <assert.h>
#include <string.h>
#include <strings.h>

#include "psAbort.h"
#include "psAssert.h"
#include "psError.h"
#include "psTrace.h"

#include "psImage.h"
#include "psFits.h"
#include "psImageBackground.h"
#include "psStats.h"
#include "psImageStructManip.h"

#include "psMemory.h"

#include "psFitsScale.h"

// Remember:
// TRUE(i.e., value in memory) = BZERO + BSCALE * FITS(i.e., value on disk)


#define MEAN_STAT PS_STAT_ROBUST_MEDIAN // Statistic to use for mean
#define STDEV_STAT PS_STAT_ROBUST_STDEV // Statistic to use for stdev


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Private (file-static) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Determine appropriate BSCALE and BZERO for an image, preserving dynamic range
static bool scaleRange(double *bscale, // Scaling, to return
                       double *bzero, // Zero point, to return
                       const psImage *image, // Image to scale
                       const psFitsOptions *options // FITS options
    )
{
    psAssert(bscale, "impossible");
    psAssert(bzero, "impossible");
    psAssert(image, "impossible");
    psAssert(options, "impossible");

    double range = pow(2.0, options->bitpix) - 1.0; // Range of values for target BITPIX, reduced by the BLANK
    double min = INFINITY, max = -INFINITY; // Minimum and maximum values
    int numCols = image->numCols, numRows = image->numRows; // Size of image

    // Determine the minimum and maximum values
#define SCALE_DETERMINE_CASE(IN, INTYPE) \
    case PS_TYPE_##INTYPE: { \
        for (int y = 0; y < numRows; y++) { \
            for (int x = 0; x < numCols; x++) { \
                ps##INTYPE value = (IN)->data.INTYPE[y][x]; /* Value of interest */ \
                if (isfinite(value)) { \
                    if (value < min) { \
                        min = value; \
                    } \
                    if (value > max) { \
                        max = value; \
                    } \
                } \
            } \
        } \
        break; \
    }

    switch (image->type.type) {
        SCALE_DETERMINE_CASE(image, F32);
        SCALE_DETERMINE_CASE(image, F64);
      default:
        psAbort("Should be unreachable.");
    }

    if (!isfinite(min) || !isfinite(max)) {
        psWarning("No valid values in image to derive BSCALE,BZERO");
        *bscale = 1.0;
        *bzero = 0.0;
        return false;
    }
    if (min == max) {
        *bscale = 1.0;
        *bzero = min;
    } else {
        *bscale = (max - min) / range ;
        *bzero = min + 0.5 * range * (*bscale);
    }

    return true;
}

// Determine appropriate BSCALE and BZERO for an image, mapping the standard deviation to the nominated number
// of bits
static bool scaleStdev(double *bscale, // Scaling, to return
                       double *bzero, // Zero point, to return
                       const psImage *image, // Image to scale
                       const psImage *mask, // Mask image
                       psImageMaskType maskVal, // Value to mask
                       const psFitsOptions *options // FITS options
    )
{
    psAssert(bscale, "impossible");
    psAssert(bzero, "impossible");
    psAssert(image, "impossible");
    psAssert(options, "impossible");

    // Measure the mean and stdev
    // psImageBackground automatically excludes pixels that are non-finite, so we don't need to bother about a
    // mask.
    psU64 seed = p_psRandomGetSystemSeed(false);
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS, seed);
    psStats *stats = psStatsAlloc(MEAN_STAT | STDEV_STAT); // Statistics object
    if (!psImageBackground(stats, NULL, image, mask, maskVal, rng)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to perform statistics on image");
        psFree(rng);
        psFree(stats);
        return false;
    }
    psFree(rng);

    double mean = psStatsGetValue(stats, MEAN_STAT); // Mean
    double stdev = psStatsGetValue(stats, STDEV_STAT); // Standard deviation
    psFree(stats);
    if (!isfinite(mean) || !isfinite(stdev)) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                "Mean (%f) or stdev (%f) is non-finite.", mean, stdev);
        return false;
    }

    long range = 1 << options->stdevBits;  // Range of values to carry standard deviation
    *bscale = stdev / (double) range;

    double imageVal;                    // Value on image
    long diskVal;                       // Corresponding value on disk
    switch (options->scaling) {
      case PS_FITS_SCALE_STDEV_POSITIVE:
        // Put (mean - N sigma) at the lowest possible value: predominantly positive images
        imageVal = mean - options->stdevNum * stdev;
        diskVal = - (1 << (options->bitpix - 1));
        break;
      case PS_FITS_SCALE_STDEV_NEGATIVE:
        // Put (mean + N sigma) at the highest possible value: predominantly negative images
        imageVal = mean + options->stdevNum * stdev;
        diskVal = (1 << (options->bitpix - 1)) - 1;
        break;
      case PS_FITS_SCALE_STDEV_BOTH:
        // Put mean right in the middle: images with an equal abundance of positive and negative values
        imageVal = mean;
        diskVal = 0;
        break;
      default:
        psAbort("Should never get here.");
    }

    *bzero = imageVal - *bscale * diskVal;

    return true;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

bool psFitsScaleDetermine(double *bscale, double *bzero, long *blank, const psImage *image,
                          const psImage *mask, psImageMaskType maskVal, const psFits *fits)
{
    PS_ASSERT_PTR_NON_NULL(bscale, false);
    PS_ASSERT_PTR_NON_NULL(bzero, false);
    PS_ASSERT_PTR_NON_NULL(blank, false);
    PS_ASSERT_IMAGE_NON_NULL(image, false);
    if (mask) {
        PS_ASSERT_IMAGE_TYPE(mask, PS_TYPE_IMAGE_MASK, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(mask, image, false);
    }
    PS_ASSERT_FITS_NON_NULL(fits, false);

    *bscale = NAN;
    *bzero = NAN;
    *blank = 0;

    psFitsOptions *options = fits->options; // FITS options
    if (!PS_IS_PSELEMTYPE_REAL(image->type.type) || !options) {
        return true;
    }

    switch (options->bitpix) {
      case 0:
        // No scaling applied
        return true;
      case 8:
      case 16:
      case 32:
      case 64:
        // Nothing to do; just allowing these values to pass through
        break;
      default:
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Target bitpix (%d) is not one of 0,8,16,32,64",
                options->bitpix);
        return false;
    }

    *blank = (1 << (options->bitpix - 1)) - 1;

    switch (options->scaling) {
      case PS_FITS_SCALE_NONE:
        // No scaling applied
        break;
      case PS_FITS_SCALE_RANGE:
        if (!scaleRange(bscale, bzero, image, options)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to set BSCALE and BZERO from range");
            return false;
        }
        break;
      case PS_FITS_SCALE_STDEV_POSITIVE:
      case PS_FITS_SCALE_STDEV_NEGATIVE:
      case PS_FITS_SCALE_STDEV_BOTH:
        if (!scaleStdev(bscale, bzero, image, mask, maskVal, options)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to set BSCALE and BZERO from stdev");
            return false;
        }
        break;
      case PS_FITS_SCALE_MANUAL:
        *bscale = options->bscale;
        *bzero = options->bzero;
        break;
      default:
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Unrecognised FITS scaling method: %x",
                options->scaling);
        return false;
    }

    if (options->bitpix == 8) {
        // FITS standard wants unsigned for BITPIX=8, two's-complement for BITPIX=16,32,64.
        *bzero -= *bscale * 128;
        *blank = 255;
    }


    psTrace("psLib.fits", 3, "BSCALE = %.10lf, BZERO = %.10lf, BLANK = %ld\n", *bscale, *bzero, *blank);
    return true;
}


psImage *psFitsScaleForDisk(const psImage *image, const psFits *fits, double bscale, double bzero,
                            psRandom *rng)
{
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    PS_ASSERT_FITS_NON_NULL(fits, NULL);

    psFitsOptions *options = fits->options; // FITS options
    if (!PS_IS_PSELEMTYPE_REAL(image->type.type) || !options || options->bitpix == 0) {
        // No scaling desired
        return psMemIncrRefCounter((psImage*)image); // Casting away "const"
    }

    int bitpix = options->bitpix;       // Bits per pixel
    psElemType outType;                 // Type for output image
    // Choosing to use signed types because those don't require BSCALE,BZERO to represent them in the FITS
    // file
    switch (bitpix) {
      case 8:
        // Note: Use unsigned integer for BITPIX=8
        outType = PS_TYPE_U8;
        break;
      case 16:
        outType = PS_TYPE_S16;
        break;
      case 32:
        outType = PS_TYPE_S32;
        break;
      case 64:
        outType = PS_TYPE_S64;
        break;
      default:
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Target bitpix (%d) is not one of 8,16,32,64", bitpix);
        return NULL;
    }

    if (bscale == 1.0 && bzero == 0.0 && !options->fuzz) {
        return psImageCopy(NULL, image, outType);
    }

    int numCols = image->numCols, numRows = image->numRows; // Size of image
    psImage *out = psImageAlloc(numCols, numRows, outType); // Output image

    if (!psMemIncrRefCounter(rng) && options->fuzz) {
        // Don't blab about which seed we're going to get --- it's not necessary for this purpose
        psU64 seed = p_psRandomGetSystemSeed(false);
        rng = psRandomAlloc(PS_RANDOM_TAUS, seed);
    }

#define SCALE_WRITE_OUT_CASE(IN, INTYPE, OUT, OUTTYPE) \
    case PS_TYPE_##OUTTYPE: { \
        double scale = 1.0 / bscale; \
        double zero = bzero; \
        /* Note: BITPIX=8 treated differently, since it uses unsigned values; the rest use signed */ \
        double min = bitpix == 8 ? 0 : -pow(2.0, options->bitpix - 1); \
        double max = bitpix == 8 ? 255 : (pow(2.0, options->bitpix - 1) - 1.0); \
        for (int y = 0; y < numRows; y++) { \
            for (int x = 0; x < numCols; x++) { \
                ps##INTYPE value = (IN)->data.INTYPE[y][x]; \
                if (!isfinite(value)) { \
                    /* This choice of "max" for non-finite pixels is mainly cosmetic --- it has to be */ \
                    /* something, and "min" would produce holes in the cores of bright stars. */ \
                    (OUT)->data.OUTTYPE[y][x] = max; \
                } else { \
                    value = (value - zero) * scale; \
                    if (options->fuzz) { \
                       /* Add random factor [0,1): adds a variance of 1/12, */ \
                       /* but preserves the expectation value */ \
                        value += psRandomUniform(rng); \
                    } \
                    /* Check for underflow and overflow; set either to max */ \
                    (OUT)->data.OUTTYPE[y][x] = (value < min || value > max ? max : value); \
                } \
            } \
        } \
        break; \
    }

#define SCALE_WRITE_IN_CASE(IN, INTYPE, OUT) \
    case PS_TYPE_##INTYPE: { \
        switch (outType) { \
            SCALE_WRITE_OUT_CASE(IN, INTYPE, OUT, U8); \
            SCALE_WRITE_OUT_CASE(IN, INTYPE, OUT, S16); \
            SCALE_WRITE_OUT_CASE(IN, INTYPE, OUT, S32); \
            SCALE_WRITE_OUT_CASE(IN, INTYPE, OUT, S64); \
          default: \
            psAbort("Should be unreachable."); \
        } \
        break; \
    }

    switch (image->type.type) {
        SCALE_WRITE_IN_CASE(image, F32, out);
        SCALE_WRITE_IN_CASE(image, F64, out);
      default:
        psAbort("Should be unreachable.");
    }

    psFree(rng);

    return out;
}


#if 0
// This function to apply BSCALE and BZERO to an image read immediately from disk should not be necessary at
// the present time, since cfitsio should apply the scaling itself in the process of reading.  However, we may
// later desire it (e.g., if we ever make our own FITS implementation).
psImage *psFitsScaleFromDisk(psFits *fits, psImage *image)
{
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);

    if (bscale == 0.0) {
        // BSCALE = 0 means don't apply anything
        return psMemIncrRefCounter(image);
    }

    psElemType inType = image->type.type; // Type for input image
    psElemType outType;                 // Type for output image
    switch (inType) {
      case PS_TYPE_S8:
      case PS_TYPE_S16:
      case PS_TYPE_S32:
      case PS_TYPE_U8:
      case PS_TYPE_U16:
        outType = PS_TYPE_F32;
        break;
      case PS_TYPE_S64:
      case PS_TYPE_U32:
      case PS_TYPE_U64:
        outType = PS_TYPE_F64;
        break;
        // Including floating-point types just in case someone wants to apply a BSCALE and BZERO to them.
      case PS_TYPE_F32:
        outType = PS_TYPE_F32;
        break;
      case PS_TYPE_F64:
        outType = PS_TYPE_F64;
        break;
      default:
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Unsupported image type: %x", inType);
        return NULL;
    }

    int numCols = image->numCols, numRows = image->numRows;
    psImage *out = psImageAlloc(numCols, numRows, outType); // Output scaled image


#define SCALE_READ_OUT_CASE(INTYPE, OUTTYPE) \
  case PS_TYPE_##OUTTYPE: { \
      for (int y = 0; y < numRows; y++) { \
          for (int x = 0; x < numCols; x++) { \
              out->data.OUTTYPE[y][x] = image->data.INTYPE[y][x] * bscale + bzero; \
          } \
      } \
      break; \
  }

#define SCALE_READ_IN_CASE(INTYPE) \
  case PS_TYPE_##INTYPE: { \
      switch (outType) { \
          SCALE_READ_OUT_CASE(INTYPE, F32); \
          SCALE_READ_OUT_CASE(INTYPE, F64); \
        default: \
          psAbort("Should never get here: type %x should be F32 or F64", outType); \
      } \
      break; \
  }

    switch (inType) {
        SCALE_READ_IN_CASE(S8);
        SCALE_READ_IN_CASE(S16);
        SCALE_READ_IN_CASE(S32);
        SCALE_READ_IN_CASE(S64);
        SCALE_READ_IN_CASE(U8);
        SCALE_READ_IN_CASE(U16);
        SCALE_READ_IN_CASE(U32);
        SCALE_READ_IN_CASE(U64);
        SCALE_READ_IN_CASE(F32);
        SCALE_READ_IN_CASE(F64);
      default:
          psAbort("Should never get here: type %x should be integer", inType);
    }

    return out;
}
#endif



psFitsScaling psFitsScalingFromString(const char *string)
{
    PS_ASSERT_STRING_NON_EMPTY(string, PS_FITS_SCALE_NONE);

    if (strcasecmp(string, "RANGE") == 0)          return PS_FITS_SCALE_RANGE;
    if (strcasecmp(string, "STDEV_POSITIVE") == 0) return PS_FITS_SCALE_STDEV_POSITIVE;
    if (strcasecmp(string, "STDEV_NEGATIVE") == 0) return PS_FITS_SCALE_STDEV_NEGATIVE;
    if (strcasecmp(string, "STDEV_BOTH") == 0)     return PS_FITS_SCALE_STDEV_BOTH;
    if (strcasecmp(string, "MANUAL") == 0)         return PS_FITS_SCALE_MANUAL;

    psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Unable to interpret FITS scaling: %s", string);
    return PS_FITS_SCALE_NONE;
}


