/** @file ppSubMatchPSFs.c
 *
 *  @brief
 *
 *  @ingroup ppSub
 *
 *  @author IfA
 *  @version $Revision: 1.3 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2009-02-18 00:31:20 $
 *  Copyright 2009 Institute for Astronomy, University of Hawaii
 */

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

#include <stdio.h>
#include <pslib.h>
#include <psmodules.h>
#include <psphot.h>

#include "ppSub.h"

#define COVAR_FRAC 0.01                 // Truncation fraction for covariance matrix

// Normalise a region on an image
static void normaliseRegion(psImage *image, // Image to normalise
                            const psRegion *region, // Region of image to normalise
                            float norm  // Normalisation
                            )
{
    if (!image) {
        return;
    }
    psAssert(region, "Expect region");
    psImage *subImage = psImageSubset(image, *region); // Sub-image
    psBinaryOp(subImage, subImage, "*", psScalarAlloc(norm, PS_TYPE_F32));
    psFree(subImage);
    return;
}

// Measure the PSF for an image
static float subImagePSF(ppSubData *data, // Processing data
                         const pmReadout *ro, // Readout for which to measure PSF
                         psArray *sources     // Sources with positions at which to measure PSF
    )
{
    psAssert(data, "Require processing data");
    pmConfig *config = data->config;    // Configuration
    psAssert(config, "Require configuration");

    psAssert(ro, "Need readout");
    psAssert(sources, "Need sources.");

    pmFPAfile *photFile = psMetadataLookupPtr(NULL, config->files, "PSPHOT.INPUT"); // Photometry file
    psAssert(photFile, "Need photometry file.");
    if (!pmFPACopy(photFile->fpa, ro->parent->parent->parent)) {
        psError(PPSUB_ERR_CONFIG, false, "Unable to copy FPA for photometry");
        return false;
    }

    pmFPAview *view = ppSubViewReadout(); // View to readout
    pmReadout *photRO = pmFPAviewThisReadout(view, photFile->fpa); // Readout to photometer

    if (psMetadataLookup(photRO->analysis, "PSPHOT.DETECTIONS")) {
        psMetadataRemoveKey(photRO->analysis, "PSPHOT.DETECTIONS");
    }
    if (psMetadataLookup(photRO->parent->parent->analysis, "PSPHOT.PSF")) {
        psMetadataRemoveKey(photRO->parent->parent->analysis, "PSPHOT.PSF");
    }

    // Extract the loaded sources from the associated readout, and generate PSF
    // Here, we assume the image is background-subtracted
    if (!psphotReadoutFindPSF(config, view, sources)) {
        psErrorStackPrint(stderr, "Unable to determine PSF");
        psWarning("Unable to determine PSF.");
        psFree(view);
        return true;
    }

    psFree(view);

    psMetadata *header = psMetadataLookupMetadata(NULL, photRO->analysis, "PSPHOT.HEADER");
    psAssert(header, "Require header.");
    float fwhm = psMetadataLookupF32(NULL, header, "FWHM_MAJ");

    if (psMetadataLookup(photRO->analysis, "PSPHOT.DETECTIONS")) {
        psMetadataRemoveKey(photRO->analysis, "PSPHOT.DETECTIONS");
    }
    if (psMetadataLookup(photRO->parent->parent->analysis, "PSPHOT.PSF")) {
        psMetadataRemoveKey(photRO->parent->parent->analysis, "PSPHOT.PSF");
    }

    return fwhm;
}

