#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

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

#include "pmSource.h"

#include "pmSourceMatch.h"

#define SOURCE_MASK (PM_SOURCE_MODE_FAIL | PM_SOURCE_MODE_SATSTAR | PM_SOURCE_MODE_BLEND | \
                     PM_SOURCE_MODE_BADPSF | PM_SOURCE_MODE_DEFECT | PM_SOURCE_MODE_SATURATED | \
                     PM_SOURCE_MODE_CR_LIMIT | PM_SOURCE_MODE_EXT_LIMIT) // Mask to apply to input sources
#define SOURCE_FAINTEST 50.0            // Faintest magnitude to consider
#define SOURCES_MAX_LEAF 2              // Maximum number of points on a tree leaf


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Private (file-static) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Extract coordinates from a source
static inline void coordsFromSource(float *x, float *y, // Coordinates to return
                                    pmSource *source // Source of interest
    )
{
    if (!source) {
        *x = NAN;
        *y = NAN;
    } else if (source->modelPSF) {
        *x = source->modelPSF->params->data.F32[PM_PAR_XPOS];
        *y = source->modelPSF->params->data.F32[PM_PAR_YPOS];
    } else {
        *x = source->peak->xf;
        *y = source->peak->yf;
    }
    return;
}

// Parse the sources into vectors for each coordinate, and a bounding box
// Returns number of valid sources
static long sourcesParse(psRegion **bounds, // Region to update with bounding box
                         psVector **x, psVector **y, // Coordinate vectors to return
                         psVector **mag, psVector **magErr, // Magnitude and error vectors to return
                         const psArray *sources // Input sources
    )
{
    psAssert(bounds, "Must be given a region for bounding box");
    psAssert(x && y, "Must be given position vectors");
    psAssert(mag && magErr, "Must be given magnitude vectors");
    psAssert(sources, "Must be given sources");

    long numSources = sources->n;              // Number of sources
    *x = psVectorRecycle(*x, numSources, PS_TYPE_F32);
    *y = psVectorRecycle(*y, numSources, PS_TYPE_F32);
    *mag = psVectorRecycle(*mag, numSources, PS_TYPE_F32);
    *magErr = psVectorRecycle(*magErr, numSources, PS_TYPE_F32);
    float xMin = INFINITY, xMax = -INFINITY, yMin = INFINITY, yMax = -INFINITY; // Bounds of sources
    long num = 0;                       // Number of valid sources
    for (long i = 0; i < numSources; i++) {
        pmSource *source = sources->data[i]; // Source of interest
        if (!source || (source->mode & SOURCE_MASK) || !isfinite(source->psfMag) ||
            !isfinite(source->errMag) || source->psfMag > SOURCE_FAINTEST) {
            continue;
        }
        float xSrc, ySrc;               // Coordinates of source
        coordsFromSource(&xSrc, &ySrc, source);
        if (xSrc < xMin) xMin = xSrc;
        if (xSrc > xMax) xMax = xSrc;
        if (ySrc < yMin) yMin = ySrc;
        if (ySrc > yMax) yMax = ySrc;

        (*x)->data.F32[num] = xSrc;
        (*y)->data.F32[num] = ySrc;
        (*mag)->data.F32[num] = source->psfMag;
        (*magErr)->data.F32[num] = source->errMag;
        num++;
    }
    (*x)->n = num;
    (*y)->n = num;
    (*mag)->n = num;
    (*magErr)->n = num;

    if (*bounds) {
        (*bounds)->x0 = xMin;
        (*bounds)->x1 = xMax;
        (*bounds)->y0 = yMin;
        (*bounds)->y1 = yMax;
    } else {
        *bounds = psRegionAlloc(xMin, xMax, yMin, yMax);
    }

    psTrace("psModules.objects", 8, "%ld sources: bounds are [%.2f:%.2f,%.2f:%.2f]\n",
            num, xMin, xMax, yMin, yMax);

    return num;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// pmSourceMatch operations
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

static void sourceMatchFree(pmSourceMatch *match)
{
    psFree(match->mag);
    psFree(match->magErr);
    psFree(match->image);
    psFree(match->index);
}

pmSourceMatch *pmSourceMatchAlloc(int num // Maximum number of images
    )
{
    pmSourceMatch *match = psAlloc(sizeof(pmSourceMatch)); // Match data
    psMemSetDeallocator(match, (psFreeFunc)sourceMatchFree);

    match->num = 0;
    match->mag = psVectorAllocEmpty(num, PS_TYPE_F32);
    match->magErr = psVectorAllocEmpty(num, PS_TYPE_F32);
    match->image = psVectorAllocEmpty(num, PS_TYPE_U32);
    match->index = psVectorAllocEmpty(num, PS_TYPE_U32);
    match->mask = psVectorAllocEmpty(num, PS_TYPE_VECTOR_MASK);

    return match;
}

void pmSourceMatchAdd(pmSourceMatch *match, // Match data
                      float mag, float magErr, // Magnitude and error
                      int image, // Image index
                      int index // Source index
    )
{
    int num = match->num;               // Number of matches
    psAssert(num <= match->mag->nalloc, "Too many matches.");
    psAssert(num <= match->magErr->nalloc, "Too many matches.");
    psAssert(num <= match->image->nalloc, "Too many matches.");
    psAssert(num <= match->index->nalloc, "Too many matches.");

    match->mag->data.F32[num] = mag;
    match->magErr->data.F32[num] = magErr;
    match->image->data.S32[num] = image;
    match->index->data.S32[num] = index;
    match->mask->data.PS_TYPE_VECTOR_MASK_DATA[num] = 0;
    match->num++;

    match->mag->n = match->magErr->n = match->image->n = match->index->n = match->mask->n = match->num;
}


psArray *pmSourceMatchSources(const psArray *sourceArrays, float radius)
{
    PS_ASSERT_ARRAY_NON_NULL(sourceArrays, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN(radius, 0.0, NULL);

    int numImages = sourceArrays->n;    // Number of images

    // First, merge the source lists, pulling out matches

    psRegion *boundsMaster = NULL;       // Bounds of source list
    psVector *xMaster = NULL, *yMaster = NULL; // Coordinates of sources
    long numMaster = 0;                 // Number in master list
    psArray *matches = NULL;            // Source matches (potential and actual matches)

    for (int i = 0; i < numImages; i++) {
        psArray *sources = sourceArrays->data[i]; // Sources in image
        if (!sources) {
            continue;
        }
        psRegion *boundsImage = NULL;   // Bounds of source list
        psVector *xImage = NULL, *yImage = NULL; // Coordinates of sources
        psVector *magImage = NULL, *magErrImage = NULL; // Magnitude and mag

        int numSources = sourcesParse(&boundsImage, &xImage, &yImage, &magImage, &magErrImage,
                                      sources); // Number of sources

        if (!boundsMaster) {
            // First run through --- can just copy
            boundsMaster = psMemIncrRefCounter(boundsImage);
            xMaster = psMemIncrRefCounter(xImage);
            yMaster = psMemIncrRefCounter(yImage);
            matches = psArrayAlloc(numSources);
            for (int j = 0; j < numSources; j++) {
                pmSourceMatch *match = pmSourceMatchAlloc(numImages); // Match data
                pmSourceMatchAdd(match, magImage->data.F32[j], magErrImage->data.F32[j], i, j);
                matches->data[j] = match;
            }
            numMaster += numSources;
            goto match_image_done;
        }

        if (boundsImage->x0 > boundsMaster->x1 || boundsImage->x1 < boundsMaster->x0 ||
            boundsImage->y0 > boundsMaster->y1 || boundsImage->y1 < boundsMaster->y0) {
            psTrace("psModules.objects", 7, "Bounds don't overlap\n");
            // Add everything in to the master list
            long size = numMaster + numSources; // New size
            xMaster = psVectorRealloc(xMaster, size);
            yMaster = psVectorRealloc(yMaster, size);
            matches = psArrayRealloc(matches, size);

            memcpy(&xMaster->data.F32[numMaster], xImage->data.F32,
                   numSources * PSELEMTYPE_SIZEOF(PS_TYPE_F32));
            memcpy(&yMaster->data.F32[numMaster], yImage->data.F32,
                   numSources * PSELEMTYPE_SIZEOF(PS_TYPE_F32));
            for (int j = 0, k = numMaster; j < numSources; j++, k++) {
                pmSourceMatch *match = pmSourceMatchAlloc(numImages); // Match data
                pmSourceMatchAdd(match, magImage->data.F32[j], magErrImage->data.F32[j], i, j);
                matches->data[k] = match;
            }

            xMaster->n = size;
            yMaster->n = size;
            matches->n = size;

            goto match_image_done;
        }

        // Match with the master list

        psTree *tree = psTreePlant(2, SOURCES_MAX_LEAF, xMaster, yMaster); // kd Tree with sources
        long numMatch = 0;              // Number of matches

        long size = numMaster + numSources; // New size
        xMaster = psVectorRealloc(xMaster, size);
        yMaster = psVectorRealloc(yMaster, size);
        matches = psArrayRealloc(matches, size);

        psVector *coords = psVectorAlloc(2, PS_TYPE_F32); // Coordinates for tree lookup
        for (int j = 0; j < numSources; j++) {
            coords->data.F32[0] = xImage->data.F32[j];
            coords->data.F32[1] = yImage->data.F32[j];
            long index = psTreeNearestWithin(tree, coords, radius); // Match index
            if (index >= 0) {
                // Record the match
                pmSourceMatch *match = matches->data[index]; // Match data
                pmSourceMatchAdd(match, magImage->data.F32[j], magErrImage->data.F32[j], i, j);
                numMatch++;
            } else {
                // Add to the master list
                pmSourceMatch *match = pmSourceMatchAlloc(numImages); // Match data
                pmSourceMatchAdd(match, magImage->data.F32[j], magErrImage->data.F32[j], i, j);
                xMaster->data.F32[numMaster] = xImage->data.F32[j];
                yMaster->data.F32[numMaster] = yImage->data.F32[j];
                matches->data[numMaster] = match;
                numMaster++;
                xMaster->n = yMaster->n = matches->n = numMaster;
            }

        }
        psFree(coords);
        psFree(tree);


    match_image_done:
        psFree(boundsImage);
        psFree(xImage);
        psFree(yImage);
        psFree(magImage);
        psFree(magErrImage);
    }

    // Now cull the matches that contain only a single star
    int numGood = 0;                    // Number of good matches
    for (int i = 0; i < matches->n; i++) {
        pmSourceMatch *match = matches->data[i]; // Match of interest
        if (match->num > 1) {
            if (i != numGood) {
                psFree(matches->data[numGood]);
                matches->data[numGood] = psMemIncrRefCounter(match);
                //                psFree(match);
            }
            numGood++;
        }
    }
    matches->n = numGood;
    for (int i = numGood; i < numMaster; i++) {
        psFree(matches->data[i]);
        matches->data[i] = NULL;
    }

    return matches;
}


// Iterate on the star magnitudes and image transparencies
// Returns the solution chi^2
float sourceMatchRelphotIterate(psVector *trans, // Transparencies
                                psVector *stars, // Star magnitudes
                                const psArray *matches, // Array of matches
                                const psVector *zp, // Zero points for each image (including airmass term)
                                const psVector *photo, // Photometric image?
                                float sysErr2 // Systematic error, squared
                                )
{
    psAssert(zp && zp->type.type == PS_TYPE_F32, "Need zero points");
    psAssert(matches, "Need list of matches");

    int numImages = zp->n;              // Number of images
    int numStars = matches->n;          // Number of stars

    psAssert(trans && trans->type.type == PS_TYPE_F32, "Need transparencies");
    psAssert(trans->n == numImages, "Not enough transparencies: %ld\n", trans->n);
    psAssert(stars && stars->type.type == PS_TYPE_F32, "Need star magnitudes");
    psAssert(stars->n == numStars, "Not enough stars: %ld\n", stars->n);
    psAssert(zp->n == numImages, "Not enough ZPs: %ld", zp->n);
    psAssert(!photo || photo->type.type == PS_TYPE_U8, "Photometric determination is wrong type");
    psAssert(!photo || photo->n == numImages, "Not enough photometric determinations: %ld", photo->n);

    // Solve the star magnitudes
    psVectorInit(stars, 0.0);
    int numGoodStars = 0;               // Number of stars with good measurements
    for (int i = 0; i < numStars; i++) {
        pmSourceMatch *match = matches->data[i]; // Matched stars
        int numMeasurements = 0;        // Number of unmasked measurements for star
        double star = 0.0, starErr = 0.0; // Accumulators for star
        for (int j = 0; j < match->num; j++) {
            if (match->mask->data.PS_TYPE_VECTOR_MASK_DATA[j]) {
                continue;
            }
            numMeasurements++;
            int index = match->image->data.U32[j]; // Image index
            float mag = match->mag->data.F32[j]; // Measured magnitude
            double magErr2 = PS_SQR(match->magErr->data.F32[j]) + sysErr2; // Error in measured magnitude
            double invErr2 = 1.0 / magErr2; // Inverse square error
            float cal = zp->data.F32[index]; // Calibration to apply to image
            if (!photo || !photo->data.U8[index]) {
                cal -= trans->data.F32[index];
            }
            star += (mag + cal) * invErr2;
            starErr += invErr2;
        }
        if (numMeasurements > 1) {
            // It's only a good star (contributing to the chi^2) if there's more than 1 measurement
            numGoodStars++;
        }
        stars->data.F32[i] = star / starErr;
    }

    // Solve for the transparencies
    // We solve even for the "photometric" images since they may jump out of that status upon iteration
    psVector *accum = psVectorAlloc(numImages, PS_TYPE_F64); // Transparency accumulator
    psVector *accumErr = psVectorAlloc(numImages, PS_TYPE_F64); // Transparency accumulator
    psVectorInit(accum, 0.0);
    psVectorInit(accumErr, 0.0);
    for (int i = 0; i < numStars; i++) {
        pmSourceMatch *match = matches->data[i]; // Matched stars
        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
            float mag = match->mag->data.F32[j]; // Measured magnitude
            double magErr2 = PS_SQR(match->magErr->data.F32[j]) + sysErr2; // Error in measured magnitude
            double invErr2 = 1.0 / magErr2; // Inverse square error
            float cal = zp->data.F32[index]; // Calibration to apply to image
            accum->data.F64[index] += (mag + cal - stars->data.F32[i]) * invErr2;
            accumErr->data.F64[index] += invErr2;
        }
    }
    for (int i = 0; i < numImages; i++) {
        trans->data.F32[i] = accum->data.F64[i] / accumErr->data.F64[i];

        psTrace("psModules.objects", 3, "Transparency for image %d: %f\n", i, trans->data.F32[i]);
    }
    psFree(accum);
    psFree(accumErr);

    // Once more through to evaluate chi^2
    float chi2 = 0.0;                   // chi^2 for iteration
    int dof = 0;                        // Degrees of freedom
    for (int i = 0; i < numStars; i++) {
        pmSourceMatch *match = matches->data[i]; // Matched stars
        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
            float mag = match->mag->data.F32[j]; // Measured magnitude
            float magErr2 = PS_SQR(match->magErr->data.F32[j]) + sysErr2; // Error in measured magnitude
            float cal = zp->data.F32[index]; // Calibration to apply to image
            if (!photo || !photo->data.U8[index]) {
                cal -= trans->data.F32[index];
            }
            float dev2 = mag + cal - stars->data.F32[index]; // Deviation from fit
            if (isfinite(dev2)) {
                chi2 += PS_SQR(dev2) / magErr2;
                dof++;
            }
        }
    }
    dof -= numGoodStars + numImages;
    chi2 /= dof;

    return chi2;
}

