/** @file pswarpLoop.c
 *
 *  @brief
 *
 *  @ingroup pswarp
 *
 *  @author IfA
 *  @version $Revision: 1.38 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2009-02-13 21:54:32 $
 *  Copyright 2009 Institute for Astronomy, University of Hawaii
 */

#include "pswarp.h"
#include <ppStats.h>
#include "pswarpFileNames.h"            // Lists of file rules used at different stages

#define WCS_NONLIN_TOL 0.001            // Non-linear tolerance for header WCS
#define PSPHOT_FIND_PSF 1               // Use psphot's findPSF function?
#define TESTING 0                       // Testing output?


// XXX these are generic functions which should be moved to psModules
// Activate a list of files
static void fileActivation(pmConfig *config, // Configuration
                           char **files, // Files to turn on/off
                           bool state   // Activation state
    )
{
    for (int i = 0; files[i] != NULL; i++) {
        pmFPAfileActivate(config->files, state, files[i]);
    }
    return;
}

/**
 * Run down the FPA hierarchy, checking files
 */
static bool ioChecksBefore(pmConfig *config ///< Configuration
                           )
{
    pmFPAview *view = pmFPAviewAlloc(0); ///< View for checking
    bool status = true;                 ///< Status of checks
    status &= pmFPAfileIOChecks(config, view, PM_FPA_BEFORE);
    view->chip = 0;
    status &= pmFPAfileIOChecks(config, view, PM_FPA_BEFORE);
    view->cell = 0;
    status &= pmFPAfileIOChecks(config, view, PM_FPA_BEFORE);
    psFree(view);
    return status;
}

/**
 * Run up the FPA hierarchy, checking files
 */
static bool ioChecksAfter(pmConfig *config ///< Configuration
                          )
{
    pmFPAview *view = pmFPAviewAlloc(0); ///< View for checking
    view->chip = view->cell = 0;
    bool status = true;                 ///< Status of checks
    status &= pmFPAfileIOChecks(config, view, PM_FPA_AFTER);
    view->cell = -1;
    status &= pmFPAfileIOChecks(config, view, PM_FPA_AFTER);
    view->chip = -1;
    status &= pmFPAfileIOChecks(config, view, PM_FPA_AFTER);
    psFree(view);
    return status;
}


/**
 * Loop over the inputs, warp them to the output skycell and then write out the output.
 */