// Scale the kernel parameters according to the PSFs
static bool subScaleKernel(ppSubData *data, // Processing data
                           psVector *kernelWidths, // Widths for kernel
                           int *kernelSize,        // Size of kernel
                           int *stampSize          // Size of stamps (footprint)
    )
{
    psAssert(data, "Require processing data");
    pmConfig *config = data->config;    // Configuration
    psAssert(config, "Require configuration");

    // Nothing to do if pre-calculated kernel exists
    pmFPAview *view = ppSubViewReadout(); // View to readout
    pmReadout *kernelRO = pmFPAfileThisReadout(config->files, view, "PPSUB.OUTPUT.KERNELS"); // RO with kernel
    if (kernelRO) {
        psFree(view);
        return true;
    }

    // Look up recipe values
    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSUB_RECIPE); // Recipe for ppSim
    psAssert(recipe, "We checked this earlier, so it should be here.");
    if (!psMetadataLookupBool(NULL, recipe, "SCALE")) {
        // No scaling requested
        psFree(view);
        return true;
    }

    // Input images
    pmReadout *inRO = pmFPAfileThisReadout(config->files, view, "PPSUB.INPUT"); // Input readout
    pmReadout *refRO = pmFPAfileThisReadout(config->files, view, "PPSUB.REF"); // Reference readout

    // Input sources
    pmReadout *inSourceRO = pmFPAfileThisReadout(config->files, view, "PPSUB.INPUT.SOURCES");
    pmReadout *refSourceRO = pmFPAfileThisReadout(config->files, view, "PPSUB.REF.SOURCES");
    if (!inSourceRO || !refSourceRO) {
        psWarning("Unable to scale kernel, since no sources were provided.");
        return true;
    }

    pmDetections *inDetections = psMetadataLookupPtr(NULL, inSourceRO->analysis, "PSPHOT.DETECTIONS"); // Input sources
    pmDetections *refDetections = psMetadataLookupPtr(NULL, refSourceRO->analysis, "PSPHOT.DETECTIONS"); // Ref sources

    psFree(view);

    if (!inDetections || !refDetections) {
        psWarning("Unable to scale kernel, since no sources were provided.");
        return true;
    }

    psArray *inSources = inDetections->allSources;
    psAssert (inSources, "missing inSources?");

    psArray *refSources = refDetections->allSources;
    psAssert (refSources, "missing refSources?");

    float inFWHM = psMetadataLookupF32(NULL, inRO->parent->parent->concepts, "CHIP.SEEING"); // FWHM for input
    if (!isfinite(inFWHM) || inFWHM == 0.0) {
        inFWHM = subImagePSF(data, inRO, inSources);
    }
    float refFWHM = psMetadataLookupF32(NULL, refRO->parent->parent->concepts, "CHIP.SEEING"); // FWHM for ref
    if (!isfinite(refFWHM) || refFWHM == 0.0) {
        refFWHM = subImagePSF(data, refRO, refSources);
    }
    psLogMsg("ppSub", PS_LOG_INFO, "Input FWHM: %f\nReference FWHM: %f\n", inFWHM, refFWHM);
    if (!isfinite(inFWHM) || !isfinite(refFWHM)) {
        psWarning("Unable to scale kernel, since unable to measure PSFs.");
        return true;
    }

    float scaleRef = psMetadataLookupF32(NULL, recipe, "SCALE.REF"); // Reference for scaling
    float scaleMin = psMetadataLookupF32(NULL, recipe, "SCALE.MIN"); // Minimum for scaling
    float scaleMax = psMetadataLookupF32(NULL, recipe, "SCALE.MAX"); // Maximum for scaling
    if (!isfinite(scaleRef) || !isfinite(scaleMin) || !isfinite(scaleMax)) {
        psError(PPSUB_ERR_ARGUMENTS, false,
                "Scale parameters (SCALE.REF=%f, SCALE.MIN=%f, SCALE.MAX=%f) not set in recipe.",
                scaleRef, scaleMin, scaleMax);
        return false;
    }

    if (!pmSubtractionParamsScale(kernelSize, stampSize, kernelWidths, inFWHM, refFWHM,
                                  scaleRef, scaleMin, scaleMax)) {
        psError(PPSUB_ERR_DATA, false, "Unable to scale parameters.");
        return false;
    }

    return true;
}


