// POIS: Pan-STARRS (or Pricey's) Optimal Image Subtraction
//
// Attempts to fit for a general convolution kernel that will match the PSFs between two images.
//
// Copyright (C) 2006  Paul A. Price (price@ifa.hawaii.edu)
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

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

#include <stdio.h>
#include <string.h>
#include <sys/time.h>
#include "pslib.h"
#include "pois.h"

#define MAXCHAR 80

// Gets the current time.
double getTime(void)
{
    struct timeval tv;

    gettimeofday(&tv, NULL);
    return(tv.tv_sec + 1.e-6 * tv.tv_usec);
}


int main(int argc, char **argv)
{
    double startTime;
    startTime = getTime();

    // Set trace levels
    (void)psTraceSetLevel(".", 0);

    // Parse the command line
    poisConfig *config = poisParseConfig(argc, argv); // Configuration values

    if (config->verbose) {
        (void)psTraceSetLevel("pois", 10);
        (void)psTraceSetLevel("pois.config", 10);
        (void)psTraceSetLevel("pois.time", 10);
        (void)psTraceSetLevel("pois.solution", 0);
        (void)psTraceSetLevel("pois.kernelBasisFunctions", 10);
        (void)psTraceSetLevel("pois.calculateEquation", 10);
        (void)psTraceSetLevel("pois.convolveImage", 10);
        (void)psTraceSetLevel("pois.findStamps", 10);
        (void)psTraceSetLevel("pois.parseConfig", 10);
        (void)psTraceSetLevel("pois.makeMask", 10);
        (void)psTraceSetLevel("pois.extractKernel", 10);
        (void)psTraceSetLevel("pois.calculateDeviations", 10);
        (void)psTraceSetLevel("pois.checkKernel", 10);
        if (config->stampFile) {
            (void)psTraceSetLevel("pois.checkStamp", 10);
        } else {
            (void)psTraceSetLevel("pois.checkStamp", 0);
        }
        // Set logging level
        psLogSetLevel(9);
    }

    // Read inputs
    psMetadata *header = NULL;          // Header for output image

    psRegion imageRegion = {0, 0, 0, 0};
    psFits *reference = psFitsOpen(config->refFile, "r");
    if (config->reverse) {
        header = psFitsReadHeader(NULL, reference);
    }
    psImage *refImage = psFitsReadImage(reference, imageRegion, 0);
    if (refImage == NULL) {
        psErrorStackPrint(stderr, "Fatal error: unable to read %s\n", config->refFile);
        exit(EXIT_FAILURE);
    }
    psTrace("pois", 3, "Read %s: %dx%d\n", config->refFile, refImage->numCols, refImage->numRows);
    psFitsClose(reference);
    // Convert to 32-bit floating point, in necessary
    if (refImage->type.type != PS_TYPE_F32) {
        psTrace("pois", 3, "Converting %s to floating point in memory....\n", config->refFile);
        psImage *temp = psImageCopy(NULL, refImage, PS_TYPE_F32);
        psFree(refImage);
        refImage = temp;
    }

    psFits *input = psFitsOpen(config->inFile, "r");
    if (! config->reverse) {
        header = psFitsReadHeader(NULL, input);
    }
    psImage *inImage = psFitsReadImage(input, imageRegion, 0);
    if (inImage == NULL) {
        psErrorStackPrint(stderr, "Fatal error: unable to read %s\n", config->inFile);
        exit(EXIT_FAILURE);
    }
    psTrace("pois", 3, "Read %s: %dx%d\n", config->inFile, inImage->numCols, inImage->numRows);
    psFitsClose(input);
    // Convert to 32-bit floating point, in necessary
    if (inImage->type.type != PS_TYPE_F32) {
        psTrace("pois", 3, "Converting %s to floating point in memory....\n", config->inFile);
        psImage *temp = psImageCopy(NULL, inImage, PS_TYPE_F32);
        psFree(inImage);
        inImage = temp;
    }

    if ((refImage->numCols != inImage->numCols) || (refImage->numRows != inImage->numRows)) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true,
                "Reference and input images must have the same dimensions: %dx%d vs %dx%d\n",
                refImage->numCols, refImage->numRows, inImage->numCols, inImage->numRows);
        exit(EXIT_FAILURE);
    }
    config->xImage = refImage->numCols;
    config->yImage = refImage->numRows;

    psTrace("pois.time", 1, "Inputs read at %f sec\n", getTime() - startTime);

    // Add a background offset to get the image all above zero
    // An image that goes negative (e.g., a background-subtracted image) produces
    // horrible results --- the matrix goes NaN because the noise is determined
    // as the square root of the value.
    if (config->background != 0.0) {
        (void)psBinaryOp(refImage, refImage, "+", psScalarAlloc(config->background, PS_TYPE_F32));
        (void)psBinaryOp(inImage, inImage, "+", psScalarAlloc(config->background, PS_TYPE_F32));
    }

    // Generate the kernel parameters
    psArray *kernelBasisFunctions = poisKernelBasisFunctions(config);

    // Generate a mask
    psImage *mask = poisMakeMask(NULL, refImage, inImage, config);

