#include <stdio.h>
#include <math.h>
#include <string.h>
#include <pslib.h>
#include <psmodules.h>

#include "ppStack.h"

//#define TESTING                         // Enable debugging output

//#define ASTROMETRY                    // Correct astrometry?

#ifdef TESTING
// Size of fake image; set by hand because it's trouble to get it from other places
#define FAKE_COLS 4861
#define FAKE_ROWS 4913
#endif

#ifdef TESTING
// Dump matches to a file
static void dumpMatches(const char *filename, // File to which to dump
                        int num,        // Number of inputs
                        psArray *matches, // Star matches
                        psVector *zp,   // Zero points
                        psVector *trans // Transparencies
                        )
{
    FILE *outMatches = fopen(filename, "w"); // Output matches
    psVector *mag = psVectorAlloc(num, PS_TYPE_F32); // Magnitudes for each star
    psVector *magErr = psVectorAlloc(num, PS_TYPE_F32); // Errors for each star
    for (int i = 0; i < matches->n; i++) {
        pmSourceMatch *match = matches->data[i]; // Match of interest
        psVectorInit(mag, NAN);
        psVectorInit(magErr, NAN);
        for (int j = 0; j < match->num; j++) {
            if (match->mask->data.PS_TYPE_VECTOR_MASK_DATA[j]) {
                continue;
            }
            int index = match->image->data.U32[j]; // Image index
            mag->data.F32[index] = match->mag->data.F32[j] - zp->data.F32[index];
            if (trans) {
                mag->data.F32[index] -= trans->data.F32[index];
            }
            magErr->data.F32[index] = match->magErr->data.F32[j];
        }
        for (int j = 0; j < num; j++) {
            fprintf(outMatches, "%f (%f) ", mag->data.F32[j], magErr->data.F32[j]);
        }
        fprintf(outMatches, "\n");
    }
    psFree(mag);
    psFree(magErr);
    fclose(outMatches);
    return;
}
#endif


bool ppStackSourcesTransparency(ppStackOptions *options, const pmFPAview *view, const pmConfig *config)
{
    PS_ASSERT_PTR_NON_NULL(options, false);
    PS_ASSERT_PTR_NON_NULL(view, false);
    PS_ASSERT_PTR_NON_NULL(config, false);

    if (!options->matchZPs && !options->photometry) {
        options->norm = psVectorAlloc(options->num, PS_TYPE_F32);
        psVectorInit(options->norm, 0.0);
        return true;
    }

    psArray *sourceLists = options->sourceLists; // Source lists for each input
    psVector *inputMask = options->inputMask; // Mask for inputs

    PS_ASSERT_ARRAY_NON_NULL(sourceLists, false);
    PS_ASSERT_VECTOR_NON_NULL(inputMask, false);
    PS_ASSERT_VECTOR_TYPE(inputMask, PS_TYPE_U8, false);
    PS_ASSERT_VECTOR_SIZE(inputMask, sourceLists->n, false);

#if defined(TESTING) && 0
    {
        // Deliberately induce a major transparency difference
        psArray *sources = sourceLists->data[1]; // Sources to correct
        for (int i = 0; i < sources->n; i++) {
            pmSource *source = sources->data[i]; // Source of interest
            if (!source) {
                continue;
            }
            source->psfMag += 1.0;
#ifdef ASTROMETRY
            if (source->modelPSF) {
                source->modelPSF->params->data.F32[PM_PAR_XPOS] += 1.0;
                source->modelPSF->params->data.F32[PM_PAR_YPOS] += 1.0;
            }
            if (source->peak) {
                source->peak->xf += 1.0;
                source->peak->yf += 1.0;
            }
#endif
        }
    }
#endif

    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSTACK_RECIPE); // ppStack recipe
    psAssert(recipe, "We've thrown an error on this before.");

    float radius = psMetadataLookupF32(NULL, recipe, "ZP.RADIUS"); // Radius (pixels) for matching sources
    int iter1 = psMetadataLookupS32(NULL, recipe, "ZP.ITER.1"); // Maximum iterations for pass 1
    int iter2 = psMetadataLookupS32(NULL, recipe, "ZP.ITER.2"); // Maximum iterations for pass 2
    float tol = psMetadataLookupF32(NULL, recipe, "ZP.TOL"); // Tolerance for zero point iterations
    int transIter = psMetadataLookupS32(NULL, recipe, "ZP.TRANS.ITER"); // Iterations for transparency
    float transRej = psMetadataLookupF32(NULL, recipe, "ZP.TRANS.REJ");// Rejection threshold for transparency
    float transThresh = psMetadataLookupF32(NULL, recipe, "ZP.TRANS.THRESH"); // Threshold for transparency

    float starRej1 = psMetadataLookupF32(NULL, recipe, "ZP.STAR.REJ.1"); // Rejection threshold for stars
    float starSys1 = psMetadataLookupF32(NULL, recipe, "ZP.STAR.SYS.1"); // Estimated systematic error
    float starRej2 = psMetadataLookupF32(NULL, recipe, "ZP.STAR.REJ.2"); // Rejection threshold for stars
    float starSys2 = psMetadataLookupF32(NULL, recipe, "ZP.STAR.SYS.2"); // Estimated systematic error

    float starLimit = psMetadataLookupF32(NULL, recipe, "ZP.STAR.LIMIT"); // Limit on star rejection fraction

    float fracMatch = psMetadataLookupF32(NULL, recipe, "ZP.MATCH"); // Fraction of images to match for star

    psMetadata *airmassZP = psMetadataLookupMetadata(NULL, recipe, "ZP.AIRMASS"); // Airmass terms
    if (!airmassZP) {
        psError(PPSTACK_ERR_CONFIG, false, "Unable to find ZP.AIRMASS in recipe.");
        return false;
    }
    psMetadata *zpTargetMenu = psMetadataLookupMetadata(NULL, recipe, "ZP.TARGET"); // Target zero point terms
    if (!zpTargetMenu) {
        psError(PPSTACK_ERR_CONFIG, false, "Unable to find ZP.TARGET in recipe.");
        return false;
    }

    int num = options->num;             // Number of inputs
    psAssert(num == sourceLists->n, "Wrong number of source lists: %ld\n", sourceLists->n);

    psVector *zp = psVectorAlloc(num, PS_TYPE_F32); // Relative zero points for each image
    psVector *zpExp = psVectorAlloc(num, PS_TYPE_F32); // Measured zero points for each image (maybe)
    int zpExpNum = 0;                                  // Number of measured zero points
    const char *filter = NULL;          // Filter name
    float airmassTerm = NAN;            // Airmass term
    float zpTarget = NAN;               // Target zero point
    int numGoodImages = 0;              // Number of good images
    for (int i = 0; i < num; i++) {
        psArray *sources = sourceLists->data[i]; // Source list
        if (!sources || sources->n == 0) {
            psLogMsg("ppStack", PS_LOG_WARN, "Image %d has no sources for transparency measurement.", i);
            options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] |= PPSTACK_MASK_CAL;
            zp->data.F32[i] = NAN;
            continue;
        }
        numGoodImages++;

        pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPSTACK.INPUT", i); // File of interest

