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

#include "ppStack.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

//#define TESTING                         // Enable test output

// Stack of sources
typedef struct {
    psRegion *bounds;                   // Bounding box for sources
    psArray *sources;                   // Sources within region
    psVector *indices;                  // Indices from tree to sources
    psArray *duplicates;                // Duplicate sources (from overlaps)
    psTree *tree;                       // kd tree with sources
} ppStackSourceList;


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

// Deallocator
static void stackSourceListFree(ppStackSourceList *list)
{
    psFree(list->bounds);
    psFree(list->sources);
    psFree(list->indices);
    psFree(list->duplicates);
    psFree(list->tree);
}

// Parse the sources into vectors for each coordinate, and a bounding box
// Returns number of valid sources
static long stackSourcesParse(psRegion **bounds, // Region to update with bounding box
                              psVector **x, psVector **y, // Coordinate vectors to return
                              psVector **indices, // Source indices to return
                              const psArray *sources // Input sources
    )
{
    psAssert(bounds, "Must be given a region for bounding box");
    psAssert(x && y, "Must be given vectors");
    psAssert(indices, "Must be given indices");
    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);
    *indices = psVectorRecycle(*indices, numSources, PS_TYPE_U32);
    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) ||
            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;
        (*indices)->data.U32[num] = i;
        num++;
    }
    (*x)->n = num;
    (*y)->n = num;
    (*indices)->n = num;

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

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

    return num;
}


// Allocator for ppStackSourceList
ppStackSourceList *ppStackSourceListAlloc(psArray *sources // Sources to add
    )
{
    psAssert(sources, "Must be given sources");

    ppStackSourceList *list = psAlloc(sizeof(ppStackSourceList));
    psMemSetDeallocator(list, (psFreeFunc)stackSourceListFree);

    // Copy the sources onto a new array, so there's no confusion between the inputs and outputs
    long num = sources->n;              // Number of sources
    list->sources = psArrayAlloc(num);
    for (long i = 0; i < num; i++) {
        list->sources->data[i] = psMemIncrRefCounter(sources->data[i]);
    }

    list->bounds = NULL;
    list->indices = NULL;
    list->duplicates = psArrayAllocEmpty(0);

    psVector *x = NULL, *y = NULL;      // Source coordinates
    stackSourcesParse(&list->bounds, &x, &y, &list->indices, sources);
    list->tree = psTreePlant(2, SOURCES_MAX_LEAF, x, y); // kd Tree with sources
    psFree(x);
    psFree(y);

    return list;
}

// Extend a list of sources
bool ppStackSourceListExtend(ppStackSourceList *list, // List to extend
                             const psArray *newSources, // Sources to add
                             const psArray *newDups // Duplicate sources to add
    )
{
    psAssert(list, "Must be given a list");
    psAssert(newSources, "Must be given sources");

    {
        psArray *oldSources = list->sources;// Old list of sources
        long numOld = oldSources->n;        // Number of old sources
        long numNew = newSources->n;        // Number of new sources

        list->sources = oldSources = psArrayRealloc(oldSources, numOld + numNew);
        for (long i = 0; i < numNew; i++) {
            psArrayAdd(oldSources, 1, newSources->data[i]);
        }
        psFree(list->tree);

        psVector *x = NULL, *y = NULL;      // Source coordinates
        stackSourcesParse(&list->bounds, &x, &y, &list->indices, list->sources);
        list->tree = psTreePlant(2, SOURCES_MAX_LEAF, x, y); // kd Tree with sources
        psFree(x);
        psFree(y);
    }

    if (newDups) {
        psArray *oldDups = list->duplicates; // Old list of duplicates
        long numOld = oldDups ? oldDups->n : 0; // Number of old duplicates
        long numNew = newDups->n;       // Number of new duplicates

        list->duplicates = oldDups = psArrayRealloc(oldDups, numOld + numNew);
        for (long i = 0; i < numNew; i++) {
            psArrayAdd(oldDups, 1, newDups->data[i]);
        }
    }

    return list;
}