bool ppSubMatchPSFs(ppSubData *data)
{
    psAssert(data, "Require processing data");
    pmConfig *config = data->config;    // Configuration
    psAssert(config, "Require configuration");

    // Look up recipe values
    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSUB_RECIPE); // Recipe for ppSim
    psAssert(recipe, "We checked this earlier, so it should be here.");

    pmFPAview *view = ppSubViewReadout(); // View to readout

    // Input images
    pmReadout *inRO = pmFPAfileThisReadout(config->files, view, "PPSUB.INPUT"); // Input readout
    pmReadout *refRO = pmFPAfileThisReadout(config->files, view, "PPSUB.REF"); // Reference readout

    // Output image holders
    pmReadout *inConv = pmFPAfileThisReadout(config->files, view, "PPSUB.INPUT.CONV"); // Input convolved
    if (!inConv) {
        pmCell *cell = pmFPAfileThisCell(config->files, view, "PPSUB.INPUT.CONV"); // Cell for convolved input
        inConv = pmReadoutAlloc(cell);
    }
    pmReadout *refConv = pmFPAfileThisReadout(config->files, view, "PPSUB.REF.CONV"); // Reference convolved
    if (!refConv) {
        pmCell *cell = pmFPAfileThisCell(config->files, view, "PPSUB.REF.CONV"); // Cell for convolved ref.
        refConv = pmReadoutAlloc(cell);
    }

    bool mdok;                          // Status of MD lookup

    // Load pre-calculated kernel, if available
    pmReadout *kernelRO = pmFPAfileThisReadout(config->files, view, "PPSUB.OUTPUT.KERNELS"); // RO with kernel

    // Sources in image, used for stamps: these must be loaded from previous analysis stages
    pmReadout *inSourceRO = pmFPAfileThisReadout(config->files, view, "PPSUB.INPUT.SOURCES");
    pmReadout *refSourceRO = pmFPAfileThisReadout(config->files, view, "PPSUB.REF.SOURCES");

    pmDetections *inDetections  = inSourceRO  ? psMetadataLookupPtr(&mdok, inSourceRO->analysis,  "PSPHOT.DETECTIONS") : NULL; // Input sources
    pmDetections *refDetections = refSourceRO ? psMetadataLookupPtr(&mdok, refSourceRO->analysis, "PSPHOT.DETECTIONS") : NULL; // Ref sources

    psFree(view);

    pmDetections *detections = NULL;    // Merged detection set
    if (inDetections && refDetections) {
        psArray *inSources  = inDetections->allSources;
        psArray *refSources = refDetections->allSources;

        psAssert (inSources, "missing in sources?");
        psAssert (refSources, "missing ref sources?");

        detections = pmDetectionsAlloc();
        float radius = psMetadataLookupF32(NULL, recipe, "SOURCE.RADIUS"); // Matching radius
        psArray *lists = psArrayAlloc(2); // Source lists
        lists->data[0] = psMemIncrRefCounter(inSources);
        lists->data[1] = psMemIncrRefCounter(refSources);
        detections->allSources = pmSourceMatchMerge(lists, radius);
        psFree(lists);
        if (!detections->allSources) {
            psError(PPSUB_ERR_DATA, false, "Unable to merge source lists");
            psFree(detections);
            return false;
        }
    }
    if (!detections && inDetections) {
        detections = psMemIncrRefCounter(inDetections);
    }
    if (!detections && refDetections) {
        detections = psMemIncrRefCounter(refDetections);
    }

    psMetadataAddPtr(inConv->analysis,  PS_LIST_TAIL, "PSPHOT.DETECTIONS", PS_META_REPLACE | PS_DATA_UNKNOWN, "Merged source list", detections);
    psMetadataAddPtr(refConv->analysis, PS_LIST_TAIL, "PSPHOT.DETECTIONS", PS_META_REPLACE | PS_DATA_UNKNOWN, "Merged source list", detections);
    psFree(detections);                    // Drop reference

    int footprint = psMetadataLookupS32(NULL, recipe, "STAMP.FOOTPRINT"); // Stamp half-size
    int stride = psMetadataLookupS32(NULL, recipe, "STRIDE"); // Size of convolution patches
    float regionSize = psMetadataLookupF32(NULL, recipe, "REGION.SIZE"); // Size of iso-kernel regs
    float spacing = psMetadataLookupF32(NULL, recipe, "STAMP.SPACING"); // Typical stamp spacing
    float threshold = psMetadataLookupF32(NULL, recipe, "STAMP.THRESHOLD"); // Threshold for stmps

    const char *typeStr = psMetadataLookupStr(NULL, recipe, "KERNEL.TYPE"); // Kernel type
    psAssert(typeStr, "We put it here in ppSubArguments.c");
    pmSubtractionKernelsType type = pmSubtractionKernelsTypeFromString(typeStr); // Type of kernel
    if (type == PM_SUBTRACTION_KERNEL_NONE) {
        psError(PPSUB_ERR_ARGUMENTS, true, "Unrecognised kernel type: %s", typeStr);
        return false;
    }

    int size = psMetadataLookupS32(NULL, recipe, "KERNEL.SIZE"); // Kernel half-size
    int order = psMetadataLookupS32(NULL, recipe, "SPATIAL.ORDER"); // Spatial polynomial order
    psVector *widths = psMetadataLookupPtr(NULL, recipe, "ISIS.WIDTHS"); // ISIS Gaussian widths
    psVector *orders = psMetadataLookupPtr(NULL, recipe, "ISIS.ORDERS"); // ISIS Polynomial orders

    int inner = psMetadataLookupS32(NULL, recipe, "INNER"); // Inner radius
    int ringsOrder = psMetadataLookupS32(NULL, recipe, "RINGS.ORDER"); // RINGS polynomial order
    int binning = psMetadataLookupS32(NULL, recipe, "SPAM.BINNING"); // Binning for SPAM kernel
    float penalty = psMetadataLookupF32(NULL, recipe, "PENALTY"); // Penalty for wideness
    int iter = psMetadataLookupS32(NULL, recipe, "ITER"); // Rejection iterations
    float rej = psMetadataLookupF32(NULL, recipe, "REJ"); // Rejection threshold
    float kernelErr = psMetadataLookupF32(NULL, recipe, "KERNEL.ERR"); // Relative systematic error in kernel
    float normFrac = psMetadataLookupF32(NULL, recipe, "NORM.FRAC"); // Fraction of window for normalisn windw
    float sysErr = psMetadataLookupF32(NULL, recipe, "SYS.ERR"); // Relative systematic error in images
    float skyErr = psMetadataLookupF32(NULL, recipe, "SKY.ERR"); // Additional error in sky
    float covarFrac = psMetadataLookupF32(NULL, recipe, "COVAR.FRAC"); // Fraction for covariance calculation

    float badFrac = psMetadataLookupF32(NULL, recipe, "BADFRAC"); // Maximum bad fraction
    float poorFrac = psMetadataLookupF32(&mdok, recipe, "POOR.FRACTION"); // Fraction for "poor"

    // Options which control the generation of optimal parameters
    bool optimum = psMetadataLookupBool(&mdok, recipe, "OPTIMUM"); // Derive optimum parameters?
    float optMin = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.MIN"); // Minimum width for search
    float optMax = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.MAX"); // Maximum width for search
    float optStep = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.STEP"); // Step for search
    psVector *optWidths = NULL;         // Vector with FWHMs for optimum search
    if (optimum) {
        optWidths = psVectorCreate(optWidths, optMin, optMax, optStep, PS_TYPE_F32);
    }
    int optOrder = psMetadataLookupS32(&mdok, recipe, "OPTIMUM.ORDER"); // Order for search
    float optThresh = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.TOL"); // Tolerance for search

    psImageMaskType maskVal = pmConfigMaskGet("MASK.VALUE", config); // Bits to mask in inputs
    psImageMaskType maskPoor = pmConfigMaskGet("CONV.POOR", config); // Bits to mask for poor pixels
    psImageMaskType maskBad = pmConfigMaskGet("BLANK", config); // Bits to mask for bad pixels

    bool dual = psMetadataLookupBool(&mdok, recipe, "DUAL"); // Dual convolution?
    pmSubtractionMode subMode;          // Subtraction mode
    if (dual) {
        subMode = PM_SUBTRACTION_MODE_DUAL;
    } else {
        int convolve = psMetadataLookupS32(NULL, config->arguments, "-convolve"); // Image number to convolve
        switch (convolve) {
          case 0:
            subMode = PM_SUBTRACTION_MODE_UNSURE;
            break;
          case 1:
            subMode = PM_SUBTRACTION_MODE_1;
            break;
          case 2:
            subMode = PM_SUBTRACTION_MODE_2;
            break;
          default:
            psError(PPSUB_ERR_ARGUMENTS, false, "Invalid value for -convolve");
            return false;
        }
    }

    if (!subScaleKernel(data, widths, &size, &footprint)) {
        psError(PPSUB_ERR_DATA, false, "Unable to scale kernel parameters");
        return false;
    }

    int threads = psMetadataLookupS32(NULL, config->arguments, "-threads"); // Number of threads
    if (threads > 0) {
        pmSubtractionThreadsInit();
    }

    if (inRO->covariance) {
        psImageCovarianceTruncate(inRO->covariance, COVAR_FRAC);
    }
    if (refRO->covariance) {
        psImageCovarianceTruncate(refRO->covariance, COVAR_FRAC);
    }

    // Data doesn't exist until it's created in pmSubtraction...
    inConv->data_exists = inConv->parent->data_exists = inConv->parent->parent->data_exists = false;
    refConv->data_exists = refConv->parent->data_exists = refConv->parent->parent->data_exists = false;

    // Match the PSFs
    bool success = false;               // Operation was successful?
    if (kernelRO) {
        success = pmSubtractionMatchPrecalc(inConv, refConv, inRO, refRO, kernelRO->analysis,
                                            stride, kernelErr, covarFrac, maskVal, maskBad, maskPoor,
                                            poorFrac, badFrac);
    } else {
        success = pmSubtractionMatch(inConv, refConv, inRO, refRO, footprint, stride, regionSize,
                                     spacing, threshold, detections ? detections->allSources : NULL,
                                     data->stamps, type, size, order, widths, orders, inner, ringsOrder,
                                     binning, penalty, optimum, optWidths, optOrder, optThresh, iter, rej,
                                     normFrac, sysErr, skyErr, kernelErr, covarFrac,
                                     maskVal, maskBad, maskPoor, poorFrac, badFrac, subMode);
    }