#if defined(TESTING) && 0
        pmReadout *fake = pmReadoutAlloc(NULL); // Fake readout
        pmPSF *psf = psMetadataLookupPtr(NULL, config->arguments, "PSF.TARGET"); // PSF for fake image
        pmReadoutFakeFromSources(fake, FAKE_COLS, FAKE_ROWS, sourceLists->data[i], 0,
                                 NULL, NULL, psf, 5, 0, false, true);
        psString name = NULL;
        psStringAppend(&name, "start_%03d.fits", i);
        pmStackVisualPlotTestImage(fake->image, name);
        psFits *fits = psFitsOpen(name, "w");
        psFree(name);
        psFitsWriteImage(fits, NULL, fake->image, 0, NULL);
        psFitsClose(fits);
        psFree(fake);
#endif

        float exptime = options->exposures->data.F32[i]; // Exposure time
        float airmass = psMetadataLookupF32(NULL, file->fpa->concepts, "FPA.AIRMASS"); // Airmass
        const char *expFilter = psMetadataLookupStr(NULL, file->fpa->concepts, "FPA.FILTER"); // Filter name
        zpExp->data.F32[i] = psMetadataLookupF32(NULL, file->fpa->concepts, "FPA.ZP"); // Exposure zero point
        psLogMsg("ppStack", PS_LOG_INFO,
                 "Image %d: %.2f sec exposure in %s at airmass %.2f with zero point %.2f",
                 i, exptime, expFilter, airmass, zpExp->data.F32[i]);
        if (!isfinite(exptime) || exptime == 0 || !isfinite(airmass) || airmass == 0 ||
            !expFilter || strlen(expFilter) == 0) {
            psError(PPSTACK_ERR_CONFIG, false,
                    "Unable to find exposure time (%f), airmass (%f) or filter (%s)",
                    exptime, airmass, expFilter);
            psFree(zp);
            return false;
        }
        if (isfinite(zpExp->data.F32[i])) {
            zpExp->data.F32[i] += 2.5 * log10(exptime);
            zpExpNum++;
        }

        if (!filter) {
            filter = expFilter;
            bool mdok;
            airmassTerm = psMetadataLookupF32(&mdok, airmassZP, filter);
            if (!mdok || !isfinite(airmassTerm)) {
                psError(PPSTACK_ERR_CONFIG, false,
                        "Unable to find airmass term (ZP.AIRMASS) for filter %s", filter);
                psFree(zp);
                return false;
            }
            zpTarget = psMetadataLookupF32(&mdok, zpTargetMenu, filter);
            if (!mdok || !isfinite(zpTarget)) {
                psError(PPSTACK_ERR_CONFIG, false,
                        "Unable to find target zero point (ZP.TARGET) for filter %s", filter);
                psFree(zp);
                return false;
            }
        } else if (strcmp(filter, expFilter) != 0) {
            psError(PPSTACK_ERR_CONFIG, false, "Filters don't match: %s vs %s", filter, expFilter);
            psFree(zp);
            return false;
        }

        zp->data.F32[i] = airmassTerm * airmass + 2.5 * log10(exptime);
    }

    if (numGoodImages == 0) {
        psLogMsg("ppStack", PS_LOG_WARN, "No images with sources to measure transparency.");
        options->quality = PPSTACK_ERR_REJECTED;
        psFree(zp);
        psFree(zpExp);
        return true;
    }
    if (numGoodImages == 1) {
        psArray *sources = NULL;        // Sources
        for (int i = 0; i < num && !sources; i++) {
            if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
                continue;
            }
            sources = sourceLists->data[i];
        }
        options->quality = PPSTACK_ERR_REJECTED;
        options->sources = psMemIncrRefCounter(sources);
        options->norm = psVectorAlloc(num, PS_TYPE_F32);
        psVectorInit(options->norm, 1.0);
        options->zp = NAN;
        psLogMsg("ppStack", PS_LOG_WARN, "Single image with sources --- cannot match transparency.");
        psFree(zp);
        psFree(zpExp);
        return true;
    }

    if (zpExpNum == numGoodImages) {
        for (int i = 0; i < num; i++) {
            zp->data.F32[i] = zpExp->data.F32[i];
        }
    }

    psArray *matches = pmSourceMatchSources(sourceLists, radius, true); // List of matches
    if (!matches) {
        psError(PPSTACK_ERR_DATA, false, "Unable to match sources");
        psFree(zp);
        return false;
    }