#ifdef TESTING
    // Write the mask out
    psFits *maskFile = psFitsOpen("mask.fits", "w");
    if (!psFitsWriteImage(maskFile, NULL, mask, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write mask: mask.fits\n");
    }
    psTrace("pois", 1, "Mask written to mask.fits\n");
    psFitsClose(maskFile);
#endif

    psArray *stamps = NULL;             // Array of stamps
    psVector *solution = NULL;          // Solution vector

    FILE *stampsFP = NULL;              // File pointer for stamps
    if (config->stampFile) {
        psTrace("pois", 5, "Opening stamp file, %s\n", config->stampFile);
        stampsFP = fopen(config->stampFile, "r");
        if (! stampsFP) {
            psError(PS_ERR_IO, true, "Unable to open stamps file, %s!\n", config->stampFile);
            exit(EXIT_FAILURE);
        }
    }

    // Iterate for a good solution
    bool badStamps = true;              // Do we have bad stamps, such that we need to continue to iterate?
    for (int iterNum = 0; iterNum < config->numIter && badStamps; iterNum++) {
        psTrace("pois", 1, "Iteration %d...\n", iterNum);

        // Find stamps
        if (config->stampFile) {
            stamps = poisReadStamps(stamps, &stampsFP, refImage, mask, config);
        } else {
            stamps = poisFindStamps(stamps, refImage, mask, config);
        }
        int numStamps = 0;
        for (int s = 0; s < stamps->n; s++) {
            poisStamp *stamp = stamps->data[s]; // Stamp of interest
            if (stamp->status != POIS_STAMP_BAD) {
                numStamps++;
            }
        }
        psTrace("pois.time", 1, "%d stamps found at %f sec\n", numStamps, getTime() - startTime);
        if (numStamps == 0) {
            fprintf(stderr, "No stamps found.  Check input parameters.\n");
            exit(EXIT_FAILURE);
        }

        // Calculate equation
        (void)poisCalculateEquation(stamps, refImage, inImage, kernelBasisFunctions, config);
        psTrace("pois.time", 1, "Equation calculated at %f sec\n", getTime() - startTime);

        // Solve the equation
        solution = poisSolveEquation(solution, stamps, config);
        psTrace("pois.time", 1, "Matrix equation solved at %f sec\n", getTime() - startTime);

#ifdef TESTING
        // Print solution
        psTrace("pois.solution", 5, "Solution:\n");
        for (int i = 0; i < solution->n - 1; i++) {
            poisKernelBasis *kernel = kernelBasisFunctions->data[i]; // The i-th kernel basis function
            psTrace("pois.solution", 7, "%d: (%d,%d) x^%d y^%d --> %f\n", i, kernel->u, kernel->v,
                    kernel->xOrder, kernel->yOrder, solution->data.F64[i]);
        }
#endif

#ifdef TESTING
        // Save the kernel postage stamp
        char kernelName[MAXCHAR];               // File name for kernel
        snprintf(kernelName, MAXCHAR, "kernel%d.fits", iterNum);
        psImage *kernelImage = poisExtractKernel(solution, kernelBasisFunctions, 0.0, 0.0, config);
        psFits *kernelFile = psFitsOpen(kernelName, "w");
        if (!psFitsWriteImage(kernelFile, NULL, kernelImage, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write kernel: %s\n", kernelName);
        }
        psTrace("pois", 1, "Kernel written to %s\n", kernelName);
        psFitsClose(kernelFile);
        psFree(kernelImage);
#endif

        // Calculate deviations for the stamps
        psVector *deviations = poisCalculateDeviations(NULL, stamps, refImage, inImage, mask,
                                                       kernelBasisFunctions, solution, config);
        badStamps = poisRejectStamps(stamps, mask, deviations, config);

        psFree(deviations);
    } // Iterating for good solution

    // If there was rejection on the last round, re-solve the equation using only the good stamps
    if (badStamps) {
        solution = poisSolveEquation(solution, stamps, config);
    }

    if (stampsFP) {
        fclose(stampsFP);
    }

    // Undo the mask for bad stamp area
    for (int j = 0; j < mask->numRows; j++) {
        for (int i = 0; i < mask->numCols; i++) {
            if (mask->data.U8[j][i] & POIS_MASK_STAMP) {
                mask->data.U8[j][i] &= ~POIS_MASK_STAMP;
            }
        }
    }

    // Check the result
    (void)poisCheckKernel(solution, kernelBasisFunctions, config);

    // Convolve the input image
    psImage *convImage = poisConvolveImage(refImage, mask, solution, kernelBasisFunctions, config);
    psTrace("pois.time", 1, "Convolution completed at %f sec\n", getTime() - startTime);

#ifdef TESTING
    // Save the convolved image
    char convName[MAXCHAR];             // Filename for convolved image
    snprintf(convName, MAXCHAR, "%s.conv", config->outFile);
    psFits *convFile = psFitsOpen(convName, "w");
    if (!psFitsWriteImage(convFile, NULL, convImage, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write convolved image: %s\n", convName);
    }
    psFitsClose(convFile);
    psTrace("pois", 1, "Convolved image written to %s\n", convName);
#endif

    // Do the subtraction
    psImage *subImage = psImageAlloc(config->xImage, config->yImage, PS_TYPE_F32);
    if (config->reverse) {
        (void)psBinaryOp(subImage, convImage, "-", inImage);
    } else {
        (void)psBinaryOp(subImage, inImage, "-", convImage);
    }
    psTrace("pois.time", 1, "Subtraction completed at %f sec\n", getTime() - startTime);

#ifdef TESTING
    {
        psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV);
        (void)psImageStats(stats, subImage, mask, 0);
        psTrace("pois", 5, "Subtracted image statistics: %f %f\n", stats->sampleMean, stats->sampleStdev);
        psFree(stats);
    }
#endif

    // Mask out the edges
    for (int y = config->yKernel; y < subImage->numRows - config->yKernel; y++) {
        for (int x = config->xKernel; x < subImage->numCols - config->xKernel; x++) {

            for (int x = 0; x < config->xKernel; x++) {
                subImage->data.F32[y][x] = 0.0;
            }
            for (int x = subImage->numCols - config->xKernel; x < subImage->numCols; x++) {
                subImage->data.F32[y][x] = 0.0;
            }
        }
    }
    for (int y = 0; y < config->yKernel; y++) {
        for (int x = 0; x < subImage->numCols; x++) {
            subImage->data.F32[y][x] = 0.0;
        }
    }
    for (int y = subImage->numRows - config->yKernel; y < subImage->numRows; y++) {
        for (int x = 0; x < subImage->numCols; x++) {
            subImage->data.F32[y][x] = 0.0;
        }
    }


    // Mask out bad pixels
    for (int y = config->yKernel; y < subImage->numRows - config->yKernel; y++) {
        for (int x = config->xKernel; x < subImage->numCols - config->xKernel; x++) {
            if (mask->data.U8[y][x]) {
                subImage->data.F32[y][x] = 0.0;
            }
        }
    }

    if (config->mask) {
        char maskName[80];              // Name of mask image
        sprintf(maskName, "%s.mask", config->outFile);
        psFits *maskFile = psFitsOpen(maskName, "w");
        if (!psFitsWriteImage(maskFile, NULL, mask, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", maskName);
        }
        psTrace("pois", 1, "Mask image written to %s\n", maskName);
        psFitsClose(maskFile);
    }


    int numNaN = psImageClipNaN(subImage, 0.0);
    if (numNaN > 0) {
        printf("Masked %d NAN pixels to zero.\n", numNaN);
    }

    psFits *subFile = psFitsOpen(config->outFile, "w");
    if (!psFitsWriteImage(subFile, header, subImage, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write subtracted image: %s\n", config->outFile);
    }
    psFitsClose(subFile);
    psTrace("pois", 1, "Subtracted image written to %s\n", config->outFile);
    psTrace("pois.time", 1, "I/O completed at %f sec\n", getTime() - startTime);

    // Clean up
    psFree(header);
    psFree(kernelBasisFunctions);
    psFree(refImage);
    psFree(inImage);
    psFree(convImage);
    psFree(subImage);
    psFree(config);
}