// Determine which images are photometric, based on estimated transparencies
// Returns number of photometric images, or -1 on error
static int sourceMatchRelphotPhotometric(psVector *photo, // Photometric determination
                                         const psVector *trans, // Estimated transparencies
                                         int transIter, // Iterations for transparency
                                         float transClip, // Clipping level for transparency
                                         float photoLevel // Level below which we declare photometric
                                         )
{
    psAssert(photo && photo->type.type == PS_TYPE_U8, "Need photometric determination");
    psAssert(trans && trans->type.type == PS_TYPE_F32, "Need transparencies");

    int numImages = photo->n;              // Number of images

    psAssert(trans->n == numImages, "Not enough transparencies: %ld", trans->n);
    psAssert(transIter >= 0, "Iterations for transparency must be non-negative: %d", transIter);
    psAssert(transClip > 0, "Clipping level for transparency must be positive: %f", transClip);
    psAssert(photoLevel > 0, "Photometric level must be positive: %f", photoLevel);

    psStats *stats = psStatsAlloc(PS_STAT_CLIPPED_MEAN | PS_STAT_CLIPPED_STDEV); // Statistics
    stats->clipIter = transIter;
    stats->clipSigma = transClip;

    if (!psVectorStats(stats, trans, NULL, NULL, 0)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to perform statistics on transparencies.");
        psFree(stats);
        return -1;
    }

    float thresh = stats->clippedMean + photoLevel * stats->clippedStdev; // Threshold for clouds
    psFree(stats);

    int numPhoto = 0;                   // Number of photometric images
    for (int i = 0; i < numImages; i++) {
        if (trans->data.F32[i] < thresh) {
            photo->data.U8[i] = 0xFF;
            numPhoto++;
        } else {
            photo->data.U8[i] = 0;
        }
    }

    return numPhoto;
}


