#include "pswarp.h"
#include <ppStats.h>

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

// Lists of file rules for the detectors and skycells, and an independent list of everything else
// Lists must end with a NULL
static char *detectorFiles[] = { "PSWARP.INPUT", "PSWARP.MASK", "PSWARP.WEIGHT", NULL };
static char *skycellFiles[] = { "PSWARP.OUTPUT", "PSWARP.OUTPUT.MASK", "PSWARP.OUTPUT.WEIGHT", NULL };
static char *photFiles[] = { "PSPHOT.INPUT", "PSPHOT.OUTPUT",
#if PSPHOT_FIND_PSF
                             "PSPHOT.INPUT.CMF",
#endif
                             "PSPHOT.RESID",
                             "PSPHOT.BACKMDL", "PSPHOT.BACKMDL.STDEV", "PSPHOT.BACKGND", "PSPHOT.BACKSUB",
                             "PSPHOT.PSF.SAVE", "SOURCE.PLOT.MOMENTS", "SOURCE.PLOT.PSFMODEL",
                             "SOURCE.PLOT.APRESID", NULL };
static char *independentFiles[] = { "PSWARP.ASTROM", // Read independently from the input pixels
                                    "PSWARP.SKYCELL", // Don't care about the skycell once we have its WCS
                                    "PSWARP.OUTPUT.SOURCES", // Save these independently so we can do the PSF
                                    NULL };

// 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;

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

    // output mask bits
    psMaskType maskValue = psMetadataLookupU8(&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();
    }

    // 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
    {
        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);
    }

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

    if (!pswarpPixelFraction(output, stats, config)) {
        // Don't write output images, and don't bother about anything else
        output->data_exists = outCell->data_exists = outChip->data_exists = false;
        psFree(cells);
        psFree(view);
        goto COMPLETED;
    }

    // Set variance factor
    {
        float varFactor = psMetadataLookupF32(NULL, output->analysis, PSWARP_ANALYSIS_VARFACTOR);
        psAssert(isfinite(varFactor), "Should be something here.");
        long goodPix = psMetadataLookupS64(NULL, output->analysis, PSWARP_ANALYSIS_GOODPIX);
        psAssert(goodPix > 0, "Should be something here.");
        varFactor /= goodPix;

        psMetadataItem *vfItem = psMetadataLookup(outCell->concepts, "CELL.VARFACTOR"); // Item to update
        psAssert(vfItem && vfItem->type == PS_TYPE_F32, "Concept should be as specified.");
        if (!isfinite(vfItem->data.F32)) {
            vfItem->data.F32 = varFactor;
        } else {
            vfItem->data.F32 *= varFactor;
        }
    }

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

    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 could use those to redetermine the PSF model.
    // Until Gene makes the necessary adaptations to psphot, we will simply redetermine the PSF model from
    // scratch.
    if (psMetadataLookupBool(&mdok, config->arguments, "PSF")) {
        fileActivation(config, photFiles, true);
        ioChecksBefore(config);

        // Tweak recipe
        psMetadata *psphotRecipe = psMetadataLookupPtr(NULL, config->recipes,
                                                       PSPHOT_RECIPE); // Recipe for psphot
        if (!psphotRecipe) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find %s recipe.", PSPHOT_RECIPE);
            return false;
        }

        psMetadataAddStr(psphotRecipe, PS_LIST_TAIL, "BREAK_POINT", PS_META_REPLACE,
                         "Break point for psphot operations", "PSFMODEL");
        psMetadataAddBool(psphotRecipe, PS_LIST_TAIL, PSPHOT_RECIPE_PSF_FAKE_ALLOW, PS_META_REPLACE,
                          "Don't allow psphot to fit a fake PSF!", false);

#if !PSPHOT_FIND_PSF
        // Get rid of the transformed sources so that we can measure the PSF
        psMetadataRemoveKey(output->analysis, "PSPHOT.SOURCES");
#endif

        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;

#if PSPHOT_FIND_PSF
        // Need to adjust the weight --- the main operation in psphotReadoutFindPSF is not determining the
        // signficance of sources, but looking at the significance of individual pixels.
        // We can adjust the weight directly since this is a deep copy
        pmReadout *photRO = pmFPAviewThisReadout(view, photFile->fpa);
        float vf = psMetadataLookupF32(NULL, photRO->parent->concepts, "CELL.VARFACTOR"); // Variance factor
        psBinaryOp(photRO->weight, photRO->weight, "*", psScalarAlloc(vf, PS_TYPE_F32));

        if (!psphotReadoutFindPSF(config, view)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to determine PSF for warped image.");
            return false;
        }
#else
        if (!psphotReadout(config, view)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to determine PSF for warped image.");
            return false;
        }
#endif

#ifdef 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

    COMPLETED:
    // 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);

        psFree(stats);
    }

    return true;
}