#ifdef TESTING
    dumpMatches("source_match.dat", num, matches, zp, NULL);
#endif

    if (options->matchZPs) {
        psVector *trans = pmSourceMatchRelphot(matches, zp, tol, iter1, starRej1, starSys1,
                                               iter2, starRej2, starSys2, starLimit,
                                               transIter, transRej, transThresh); // Transparencies per image
        if (!trans) {
            psError(PPSTACK_ERR_DATA, false, "Unable to measure transparencies");
            return false;
        }
        for (int i = 0; i < trans->n; i++) {
            if (!isfinite(trans->data.F32[i])) {
                inputMask->data.U8[i] |= PPSTACK_MASK_CAL;
            }
        }
        // M = m + c0 + c1 * airmass - 2.5log(t) + transparency
        // Want sources to have m corresponding to airmass = 1 and t = sumExpTime and transparency = 0
        // m_0 + c1 * airmass_0 + 2.5log(t_0) - trans_0 = m_1 + c1 * airmass_1 + 2.5log(t_1) - trans_1
        // m_0 = m_1 + zp_1 + trans_1 - c1 * airmass_0 - 2.5log(t_0)
        // We don't need to know the magnitude zero point for the filter, since it cancels out
        if (options->matchZPs) {
            options->norm = psVectorAlloc(num, PS_TYPE_F32);
            for (int i = 0; i < num; i++) {
                if (inputMask->data.U8[i] || !isfinite(trans->data.F32[i])) {
                    continue;
                }
                psArray *sources = sourceLists->data[i]; // Sources of interest
                float magCorr = zp->data.F32[i] + trans->data.F32[i] - 2.5*log10(options->sumExposure);
                if (zpExpNum == numGoodImages) {
                    // Using measured zero points, so attempt to set target zero point
                    magCorr -= zpTarget;
                }
                options->norm->data.F32[i] = -magCorr;
                psLogMsg("ppStack", PS_LOG_INFO,
                         "Applying scale correction to image %d: %f mag (%f)\n",
                         i, magCorr, trans->data.F32[i]);

                for (int j = 0; j < sources->n; j++) {
                    pmSource *source = sources->data[j]; // Source of interest
                    if (!source) {
                        continue;
                    }
                    source->psfMag -= magCorr;
                }
            }
        }

        if (zpExpNum == numGoodImages) {
            // Producing image with target zero point
            options->zp = zpTarget;
        } else {
            options->zp = NAN;
        }


#ifdef TESTING
        dumpMatches("source_mags.dat", num, matches, zp, trans);
#endif
        psFree(trans);

#ifdef TESTING
        // Double check: all transparencies should be zero
        {
            psArray *matches = pmSourceMatchSources(sourceLists, radius, true); // List of matches
            if (!matches) {
                psError(PPSTACK_ERR_DATA, false, "Unable to match sources");
                psFree(zp);
                return false;
            }
            psVector *trans = pmSourceMatchRelphot(matches, zp, tol, iter1, starRej1, starSys1,
                                                   iter2, starRej2, starSys2, starLimit,
                                                   transIter, transRej, transThresh); // Transparencies
            for (int i = 0; i < num; i++) {
                fprintf(stderr, "Transparency of image %d: %f\n", i, trans->data.F32[i]);
            }
            psFree(trans);
            psFree(matches);
        }
#endif
    }

    psFree(zp);
    psFree(zpExp);