# ifdef TESTING
    // XXX for testing
    psphotSaveImage (NULL, refRO->image,    "refRO.im.fits");
    psphotSaveImage (NULL, refRO->variance, "refRO.wt.fits");
    psphotSaveImage (NULL, refRO->mask,     "refRO.mk.fits");

    psphotSaveImage (NULL, inRO->image,    "inRO.im.fits");
    psphotSaveImage (NULL, inRO->variance, "inRO.wt.fits");
    psphotSaveImage (NULL, inRO->mask,     "inRO.mk.fits");

    psphotSaveImage (NULL, inConv->image,    "inConv.im.fits");
    psphotSaveImage (NULL, inConv->variance, "inConv.wt.fits");
    psphotSaveImage (NULL, inConv->mask,     "inConv.mk.fits");

    psphotSaveImage (NULL, refConv->image,    "refConv.im.fits");
    psphotSaveImage (NULL, refConv->variance, "refConv.wt.fits");
    psphotSaveImage (NULL, refConv->mask,     "refConv.mk.fits");
# endif

    psFree(optWidths);
    pmSubtractionThreadsFinalize();

    if (!success) {
        psErrorCode error = psErrorCodeLast(); // Error code
        if (error == PM_ERR_STAMPS) {
            psErrorStackPrint(stderr, "Unable to find stamps");
            psWarning("Unable to find stamps --- suspect bad data quality.");
            ppSubDataQuality(data, error, PPSUB_FILES_ALL);
            return true;
        } else if (error == PM_ERR_SMALL_AREA) {
            psErrorStackPrint(stderr, "Insufficient area for PSF matching");
            psWarning("Insufficient area for PSF matching --- suspect bad data quality.");
            ppSubDataQuality(data, error, PPSUB_FILES_ALL);
            return true;
        } else {
            psError(PPSUB_ERR_DATA, false, "Unable to match images.");
            return false;
        }
    }

    // Need to be careful with the normalisation
    // We will normalise everything to the normalisation of the *input* image
    {
        // Since the entries are MULTI, we have to retrieve them differently
        psMetadataIterator *regIter = psMetadataIteratorAlloc(inConv->analysis, PS_LIST_HEAD,
                                                              "^" PM_SUBTRACTION_ANALYSIS_REGION "$");
        psMetadataIterator *modeIter = psMetadataIteratorAlloc(inConv->analysis, PS_LIST_HEAD,
                                                              "^" PM_SUBTRACTION_ANALYSIS_MODE "$");
        psMetadataIterator *normIter = psMetadataIteratorAlloc(inConv->analysis, PS_LIST_HEAD,
                                                              "^" PM_SUBTRACTION_ANALYSIS_NORM "$");
        psMetadataItem *regItem;        // Item with region
        while ((regItem = psMetadataGetAndIncrement(regIter))) {
            psAssert(regItem->type == PS_DATA_REGION && regItem->data.V, "Expect region type");
            psRegion *region = regItem->data.V; // Region of interest
            psMetadataItem *modeItem = psMetadataGetAndIncrement(modeIter); // Item with mode
            psAssert(modeItem && modeItem->type == PS_TYPE_S32, "Expect subtraction mode");
            pmSubtractionMode mode = modeItem->data.S32; // Subtraction mode
            psMetadataItem *normItem = psMetadataGetAndIncrement(normIter); // Item with normalisation
            psAssert(normItem && normItem->type == PS_TYPE_F32 && isfinite(normItem->data.F32),
                     "Expect normalisation");
            float norm = normItem->data.F32; // Normalisation

            switch (mode) {
              case PM_SUBTRACTION_MODE_1: // Convolved the input to match template
              case PM_SUBTRACTION_MODE_DUAL: // Convolved both; template should have flux conserved
                psLogMsg("ppSub", PS_LOG_INFO, "Correcting image for normalisation of %f\n", norm);
                normaliseRegion(inConv->image, region, 1.0 / norm);
                normaliseRegion(refConv->image, region, 1.0 / norm);
                normaliseRegion(inConv->variance, region, 1.0 / PS_SQR(norm));
                normaliseRegion(refConv->variance, region, 1.0 / PS_SQR(norm));
                break;
              case PM_SUBTRACTION_MODE_2:       // Convolved the template to match input
                // We're already happy!
                psLogMsg("ppSub", PS_LOG_INFO, "Image normalisation is correct\n");
                break;
              default:
                psAbort("Invalid subtraction mode: %x", mode);
            }
        }
        psFree(regIter);
        psFree(modeIter);
        psFree(normIter);
    }


    pmConceptsCopyFPA(inConv->parent->parent->parent, inRO->parent->parent->parent, true, true);
    pmConceptsCopyFPA(refConv->parent->parent->parent, refRO->parent->parent->parent, true, true);

    if (inConv->covariance) {
        psImageCovarianceTruncate(inConv->covariance, COVAR_FRAC);
    }
    if (refConv->covariance) {
        psImageCovarianceTruncate(refConv->covariance, COVAR_FRAC);
    }

    if (inConv->variance) {
        psImageCovarianceTransfer(inConv->variance, inConv->covariance);
    }
    if (refConv->variance) {
        psImageCovarianceTransfer(refConv->variance, refConv->covariance);
    }

    return true;
}
