#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 "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
    )
{
    assert(bscale);
    assert(bzero);
    assert(image);
    assert(options);

    double range = pow(2.0, options->bitpix); // Range of values for target BITPIX
    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 - 1.0);
        *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 psFitsOptions *options // FITS options
    )
{
    assert(bscale);
    assert(bzero);
    assert(image);
    assert(options);

    // 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, NULL, 0, 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;


    return true;
}

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

bool psFitsScaleDetermine(double *bscale, // Scaling, to return
                          double *bzero, // Zero point, to return
                          const psImage *image, // Image to scale
                          const psFits *fits // FITS options
    )
{
    PS_ASSERT_PTR_NON_NULL(bscale, false);
    PS_ASSERT_PTR_NON_NULL(bzero, false);
    PS_ASSERT_IMAGE_NON_NULL(image, false);
    PS_ASSERT_FITS_NON_NULL(fits, false);

    *bscale = 1.0;
    *bzero = 0.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;
    }

    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, 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;
    }

    psTrace("psLib.fits", 3, "BSCALE = %.10lf, BZERO = %.10lf\n", *bscale, *bzero);
    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:
        outType = PS_TYPE_S8;
        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: { \
        ps##INTYPE scale = 1.0 / bscale; \
        ps##INTYPE zero = bzero; \
        ps##INTYPE min = - (1 << (options->bitpix - 1)); \
        ps##INTYPE max = (1 << (options->bitpix - 1)) - 1; \
        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 */ \
                    (OUT)->data.OUTTYPE[y][x] = (value < min ? 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, S8); \
            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;
}