// Compare an array of sources with a list
// Return number inside the list
static long stackSourceListCompare(ppStackSourceList **merge, // Merge target, modified
                                   psArray *lists, // Array of source lists
                                   int listIndex, // Index of list to compare with sources
                                   psArray **sourcesPtr, // Sources to compare with list, modified
                                   float radius, // Matching radius
                                   int minOverlap, // Minimum number of sources
                                   int iter, // Clipping iterations
                                   float rej // Clipping rejection level
    )
{
    psAssert(merge, "Expected a source list pointer");
    psAssert(radius > 0, "Silly radius: %f", radius);
    psAssert(lists, "Expected an array of source lists");
    psAssert(sourcesPtr && *sourcesPtr, "Expected a source array");

    ppStackSourceList *list = lists->data[listIndex]; // List of interest
    if (!list) {
        return 0;
    }

    psArray *sources = *sourcesPtr;     // Sources to compare with list
    psVector *x = NULL, *y = NULL;      // Coordinates of sources
    psVector *indices = NULL;           // Indices for sources
    psRegion *bounds = NULL;            // Bounding box for sources
    long numSources = stackSourcesParse(&bounds, &x, &y, &indices, sources); // Number of (good) sources

    if (bounds->x0 > list->bounds->x1 || bounds->x1 < list->bounds->x0 ||
        bounds->x0 > list->bounds->x1 || bounds->x1 < list->bounds->x0) {
        // Bounds don't overlap, so the sources don't either
        psTrace("ppStack.sources", 7, "Bounds don't overlap\n");
        psFree(x);
        psFree(y);
        psFree(indices);
        psFree(bounds);
        return 0;
    }
    psFree(bounds);

    long numIn = 0;                     // Number of sources in list
    long numOut = 0;                    // Number of sources outside list

#ifdef TESTING
    static int fileNum = 0;             // Number of file
    psString filename = NULL;   // Name of file
    psStringAppend(&filename, "source_match_%d.txt", fileNum++);
    FILE *file = fopen(filename, "w"); // File to write
    psFree(filename);
    fprintf(file, "# x y mag1 mag2\n");
#endif

    psVector *magDiff = psVectorAlloc(numSources, PS_TYPE_F32); // Magnitude differences
    psVector *coords = psVectorAlloc(2, PS_TYPE_F32); // Coordinates of source
    psArray *outside = psArrayAllocEmpty(numSources); // Sources that don't correlate
    psArray *inside = psArrayAllocEmpty(numSources); // Sources that do correlate
    for (long i = 0; i < numSources; i++) {
        coords->data.F32[0] = x->data.F32[i];
        coords->data.F32[1] = y->data.F32[i];
        long match = psTreeNearestWithin(list->tree, coords, radius); // Index of match
        if (match < 0) {
            numOut++;
            psArrayAdd(outside, 1, sources->data[indices->data.U32[i]]);
            continue;
        }

        pmSource *listSource = list->sources->data[list->indices->data.U32[match]]; // Source in list
        pmSource *source = sources->data[indices->data.U32[i]]; // Source of interest
        float listMag = listSource->psfMag; // Magnitude of source in list
        float sourceMag = source->psfMag; // Magnitude of source
        if (!(listSource->mode & SOURCE_MASK) && !(source->mode & SOURCE_MASK) &&
            isfinite(listMag) && isfinite(sourceMag)) {
#ifdef TESTING
            fprintf(file, "%f %f %f %f %.4x %.4x\n", x->data.F32[i], y->data.F32[i],
                    listSource->psfMag, source->psfMag, listSource->mode, source->mode);
#endif
            magDiff->data.F32[numIn] = listSource->psfMag - source->psfMag;
            psArrayAdd(inside, 1, source);
            numIn++;
        }
    }
    magDiff->n = numIn;

#ifdef TESTING
    fclose(file);
#endif

    psFree(coords);
    psFree(x);
    psFree(y);
    psFree(indices);

    psTrace("ppStack.sources", 7, "%ld sources (vs %d) overlap list %d\n", numIn, minOverlap, listIndex);

    if (numIn > minOverlap) {
        // There's sufficient overlap --- calculate the magnitude difference
        psStats *stats = psStatsAlloc(iter == 0 ? (PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV) :
                                      (PS_STAT_CLIPPED_MEAN | PS_STAT_CLIPPED_STDEV));
        stats->clipSigma = rej;
        stats->clipIter = iter;
        if (!psVectorStats(stats, magDiff, NULL, NULL, 0)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to determine mean magnitude difference");
            psFree(magDiff);
            psFree(stats);
            return -1;
        }
        float meanDiff = iter == 0 ? stats->sampleMean : stats->clippedMean; // Mean magnitude difference
        float stdev = iter == 0 ? stats->sampleStdev : stats->clippedStdev; // Standard deviation
        long numStats = iter == 0 ? numIn : stats->clippedNvalues; // Number used for statistics
        psLogMsg("ppStack.sources", PS_LOG_INFO, "Magnitude difference from %ld overlaps is %f +/- %f\n",
                 numStats, meanDiff, stdev);
        psFree(stats);

        // Merge our sources into the list

        if (*merge) {
            // These sources act as a bridge from one list to another
            // Correct magnitudes to be on the system defined by the first source list
            // Correct the list to 'sources' (because 'sources' has been corrected to the merge target)
            long num = list->sources->n; // Number of sources
            for (long j = 0; j < num; j++) {
                pmSource *source = list->sources->data[j];
                source->psfMag -= meanDiff;
            }
            for (long j = 0; j < list->duplicates->n; j++) {
                pmSource *source = list->duplicates->data[j];
                source->psfMag -= meanDiff;
            }

            // Suck sources from this list into the first one we found.
            psTrace("ppStack.sources", 4, "Bridge: Merging lists.");

            ppStackSourceListExtend(*merge, list->sources, NULL);
            psFree(list);
            lists->data[listIndex] = NULL;
        } else {
            // Correct 'sources' magnitudes to match the list magnitudes
            for (long j = 0; j < numOut; j++) {
                pmSource *source = outside->data[j]; // New source
                source->psfMag += meanDiff;
            }
            for (long j = 0; j < numIn; j++) {
                pmSource *source = inside->data[j]; // New source
                source->psfMag += meanDiff;
            }
            // Put the sources that didn't correlate into the list
            psTrace("ppStack.sources", 4, "Overlap: Extending list");
            ppStackSourceListExtend(list, outside, inside);
            *merge = list;
        }

        // We only need to consider sources outside the list, since the lists are (mostly) disjoint
        psFree(*sourcesPtr);
        *sourcesPtr = outside;
    } else {
        psFree(outside);
    }
    psFree(inside);
    psFree(magDiff);

    return numIn;
}