#ifdef ASTROMETRY
    // Position offsets
    {
        psArray *offsets = pmSourceMatchRelastro(matches, num, tol, iter1, starRej1,
                                                  iter2, starRej2, starLimit); // Shifts for each image
        if (!offsets) {
            psError(PPSTACK_ERR_DATA, false, "Unable to measure offsets");
            return false;
        }
        for (int i = 0; i < num; i++) {
            if (options->inputMask->data.U8[i]) {
                continue;
            }
            psArray *sources = sourceLists->data[i]; // Sources of interest
            psVector *offset = offsets->data[i];                      // Offsets for image
            float dx = offset->data.F32[0], dy = offset->data.F32[1]; // Offsets to apply
            if (!isfinite(dx) || !isfinite(dy)) {
                continue;
            }
            psLogMsg("ppStack", PS_LOG_INFO, "Applying astrometric correction to image %d: %f,%f\n",
                     i, dx, dy);
            for (int j = 0; j < sources->n; j++) {
                pmSource *source = sources->data[j]; // Source of interest
                if (!source) {
                    continue;
                }
                if (source->modelPSF) {
                    source->modelPSF->params->data.F32[PM_PAR_XPOS] -= dx;
                    source->modelPSF->params->data.F32[PM_PAR_YPOS] -= dy;
                }
                if (source->peak) {
                    source->peak->xf -= dx;
                    source->peak->yf -= dy;
                }
            }
        }
        psFree(offsets);
    }
#endif

#if (defined TESTING && defined ASTROMETRY)
        // Double check: all offsets should be zero
        {
            psArray *matches = pmSourceMatchSources(sourceLists, radius, true); // List of matches
            if (!matches) {
                psError(PPSTACK_ERR_DATA, false, "Unable to match sources");
                psFree(zp);
                return false;
            }
            psArray *offsets = pmSourceMatchRelastro(matches, num, tol, iter1, starRej1,
                                                     iter2, starRej2, starLimit); // Shifts for each image
            for (int i = 0; i < num; i++) {
                psVector *offset = offsets->data[i]; // Offsets for image
                fprintf(stderr, "Offset of image %d: %f,%f\n", i, offset->data.F32[0], offset->data.F32[1]);
            }
            psFree(offsets);
            psFree(matches);
        }
#endif


    if (options->photometry) {
        // Save best matches for future photometry
        psArray *sourcesBest = options->sources = psArrayAllocEmpty(matches->n);

        // XXX something of a hack: require at least 2 detections or the nominated fraction of the max possible
        int minMatches = PS_MAX(2, fracMatch * num);// Minimum number of matches required
        for (int i = 0; i < matches->n; i++) {
            pmSourceMatch *match = matches->data[i]; // Match of interest
            if (match->num < minMatches) {
                continue;
            }

            // We need to grab a single instance of this source: just take the first available
            int image = match->image->data.S32[0]; // Index of image
            int index = match->index->data.S32[0]; // Index of source within image
            psArray *sources = sourceLists->data[image]; // Sources for image
            pmSource *source = sources->data[index]; // Source of interest

            psArrayAdd(sourcesBest, sourcesBest->n, source);
        }
        psLogMsg("ppStack", PS_LOG_INFO, "Selected %ld sources for photometry analysis", sourcesBest->n);
    }

    psFree(matches);

    return true;
}
