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

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

#include "ppStack.h"

#define THREAD_WAIT 10000               // Microseconds to wait if thread is not available


static void stackThreadFree(ppStackThread *thread)
{
    psFree(thread->readouts);
    return;
}

ppStackThread *ppStackThreadAlloc(psArray *readouts)
{
    ppStackThread *thread = psAlloc(sizeof(ppStackThread));
    psMemSetDeallocator(thread, (psFreeFunc)stackThreadFree);

    thread->readouts = psMemIncrRefCounter(readouts);
    thread->read = false;
    thread->busy = false;
    thread->firstScan = 0;
    thread->lastScan = 0;
    return thread;
}

static void stackThreadDataFree(ppStackThreadData *stack)
{
    psFree(stack->threads);
    for (int i = 0; i < stack->imageFits->n; i++) {
        psFitsClose(stack->imageFits->data[i]);
        psFitsClose(stack->maskFits->data[i]);
        psFitsClose(stack->varianceFits->data[i]);
        stack->imageFits->data[i] = stack->maskFits->data[i] = stack->varianceFits->data[i] = NULL;
    }
    psFree(stack->imageFits);
    psFree(stack->maskFits);
    psFree(stack->varianceFits);
    return;
}

ppStackThreadData *ppStackThreadDataSetup(const psArray *cells, const psArray *imageNames,
                                          const psArray *maskNames, const psArray *varianceNames,
                                          const psArray *covariances, const pmConfig *config)
{
    PS_ASSERT_ARRAY_NON_NULL(cells, NULL);
    PS_ASSERT_ARRAYS_SIZE_EQUAL(cells, imageNames, NULL);
    PS_ASSERT_ARRAYS_SIZE_EQUAL(cells, maskNames, NULL);
    PS_ASSERT_ARRAYS_SIZE_EQUAL(cells, varianceNames, NULL);
    PS_ASSERT_ARRAYS_SIZE_EQUAL(cells, covariances, NULL);

    ppStackThreadData *stack = psAlloc(sizeof(ppStackThreadData)); // Thread data, to return
    psMemSetDeallocator(stack, (psFreeFunc)stackThreadDataFree);

    stack->lastScan = 0;

    int numInputs = cells->n;           // Number of inputs
    stack->imageFits  = psArrayAlloc(numInputs);
    stack->maskFits   = psArrayAlloc(numInputs);
    stack->varianceFits = psArrayAlloc(numInputs);
    for (int i = 0; i < numInputs; i++) {
        if (!cells->data[i]) {
            // Bad image
            continue;
        }
        // Resolved names
        psString imageResolved = pmConfigConvertFilename(imageNames->data[i], config, false, false);
        psString maskResolved = pmConfigConvertFilename(maskNames->data[i], config, false, false);
        psString varianceResolved = pmConfigConvertFilename(varianceNames->data[i], config, false, false);
        stack->imageFits->data[i] = psFitsOpen(imageResolved, "r");
        stack->maskFits->data[i] = psFitsOpen(maskResolved, "r");
        stack->varianceFits->data[i] = psFitsOpen(varianceResolved, "r");
        psFree(imageResolved);
        psFree(maskResolved);
        psFree(varianceResolved);
        if (!stack->imageFits->data[i] || !stack->maskFits->data[i] || !stack->varianceFits->data[i]) {
            psError(PS_ERR_UNKNOWN, false, "Unable to open convolved files %s, %s, %s",
                    (char*)imageNames->data[i], (char*)maskNames->data[i], (char*)varianceNames->data[i]);
            return NULL;
        }
    }

    int numThreads = psMetadataLookupS32(NULL, config->arguments, "-threads"); // Number of threads

    // Generate readouts for each input file in each file group
    psArray *threads = psArrayAlloc(numThreads + 1); // A stack for each thread
    for (int i = 0; i < threads->n; i++) {
        psArray *readouts = psArrayAlloc(numInputs); // Input readouts
        for (int j = 0; j < numInputs; j++) {
            pmCell *cell = cells->data[j]; // Cell with data
            if (!cell) {
                continue;
            }
            pmReadout *ro = pmReadoutAlloc(cell); // Readout for thread
            ro->covariance = psMemIncrRefCounter(covariances->data[j]);
            readouts->data[j] = ro;
        }
        threads->data[i] = ppStackThreadAlloc(readouts);
        psFree(readouts);               // Drop reference
    }
    stack->threads = threads;

    return stack;
}