psArray *ppStackSourceListAdd(psArray *lists, psArray *sources, const pmConfig *config)
{
    PS_ASSERT_PTR_NON_NULL(sources, NULL);
    PS_ASSERT_PTR_NON_NULL(config, NULL);

    if (!lists) {
        lists = psArrayAllocEmpty(1);
    }
    if (lists->n == 0) {
        ppStackSourceList *list = ppStackSourceListAlloc(sources);
        psArrayAdd(lists, 1, list);
        psFree(list);                   // Drop reference
        return lists;
    }

    // Read recipe values
    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, "SOURCE.RADIUS"); // Exclusion radius
    int minOverlap = psMetadataLookupS32(NULL, recipe, "SOURCE.MIN"); // Minimum number
    int iter = psMetadataLookupS32(NULL, recipe, "SOURCE.ITER"); // Clipping iterations
    float rej = psMetadataLookupF32(NULL, recipe, "SOURCE.REJ"); // Clipping rejection

    // Look for overlaps with existing lists
    int numLists = lists->n;            // Number of source lists on the stack
    ppStackSourceList *merge = NULL;    // Merged source list

    psMemIncrRefCounter(sources);       // So we can play with the pointer
    for (int i = 0; i < numLists; i++) {
        long numIn = stackSourceListCompare(&merge, lists, i, &sources, radius, minOverlap,
                                            iter, rej); // Number of overlapping sources
        if (numIn == -1) {
            psError(PS_ERR_UNKNOWN, false, "Unable to compare sources.");
            psFree(sources);
            psFree(lists);
            return NULL;
        }
    }
    psFree(sources);

    if (!merge) {
        // The source list is disjoint, so throw them into a new list
        psTrace("ppStack.sources", 4, "Disjoint source: Creating new list");
        ppStackSourceList *list = ppStackSourceListAlloc(sources);
        psArrayAdd(lists, 1, list);
        psFree(list);                   // Drop reference
    }

    return lists;
}