bool pswarpLoop(pmConfig *config)
{
    bool status;

    const char *skyCamera = psMetadataLookupStr(NULL, config->arguments,
                                                "SKYCELL.CAMERA");  ///< Name of camera for skycell
    pmConfigCamerasCull(config, skyCamera);
    pmConfigRecipesCull(config, "PSWARP,PPSTATS,PSPHOT,PSASTRO,MASKS,JPEG");


    // load the recipe
    psMetadata *recipe = psMetadataLookupPtr (&status, config->recipes, PSWARP_RECIPE);
    if (!recipe) {
        psError(PSPHOT_ERR_CONFIG, false, "missing recipe %s", PSWARP_RECIPE);
        return false;
    }

    if (!pswarpSetMaskBits(config)) {
        psError(PS_ERR_IO, false, "failed to set mask bits");
        return NULL;
    }

    // output mask bits
    psImageMaskType maskValue = psMetadataLookupImageMask(&status, recipe, "MASK.OUTPUT");
    psAssert (status, "MASK.OUTPUT was not defined");

    // select the input data sources
    pmFPAfile *input = psMetadataLookupPtr(NULL, config->files, "PSWARP.INPUT");
    if (!input) {
        psError(PSWARP_ERR_CONFIG, true, "Can't find input data!\n");
        return false;
    }

    // use the external astrometry source if supplied
    pmFPAfile *astrom = psMetadataLookupPtr(NULL, config->files, "PSWARP.ASTROM");
    if (!astrom) {
        astrom = input;
    }

    if (astrom->camera != input->camera) {
        psError(PS_ERR_UNKNOWN, true, "Input camera and astrometry camera do not match.");
        return false;
    }

    // select the output readout
    pmFPAview *view = pmFPAviewAlloc(0);
    view->chip = 0;
    view->cell = 0;
    view->readout = 0;
    pmReadout *output = pmFPAfileThisReadout(config->files, view, "PSWARP.OUTPUT");
    if (!output) {
        psError(PSWARP_ERR_CONFIG, true, "Can't find output data!\n");
        return false;
    }
    psFree (view);

    bool mdok;                          ///< Status of MD lookup
    const char *statsName = psMetadataLookupStr(&mdok, config->arguments, "STATS"); ///< Filename for statistics
    psMetadata *stats = NULL;           ///< Container for statistics
    FILE *statsFile = NULL;             ///< File stream for statistics
    if (mdok && statsName && strlen(statsName) > 0) {
        psString resolved = pmConfigConvertFilename(statsName, config, true, true);
        statsFile = fopen(resolved, "w");
        if (!statsFile) {
            psError(PS_ERR_IO, true, "Unable to open statistics file %s for writing.\n", resolved);
            psFree(resolved);
            return false;
        }
        psFree(resolved);
        stats = psMetadataAlloc();
        psMetadataAddS32(stats, PS_LIST_TAIL, "QUALITY", 0, "No problems", 0);
    }

    // Turn all skycell files on to generate them, and then turn them off for the loop over the input images
    // the input, which is in a different format.
    {
        fileActivation(config, detectorFiles, false);
        fileActivation(config, photFiles, false);
        fileActivation(config, independentFiles, false);
        fileActivation(config, skycellFiles, true);
        ioChecksBefore(config);
        fileActivation(config, skycellFiles, false);
    }

    // Read the input astrometry
    // XXX rather than use the activations here, this should just explicitly loop over the desired filerule
    {
        pmFPAfileActivate(config->files, true, "PSWARP.ASTROM");

        pmChip *chip;
        pmFPAview *view = pmFPAviewAlloc(0);
        pmFPAfileIOChecks (config, view, PM_FPA_BEFORE);
        while ((chip = pmFPAviewNextChip (view, input->fpa, 1)) != NULL) {
            psTrace ("pswarp", 4, "Chip %d: %x %x\n", view->chip, chip->file_exists, chip->process);
            if (!chip->process || !chip->file_exists) { continue; }
            pmFPAfileIOChecks(config, view, PM_FPA_BEFORE);
            pmCell *cell;
            while ((cell = pmFPAviewNextCell (view, input->fpa, 1)) != NULL) {
                psTrace ("pswarp", 4, "Cell %d: %x %x\n", view->cell, cell->file_exists, cell->process);
                if (!cell->process || !cell->file_exists) { continue; }
                pmFPAfileIOChecks (config, view, PM_FPA_BEFORE);
                pmFPAfileIOChecks (config, view, PM_FPA_AFTER);
            }
            pmFPAfileIOChecks (config, view, PM_FPA_AFTER);
        }
        pmFPAfileIOChecks (config, view, PM_FPA_AFTER);
        psFree(view);

        fileActivation(config, detectorFiles, true);
        pmFPAfileActivate(config->files, false, "PSWARP.ASTROM");
    }

    // Turn on the source output --- we need to get rid of these so that we can measure the PSF
    pmFPAfileActivate(config->files, true, "PSWARP.OUTPUT.SOURCES");

    // Don't care about the skycell anymore --- we've read it, and that's all we need to do.
    pmFPAfileActivate(config->files, false, "PSWARP.SKYCELL");
    view = pmFPAviewAlloc(0);

    // find the FPA phu
    bool bilevelAstrometry = false;
    pmHDU *phu = pmFPAviewThisPHU(view, astrom->fpa);
    if (phu) {
        char *ctype = psMetadataLookupStr(NULL, phu->header, "CTYPE1");
        if (ctype) {
            bilevelAstrometry = !strcmp (&ctype[4], "-DIS");
        }
    }
    if (bilevelAstrometry) {
        if (!pmAstromReadBilevelMosaic(input->fpa, phu->header)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to read bilevel mosaic astrometry for input FPA.");
            psFree(view);
            psFree(stats);
            return false;
        }
    }

    psList *cells = psListAlloc(NULL);  // List of cells, for concepts averaging

    // files associated with the science image
    pmFPAfileIOChecks (config, view, PM_FPA_BEFORE);

    pmChip *chip;
    while ((chip = pmFPAviewNextChip (view, input->fpa, 1)) != NULL) {
        psTrace ("pswarp", 4, "Chip %d: %x %x\n", view->chip, chip->file_exists, chip->process);
        if (!chip->process || !chip->file_exists) { continue; }
        pmFPAfileIOChecks (config, view, PM_FPA_BEFORE);

        // read WCS data from the corresponding header
        pmHDU *hdu = pmFPAviewThisHDU (view, astrom->fpa);
        if (bilevelAstrometry) {
            if (!pmAstromReadBilevelChip (chip, hdu->header)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to read bilevel chip astrometry for input FPA.");
                psFree(view);
                psFree(stats);
                return false;
            }
        } else {
            // we use a default FPA pixel scale of 1.0
            if (!pmAstromReadWCS (input->fpa, chip, hdu->header, 1.0)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to read WCS astrometry for input FPA.");
                psFree(view);
                psFree(stats);
                return false;
            }
        }

        pmCell *cell;
        while ((cell = pmFPAviewNextCell (view, input->fpa, 1)) != NULL) {
            psTrace ("pswarp", 4, "Cell %d: %x %x\n", view->cell, cell->file_exists, cell->process);
            if (!cell->process || !cell->file_exists) { continue; }
            pmFPAfileIOChecks (config, view, PM_FPA_BEFORE);

            psListAdd(cells, PS_LIST_TAIL, cell);

            // process each of the readouts
            pmReadout *readout;
            while ((readout = pmFPAviewNextReadout(view, input->fpa, 1)) != NULL) {
                pmFPAfileIOChecks(config, view, PM_FPA_BEFORE);
                if (!readout->data_exists) {
                    continue;
                }

                // Copy the sources from the astrometry carrier to the input, so they can be accessed by
                // pswarpTransformReadout
                pmReadout *astromRO = pmFPAviewThisReadout(view, astrom->fpa); // Readout for astrometry
                psArray *sources = psMetadataLookupPtr(&mdok, astromRO->analysis,
                                                       "PSPHOT.SOURCES"); // Sources from astrometry
                if (sources) {
                    psMetadataAddPtr(readout->analysis, PS_LIST_TAIL, "PSPHOT.SOURCES", PS_DATA_ARRAY,
                                     "Sources from input astrometry", sources);
                }

                pswarpTransformReadout(output, readout, config);

                pmFPAfileIOChecks(config, view, PM_FPA_AFTER);
            }
            pmFPAfileIOChecks(config, view, PM_FPA_AFTER);
        }
        pmFPAfileIOChecks(config, view, PM_FPA_AFTER);
    }

    if (!output->data_exists) {
        psWarning("No overlap between input and skycell.");
        if (stats) {
            psMetadataAddS32(stats, PS_LIST_TAIL, "QUALITY", PS_META_REPLACE,
                             "No overlap between input and skycell", PSWARP_ERR_NO_OVERLAP);
        }
        psphotFilesActivate(config, false);
        psFree(cells);
        psFree(view);
        goto DONE;
    }

    pmCell *outCell = output->parent;   ///< Output cell
    pmChip *outChip = outCell->parent;  ///< Output chip
    pmFPA *outFPA = outChip->parent;    ///< Output FP

    if (!pswarpPixelsLit(output, stats, config)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to calculate pixel regions.");
        psFree(cells);
        psFree(view);
        return false;
    }

    // Correct image for change in the plate scale
    {
        psAssert(input && input->fpa && input->fpa->toSky, "Require astrometry for input");
        psAssert(outFPA && outFPA && outFPA->toSky, "Require astrometry for output");

        double inScale = input->fpa->toSky->Xs + input->fpa->toSky->Ys; // Plate scale for input
        double outScale = outFPA->toSky->Xs + outFPA->toSky->Ys; // Plate scale for output
        float correction = PS_SQR(outScale / inScale); // Correction factor to apply to image
        psLogMsg("pswarp", PS_LOG_INFO, "Correcting flux by %f to account for pixel scales", correction);

        psBinaryOp(output->image, output->image, "*", psScalarAlloc(correction, PS_TYPE_F32));
        if (output->variance) {
            psBinaryOp(output->variance, output->variance, "*",
                       psScalarAlloc(PS_SQR(correction), PS_TYPE_F32));
        }
    }


    // Set covariance matrix for output
    {
        psList *covariances = psMetadataLookupPtr(&mdok, output->analysis,
                                                  PSWARP_ANALYSIS_COVARIANCES); // Covariance matrices
        psAssert(covariances, "Should be there");
        psArray *covars = psListToArray(covariances); // Array of covariance matrices
        output->covariance = psImageCovarianceAverage(covars);
        psFree(covars);
        psMetadataRemoveKey(output->analysis, PSWARP_ANALYSIS_COVARIANCES);
        psImageCovarianceTransfer(output->variance, output->covariance);
    }

    if (!pmConceptsAverageCells(outCell, cells, NULL, NULL, false)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to average cell concepts.");
        psFree(stats);
        psFree(cells);
        psFree(view);
        return false;
    }
    psFree(cells);

    psRegion *trimsec = psMetadataLookupPtr(NULL, outCell->concepts, "CELL.TRIMSEC"); ///< Trim section
    trimsec->x0 = trimsec->x1 = trimsec->y0 = trimsec->y1 = 0; ///< All pixels

    if (!psMetadataCopy(outFPA->concepts, input->fpa->concepts)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to copy FPA concepts from input to output.");
        psFree(stats);
        psFree(view);
        return false;
    }

    pmHDU *hdu = outFPA->hdu;           ///< HDU for the output warped image

    // Copy header from target
    {
        pmFPAview *skyView = pmFPAviewAlloc(0); ///< View into skycell
        skyView->chip = skyView->cell = 0;
        pmCell *cell = pmFPAfileThisCell(config->files, skyView, "PSWARP.SKYCELL"); // Skycell cell
        psFree(skyView);
        pmHDU *skyHDU = pmHDUFromCell(cell); ///< HDU
        if (!skyHDU) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find skycell HDU.");
            psFree(view);
            return false;
        }
        hdu->header = psMetadataCopy(hdu->header, skyHDU->header);
    }

    pswarpVersionHeader(hdu->header);

    if (!pmAstromWriteWCS(hdu->header, outFPA, outChip, WCS_NONLIN_TOL)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to generate WCS header.");
        psFree(stats);
        return false;
    }

    pmFPAfileIOChecks(config, view, PM_FPA_AFTER);

    // Done with the detector side of things
    fileActivation(config, detectorFiles, false);
    fileActivation(config, independentFiles, false);

    // We need a new PSF model for the warped frame.  It would be good to generate this analytically, but
    // that's going to be tricky.  We have a list of sources, so we use those to redetermine the PSF model.

    if (psMetadataLookupBool(&mdok, config->arguments, "PSF")) {
        fileActivation(config, photFiles, true);
        ioChecksBefore(config);

        // supply the readout and fpa of interest to psphot
        pmFPAfile *photFile = psMetadataLookupPtr(NULL, config->files, "PSPHOT.INPUT");
        pmFPACopy(photFile->fpa, outFPA);

        pmFPAview *view = pmFPAviewAlloc(0); ///< View into skycell
        view->chip = view->cell = view->readout = 0;

        // grab the sources of interest from the storage location (pmFPAfile PSPHOT.INPUT.CMF)
        psArray *sources = psphotLoadPSFSources (config, view);
        if (!sources) {
            psError(PS_ERR_UNKNOWN, false, "No sources supplied to measure PSF");
            return false;
        }

        // measure the PSF using these sources
        if (!psphotReadoutFindPSF(config, view, sources)) {
            // This is likely a data quality issue
            // XXX Split into multiple cases using error codes?
            psErrorStackPrint(stderr, "Unable to determine PSF");
            psWarning("Unable to determine PSF --- suspect bad data quality.");
            if (stats && psMetadataLookupS32(NULL, stats, "QUALITY") == 0) {
                psMetadataAddS32(stats, PS_LIST_TAIL, "QUALITY", PS_META_REPLACE,
                                 "Unable to determine PSF", psErrorCodeLast());
            }
            psErrorClear();
            psphotFilesActivate(config, false);
        }

// XXX EAM : put this in a visualization function
#if (TESTING)
        {
            #define PSF_SIZE 20         ///< Half-size of PSF
            #define PSF_FLUX 10000      ///< Central flux for PSF
            pmChip *photChip = pmFPAviewThisChip(view, photFile->fpa);
            pmPSF *psf = psMetadataLookupPtr(NULL, photChip->analysis, "PSPHOT.PSF");
            psImage *image = psImageAlloc(2 * PSF_SIZE + 1, 2 * PSF_SIZE + 1, PS_TYPE_F32);
            psImageInit(image, 0);
            pmModel *model = pmModelFromPSFforXY(psf, PSF_SIZE, PSF_SIZE, PSF_FLUX);
            pmModelAdd(image, NULL, model, PM_MODEL_OP_FULL, 0);
            psFree(model);
            psFits *fits = psFitsOpen("psf.fits", "w");
            psFitsWriteImage(fits, NULL, image, 0, NULL);
            psFitsClose(fits);
            psFree(image);
        }
#endif

        psFree(view);
    }

    fileActivation(config, skycellFiles, true);

    // Perform statistics on the output image
    if (stats) {
        if (!ppStatsFPA(stats, output->parent->parent->parent, view, maskValue, config)) {
            psWarning("Unable to perform statistics on warped image.");
        }
    }

    // Add MD5 information for readout
    const char *chipName = psMetadataLookupStr(NULL, output->parent->parent->concepts, "CHIP.NAME");
    const char *cellName = psMetadataLookupStr(NULL, output->parent->concepts, "CELL.NAME");
    psString headerName = NULL; ///< Header name for MD5
    psStringAppend(&headerName, "MD5_%s_%s_%d", chipName, cellName, view->readout);
    psVector *md5 = psImageMD5(output->image); ///< md5 hash
    psString md5string = psMD5toString(md5); ///< String
    psFree(md5);
    psMetadataAddStr(hdu->header, PS_LIST_TAIL, headerName, PS_META_REPLACE,
                     "Image MD5", md5string);
    psFree(md5string);
    psFree(headerName);
    psFree(view);

    // Ensure everything is written out, at every level
    ioChecksAfter(config);

    // Now done with the skycell side of things

 DONE:

    // Write out summary statistics
    if (stats) {
        psMetadataAddF32(stats, PS_LIST_TAIL, "DT_WARP", 0, "Time for warp completion",
                         psTimerMark("pswarp"));

        const char *statsMDC = psMetadataConfigFormat(stats);
        if (!statsMDC) {
            psError(PS_ERR_IO, false, "Unable to get statistics MDC file.\n");
            psFree(stats);
            fclose(statsFile);
            return false;
        }
        fprintf(statsFile, "%s", statsMDC);
        psFree((void*)statsMDC);
        fclose(statsFile);
        pmConfigRunFilenameAddWrite(config, "STATS", statsName);

        psFree(stats);
    }

    // Dump configuration
    psString dump_file = psMetadataLookupStr(&status, config->arguments, "DUMP_CONFIG");
    if (dump_file) {
        pmConfigDump(config, dump_file);
    }

    return true;
}