// Reject star measurements
// Returns the fraction of measurements that were rejected
static float sourceMatchRelphotReject(psVector *trans, // Transparencies
                                      psVector *stars, // Star magnitudes
                                      const psArray *matches, // Array of matches
                                      const psVector *zp, // Zero points for each image
                                      const psVector *photo, // Photometric image?
                                      float starClip, // Clipping for stars
                                      float sysErr2 // Systematic error squared
                                )
{
    psAssert(zp && zp->type.type == PS_TYPE_F32, "Need zero points");
    psAssert(matches, "Need list of matches");

    int numImages = zp->n;              // Number of images
    int numStars = matches->n;          // Number of stars

    psAssert(trans && trans->type.type == PS_TYPE_F32, "Need transparencies");
    psAssert(trans->n == numImages, "Not enough transparencies: %ld\n", trans->n);
    psAssert(stars && stars->type.type == PS_TYPE_F32, "Need star magnitudes");
    psAssert(stars->n == numStars, "Not enough stars: %ld\n", stars->n);
    psAssert(zp->n == numImages, "Not enough ZPs: %ld", zp->n);
    psAssert(!photo || photo->type.type == PS_TYPE_U8, "Photometric determination is wrong type");
    psAssert(!photo || photo->n == numImages, "Not enough photometric determinations: %ld", photo->n);

    starClip = PS_SQR(starClip);

    int numRejected = 0;                // Number rejected
    int numMeasurements = 0;            // Number of measurements
    for (int i = 0; i < numStars; i++) {
        pmSourceMatch *match = matches->data[i]; // Matched stars
        for (int j = 0; j < match->num; j++) {
            if (match->mask->data.PS_TYPE_VECTOR_MASK_DATA[j]) {
                continue;
            }
            numMeasurements++;
            int index = match->image->data.U32[j]; // Image index
            float mag = match->mag->data.F32[j]; // Measured magnitude
            float magErr = match->magErr->data.F32[j]; // Error in measured magnitude
            float cal = zp->data.F32[index]; // Calibration to apply to image
            if (!photo || !photo->data.U8[index]) {
                cal -= trans->data.F32[index];
            }
            float dev = mag + cal - stars->data.F32[i]; // Deviation
            if (PS_SQR(dev) > starClip * (PS_SQR(magErr) + sysErr2)) {
                numRejected++;
                match->mask->data.PS_TYPE_VECTOR_MASK_DATA[j] = 0xFF;
            }
        }
    }

    return (float)numRejected / (float)numMeasurements;
}


