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

#include "ppStack.h"

//#define TESTING                         // Enable debugging output


bool ppStackSourcesTransparency(const psArray *sourceLists, const pmFPAview *view, const pmConfig *config)
{
    PS_ASSERT_ARRAY_NON_NULL(sourceLists, false);
    PS_ASSERT_PTR_NON_NULL(view, false);
    PS_ASSERT_PTR_NON_NULL(config, false);

#ifdef TESTING
    {
        // 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 iter = psMetadataLookupS32(NULL, recipe, "ZP.ITER"); // Maximum iterations
    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 starRej = psMetadataLookupF32(NULL, recipe, "ZP.STAR.REJ"); // Rejection threshold for stars
    float starLimit = psMetadataLookupF32(NULL, recipe, "ZP.STAR.LIMIT"); // Limit on star rejection fraction
    float starSys = psMetadataLookupF32(NULL, recipe, "ZP.STAR.SYS"); // Estimated systematic error

    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

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

    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); // Transparencies for each image

#ifdef TESTING
    {
        // Dump the corrected magnitudes
        FILE *outMatches = fopen("source_match.dat", "w"); // Output matches
        psVector *mag = psVectorAlloc(num, PS_TYPE_F32); // Magnitudes for each star
        for (int i = 0; i < matches->n; i++) {
            pmSourceMatch *match = matches->data[i]; // Match of interest
            psVectorInit(mag, 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] - trans->data.F32[index];
            }
            for (int j = 0; j < num; j++) {
                fprintf(outMatches, "%f ", mag->data.F32[j]);
            }
            fprintf(outMatches, "\n");
        }
        psFree(mag);
        fclose(outMatches);
    }
#endif

    psFree(matches);
    if (!trans) {
        psError(PS_ERR_UNKNOWN, false, "Unable to measure transparencies");
        return false;
    }

    // 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
    for (int i = 0; i < num; i++) {
        psArray *sources = sourceLists->data[i]; // Sources of interest
        float magCorr = airmassTerm - 2.5*log10(sumExpTime) - zp->data.F32[i] - trans->data.F32[i];
        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);
        psFree(trans);
    }
#endif
    return true;
}
