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

#include "ppStack.h"

//#define TESTING                         // Enable debugging output

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

    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;
        }
    }
#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(PS_ERR_UNKNOWN, false, "Unable to find ZP.AIRMASS in recipe.");
        return false;
    }

    int num = psMetadataLookupS32(NULL, config->arguments, "INPUTS.NUM"); // Number of inputs
    psAssert(num == sourceLists->n, "Wrong number of source lists: %ld\n", sourceLists->n);

    psVector *zp = psVectorAlloc(num, PS_TYPE_F32); // Zero points for each image
    const char *filter = NULL;          // Filter name
    float airmassTerm = NAN;            // Airmass term
    float sumExpTime = 0.0;             // Sum of the exposure time
    for (int i = 0; i < num; i++) {
        pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPSTACK.INPUT", i); // File of interest
        pmCell *cell = pmFPAviewThisCell(view, file->fpa); // Cell 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 = psMetadataLookupF32(NULL, cell->concepts, "CELL.EXPOSURE"); // Exposure time
        float airmass = psMetadataLookupF32(NULL, file->fpa->concepts, "FPA.AIRMASS"); // Airmass
        const char *expFilter = psMetadataLookupStr(NULL, file->fpa->concepts, "FPA.FILTER"); // Filter name
        if (!isfinite(exptime) || exptime == 0 || !isfinite(airmass) || airmass == 0 ||
            !expFilter || strlen(expFilter) == 0) {
            psError(PS_ERR_UNEXPECTED_NULL, false,
                    "Unable to find exposure time (%f), airmass (%f) or filter (%s)",
                    exptime, airmass, expFilter);
            psFree(zp);
            return false;
        }

        if (!filter) {
            filter = expFilter;
            airmassTerm = psMetadataLookupF32(NULL, airmassZP, filter);
            if (!isfinite(airmassTerm)) {
                psError(PS_ERR_BAD_PARAMETER_VALUE, false,
                        "Unable to find airmass term (ZP.AIRMASS) for filter %s", filter);
                psFree(zp);
                return false;
            }
        } else if (strcmp(filter, expFilter) != 0) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, false, "Filters don't match: %s vs %s", filter, expFilter);
            psFree(zp);
            return false;
        }

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

    options->sumExposure = sumExpTime;

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

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

    psVector *trans = pmSourceMatchRelphot(matches, zp, tol, iter1, starRej1, starSys1,
                                           iter2, starRej2, starSys2, starLimit,
                                           transIter, transRej, transThresh); // Transparencies for each image
    if (!trans) {
        psError(PS_ERR_UNKNOWN, false, "Unable to measure transparencies");
        return false;
    }

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

    for (int i = 0; i < trans->n; i++) {
        if (!isfinite(trans->data.F32[i])) {
            inputMask->data.U8[i] |= PPSTACK_MASK_CAL;
        }
    }

    // Save best matches SOMEWHERE for future photometry
    // XXX this is a really poor output location; clean up the pmFPAfiles used in ppStack
    pmCell *sourcesCell = pmFPAfileThisCell(config->files, view, "PPSTACK.OUTPUT");
    psArray *sourcesBest = 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);
    }
    psMetadataAdd(sourcesCell->analysis, PS_LIST_TAIL, "PSPHOT.SOURCES", PS_DATA_ARRAY | PS_META_REPLACE,
                  "psphot sources", sourcesBest);
    psLogMsg("ppStack", PS_LOG_INFO, "Selected %ld sources for photometry analysis", sourcesBest->n);
    psFree(sourcesBest);

    psFree(matches);

    // 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
    options->norm = psVectorAlloc(num, PS_TYPE_F32);
    for (int i = 0; i < num; i++) {
        if (!isfinite(trans->data.F32[i])) {
            continue;
        }
        psArray *sources = sourceLists->data[i]; // Sources of interest
        float magCorr = airmassTerm - 2.5*log10(sumExpTime) - zp->data.F32[i] - trans->data.F32[i];
        options->norm->data.F32[i] = magCorr;
        psLogMsg("ppStack", PS_LOG_INFO, "Applying magnitude correction to image %d: %f\n", i, magCorr);

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

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

    return true;
}