psArray *ppStackSourceListCombine(psArray *lists, const pmConfig *config)
{
    PS_ASSERT_ARRAY_NON_NULL(lists, NULL);
    PS_ASSERT_PTR_NON_NULL(config, NULL);

    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, "SOURCE.RADIUS"); // Exclusion radius

    int numLists = 0;                   // Number of lists
    ppStackSourceList *firstList = NULL; // Last list
    for (int i = 0; i < lists->n; i++) {
        ppStackSourceList *list = lists->data[i]; // List of interest
        if (list) {
            numLists++;
            if (!firstList) {
                firstList = list;
            }
        }
    }

    if (!firstList || numLists == 0) {
        psError(PS_ERR_UNKNOWN, true, "No lists found.");
        return NULL;
    }

#if 0
    if (numLists == 1) {
        // Trivial case
        return psMemIncrRefCounter(firstList->sources);
    }
#endif

    for (int i = lists->n - 1; lists->data[i] != firstList; i--) {
        // There may be some small overlap (we know it's less than SOURCE.MIN), so attempt to correct
        // magnitudes of these first
        ppStackSourceList *lastList = lists->data[i]; // Last source list in the array
        if (!lastList) {
            continue;
        }
        psArray *sources = psMemIncrRefCounter(lastList->sources); // Sources of interest
        ppStackSourceList *merge = NULL; // Merge target
        for (int j = 0; j < i; j++) {
            ppStackSourceList *list = lists->data[i]; // List to compare
            if (!list) {
                continue;
            }
            // Note: setting iterations = 0, so no clipping (figure not enough values to do clipping)
            long numIn = stackSourceListCompare(&merge, lists, j, &sources, radius, 0,
                                                1, NAN); // Number of overlapping sources
            if (numIn < 0) {
                psError(PS_ERR_UNKNOWN, false, "Unable to compare sources.");
                psFree(sources);
                return NULL;
            }
            if (numIn == 0) {
                // It doesn't match anything at all.  Merge it in to the first list.
                psTrace("ppStack.sources", 4, "Disjoint list: Merging into first");
                ppStackSourceListExtend(firstList, list->sources, NULL);
            }
        }
        psFree(sources);
    }

    return psMemIncrRefCounter(firstList->sources);
}



void ppStackSourcesPrint(const psArray *sources)
{
    static int num = 0;                 // Number of source array
    psString filename = NULL;           // Name of file
    psStringAppend(&filename, "sources_%d.reg", num++);

    const char *goodColour = "green";   // Colour for good sources
    const char *badColour = "red";      // Colour for bad sources
    float radius = 5;                   // Radius for circle

    FILE *file = fopen(filename, "w");  // File to which to write
    psFree(filename);
    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i]; // Source of interest
        if (!source) {
            continue;
        }

        float x, y;                     // Source coordinates
        coordsFromSource(&x, &y, source);

        bool bad = (source->mode & SOURCE_MASK) || !isfinite(source->psfMag) ||
            source->psfMag > SOURCE_FAINTEST; // Is this a bad source?

        fprintf(file, "image; circle(%f,%f,%f) # color = %s\n", x + 1.0, y + 1.0, radius,
                bad ? badColour : goodColour);
    }
    fclose(file);

    return;
}