ppStackThread *ppStackThreadRead(bool *status, ppStackThreadData *stack, pmConfig *config, int numChunk,
                                 int overlap)
{
    assert(status);
    PS_ASSERT_PTR_NON_NULL(stack, NULL);
    PS_ASSERT_PTR_NON_NULL(config, NULL);

    *status = true;

    psArray *threads = stack->threads;  // Threads for reading
    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSTACK_RECIPE); // Recipe
    if (!recipe) {
        psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find recipe %s", PPSTACK_RECIPE);
        return NULL;
    }
    int rows = psMetadataLookupS32(NULL, recipe, "ROWS"); // Number of rows to read per chunk
    if (rows <= 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, false, "ROWS is not set in the recipe.");
        return NULL;
    }

    // Select an available group
    while (true) {
        // check for any groups which can read data
        for (int j = 0; j < threads->n; j++) {
            ppStackThread *thread = threads->data[j];
            if (thread->read) {
                continue;
            }

            thread->firstScan = stack->lastScan; // Overlap is taken care of in pmReadoutReadChunk
            thread->lastScan = stack->lastScan + rows;
            stack->lastScan = thread->lastScan;

            psArray *readouts = thread->readouts;

            psTimerStart ("ppStackReadChunk");

            psTrace("ppStack", 2, "Reading data for chunk %d into group %d....\n", numChunk, j);
            for (int i = 0; i < readouts->n; i++) {
                pmReadout *ro = readouts->data[i]; // Input readout
                if (!ro) {
                    continue;
                }

                // override the recorded last scan
                ro->thisImageScan  = thread->firstScan;
                ro->thisVarianceScan = thread->firstScan;
                ro->thisMaskScan   = thread->firstScan;
                ro->lastImageScan  = thread->lastScan;
                ro->lastMaskScan   = thread->lastScan;
                ro->lastVarianceScan = thread->lastScan;
                ro->forceScan      = true;

                psFits *imageFits  = stack->imageFits->data[i]; // FITS file for image
                psFits *maskFits   = stack->maskFits->data[i]; // FITS file for mask
                psFits *varianceFits = stack->varianceFits->data[i]; // FITS file for variance


                bool keepReading = false;
                if (pmReadoutMore(ro, imageFits, 0, rows, config)) {
                    keepReading = true;
                    if (!pmReadoutReadChunk(ro, imageFits, 0, rows, overlap, config)) {
                        psError(PS_ERR_IO, false, "Unable to read chunk %d for file PPSTACK.INPUT %d",
                                numChunk, i);
                        *status = false;
                        return NULL;
                    }
                }

                if (pmReadoutMoreMask(ro, maskFits, 0, rows, config)) {
                    keepReading = true;
                    if (!pmReadoutReadChunkMask(ro, maskFits, 0, rows, overlap, config)) {
                        psError(PS_ERR_IO, false, "Unable to read chunk %d for file PPSTACK.INPUT.MASK %d",
                                numChunk, i);
                        *status = false;
                        return NULL;
                    }
                }

                if (pmReadoutMoreVariance(ro, varianceFits, 0, rows, config)) {
                    keepReading = true;
                    if (!pmReadoutReadChunkVariance(ro, varianceFits, 0, rows, overlap, config)) {
                        psError(PS_ERR_IO, false,
                                "Unable to read chunk %d for file PPSTACK.INPUT.VARIANCE %d",
                                numChunk, i);
                        *status = false;
                        return NULL;
                    }
                }
                if (!keepReading) {
                    return NULL;
                }
            }

            thread->read = thread->busy = true;
            return thread;
        }

        // Check for threads that are ready to read
        bool wait = true;
        for (int j = 0; j < threads->n; j++) {
            ppStackThread *thread = threads->data[j];
            if (thread->busy) {
                continue;
            }
            thread->read = false;
            wait = false;
        }
        if (wait) {
            // No threads currently available
            usleep(THREAD_WAIT);
        }
    }
    return NULL;
}


void ppStackThreadInit(void)
{
    static bool threaded = false;       // Are we running threaded?
    if (threaded) {
        psAbort("Already running threaded.");
    }
    threaded = true;

    {
        psThreadTask *task = psThreadTaskAlloc("PPSTACK_INITIAL_COMBINE", 6);
        task->function = &ppStackReadoutInitialThread;
        psThreadTaskAdd(task);
        psFree(task);
    }

    {
        psThreadTask *task = psThreadTaskAlloc("PPSTACK_INSPECT", 2);
        task->function = &ppStackInspect;
        psThreadTaskAdd(task);
        psFree(task);
    }

    {
        psThreadTask *task = psThreadTaskAlloc("PPSTACK_FINAL_COMBINE", 7);
        task->function = &ppStackReadoutFinalThread;
        psThreadTaskAdd(task);
        psFree(task);
    }

    return;
}