psVector *pmSourceMatchRelphot(const psArray *matches, // Array of matches
                               const psVector *zp, // Zero points for each image (including airmass term)
                               int maxIter, // Maximum number of iterations
                               float tol, // Relative tolerance for convergence
                               float rejLimit, // Limit on rejection between iterations
                               int transIter, // Clipping iterations for transparency
                               float transClip, // Clipping level for transparency
                               float photoLevel, // Level at which we declare image is photometric
                               float starClip, // Clipping for stars
                               float sysErr // Systematic error in measurements
                               )
{
    PS_ASSERT_ARRAY_NON_NULL(matches, NULL);
    PS_ASSERT_VECTOR_NON_NULL(zp, NULL);
    PS_ASSERT_VECTOR_TYPE(zp, PS_TYPE_F32, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN(transClip, 0.0, NULL);

    sysErr *= sysErr;

    int numImages = zp->n;              // Number of images
    int numStars = matches->n;          // Number of stars
    psVector *trans = psVectorAlloc(numImages, PS_TYPE_F32); // Transparencies for each image, magnitudes
    psVectorInit(trans, 0.0);
    psVector *photo = psVectorAlloc(numImages, PS_TYPE_U8); // Photometric determination for each image
    psVectorInit(photo, 0);
    psVector *stars = psVectorAlloc(numStars, PS_TYPE_F32); // Magnitudes for each star

    float chi2 = sourceMatchRelphotIterate(trans, stars, matches, zp, photo, sysErr); // chi^2 for solution
    psTrace("psModules.objects", 1, "Initial: chi^2 = %f\n", chi2);
    float lastChi2 = INFINITY;          // chi^2 on last iteration
    float fracRej = INFINITY;        // Fraction of measurements rejected
    for (int i = 0; i < maxIter && (fabsf(lastChi2 - chi2) > tol * chi2 || fracRej > rejLimit); i++) {
        lastChi2 = chi2;

        // Identify photometric nights
        int numPhoto = sourceMatchRelphotPhotometric(photo, trans, transIter, transClip,
                                                     photoLevel); // Number of photometric images
        if (numPhoto < 0) {
            psError(PS_ERR_UNKNOWN, false, "Unable to perform photometric determination");
            psFree(trans);
            psFree(photo);
            psFree(stars);
            return NULL;
        }
        psTrace("psModules.objects", 3, "Determined %d/%d are photometric", numPhoto, numImages);

        fracRej = sourceMatchRelphotReject(trans, stars, matches, zp, photo, starClip, sysErr);
        psTrace("psModules.objects", 3, "%f%% of measurements rejected", fracRej * 100);

        chi2 = sourceMatchRelphotIterate(trans, stars, matches, zp, photo, sysErr); // chi^2 for solution
        psTrace("psModules.objects", 1, "iter = %d: chi^2 = %f rejected = %f\n", i, chi2, fracRej);
    }

    psFree(photo);
    psFree(stars);

    if (fabsf(lastChi2 - chi2) > tol * chi2 || fracRej > rejLimit) {
        psWarning("Unable to converge to relphot solution (%f,%f)", (lastChi2 - chi2) / chi2, fracRej);
    }

    return trans;
}
