# include "pswarp.h"

# define THREADED 1
# define SOURCE_ARRAY_BUFFER 100         // Size to grow the array of sources at a time

// NOTE: in this function, the coordinates are transformed from the OUTPUT to the INPUT
bool pswarpTransformReadout_Threaded(pmReadout *output, pmReadout *input, pmConfig *config)
{
    // XXX this implementation currently ignores the use of the region
    psImage *region = NULL;

    psTimerStart ("warp");

    // Get warp parameters
    bool mdok;
    int nGridX = psMetadataLookupS32(NULL, config->arguments, "GRID.NX");
    int nGridY = psMetadataLookupS32(NULL, config->arguments, "GRID.NY");
    psImageInterpolateMode interpolationMode = psMetadataLookupS32(NULL, config->arguments, "INTERPOLATION.MODE");

    // load the recipe
    psMetadata *recipe = psMetadataLookupPtr (NULL, config->recipes, PSWARP_RECIPE);
    psAssert (recipe, "missing recipe %s", PSWARP_RECIPE);

    // output mask bits
    psMaskType maskIn   = psMetadataLookupU8(&mdok, recipe, "MASK.INPUT");
    psMaskType maskPoor = pmConfigMaskGet("POOR.WARP", config);
    psMaskType maskBad  = pmConfigMaskGet("BAD.WARP", config);
    psAssert (mdok, "MASK.INPUT was not defined");

    // Flux fraction for "poor"
    float poorFrac = psMetadataLookupF32(NULL, config->arguments, "POOR.FRAC");

    // find the output pixel range
    int minX, minY, maxX, maxY;
    pswarpMatchRange (&minX, &minY, &maxX, &maxY, input, output);

    // pswarpMapGridFromImage builds a set of locally-linear maps which convert the
    // output coordinates to input coordinates
    pswarpMapGrid *grid = pswarpMapGridFromImage (input, output, nGridX, nGridY);

    // XXX need to modify the grid based on this result and force the maxError < XXX
    double maxError = pswarpMapGridMaxError (grid);
    psLogMsg ("pswarp", 3, "maximum error using this grid sampling: %f\n", maxError);

    int gridX, gridY;
    pswarpMapGridSetGrid (grid, minX, minY, &gridX, &gridY);
# if (0)
    // XXX these asserts probably belong in the function (don't have outCol0,Row0 anyway)
    assert ((int)(minX - outCol0) >= 0);
    assert ((int)(maxX - outCol0) <= output->image->numCols);
    assert ((int)(minY - outRow0) >= 0);
    assert ((int)(maxY - outRow0) <= output->image->numRows);
# endif

    psImage *inImage = input->image;    // Input image
    psImage *inVar   = input->weight;   // Input weight map
    psImage *inMask  = input->mask;     // Input mask

    // Interpolation options
    psImageInterpolateOptions *interp = psImageInterpolateOptionsAlloc(interpolationMode, inImage,
                                                                       inVar, inMask, maskIn, NAN, NAN,
                                                                       maskBad, maskPoor, poorFrac);

    if (inVar && !output->weight) {
        output->weight = psImageAlloc(output->image->numCols, output->image->numRows, PS_TYPE_F32);
        psImageInit(output->weight, NAN);
    }
    if ((inMask || maskPoor || maskBad) && !output->mask) {
        output->mask = psImageAlloc(output->image->numCols, output->image->numRows, PS_TYPE_MASK);
        psImageInit(output->mask, maskBad);
    }

    // total number of good pixels across all tiles (summed below)
    int goodPixels = 0;

    // create jobs and supply them to the threads
    for (gridY = 0; gridY < nGridX; gridY++) {
        for (gridX = 0; gridX < nGridY; gridX++) {

            pswarpTransformTileArgs *args = pswarpTransformTileArgsAlloc();

            // these items are just views to the data; they are not freed with args
            args->input = input;
            args->output = output;
            args->grid = grid;
            args->interp = interp;
            args->region = region;

            args->gridX = gridX;
            args->gridY = gridY;
            args->goodPixels = 0;

# if (THREADED)
            // allocate a job
            psThreadJob *job = psThreadJobAlloc ("PSWARP_TRANSFORM_TILE", 0);

            // construct the arguments for this job
            // job is pswarpTransformTile (gridX, gridY);
            psArrayAdd (job->args, 1, args);
            psThreadJobAddPending (job);
            psFree(job);
# else
            pswarpTransformTile (args);
            goodPixels += args->goodPixels;
# endif
            psFree (args);
        }
    }

# if (THREADED)
    // wait here for the threaded jobs to finish
    if (!psThreadPoolWait(false)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to interpolate image.");
        return false;
    }

    // each job records its own goodPixel values; sum them here
    // we have only supplied one type of job, so we can assume the types here
    psThreadJob *job = NULL;
    while ((job = psThreadJobGetDone()) != NULL) {
        pswarpTransformTileArgs *args = job->args->data[0];
        goodPixels += args->goodPixels;
        psFree (job);
    }
# endif
    psFree(interp);
    psFree(grid);


    // Store the variance factor and number of good pixels
    if (goodPixels > 0) {
        float varFactor = psImageInterpolateVarianceFactor(inImage->numCols / 2.0 + inImage->col0,
                                                           inImage->numRows / 2.0 + inImage->row0,
                                                           interp); // Variance factor: large --> small scale
        psMetadataItem *vfItem = psMetadataLookup(output->analysis, PSWARP_ANALYSIS_VARFACTOR);
        if (vfItem) {
            psMetadataItem *goodpixItem = psMetadataLookup(output->analysis, PSWARP_ANALYSIS_GOODPIX);
            psAssert(goodpixItem, "It should be where we left it!");
            psAssert(vfItem->type == PS_TYPE_F32 && goodpixItem->type == PS_TYPE_S64,
                     "Should be the type we said.");

            vfItem->data.F32 += varFactor * goodPixels;
            goodpixItem->data.S64 += goodPixels;
        } else {
            psMetadataAddF32(output->analysis, PS_LIST_TAIL, PSWARP_ANALYSIS_VARFACTOR, 0,
                             "Variance factor weighted by the good pixels", varFactor * goodPixels);
            psMetadataAddS64(output->analysis, PS_LIST_TAIL, PSWARP_ANALYSIS_GOODPIX, 0,
                             "Number of good pixels", goodPixels);
        }
    }

    // Transform sources
    psArray *inSources = psMetadataLookupPtr(&mdok, input->analysis, "PSPHOT.SOURCES"); // Sources in source
    if (goodPixels > 0 && mdok && inSources) {
        pswarpMapGrid *sourceGrid = pswarpMapGridFromImage(output, input, nGridX, nGridY); // Grid for sources

        psArray *outSources = psMemIncrRefCounter(psMetadataLookupPtr(&mdok, output->analysis,
                                                                      "PSPHOT.SOURCES")); // Target sources
        if (!outSources) {
            outSources = psArrayAllocEmpty(SOURCE_ARRAY_BUFFER);
            psMetadataAddPtr(output->analysis, PS_LIST_TAIL, "PSPHOT.SOURCES", PS_DATA_ARRAY,
                             "Warped sources", outSources);
        }

        for (int i = 0; i < inSources->n; i++) {
            pmSource *source = inSources->data[i]; // Source of interest
            pmModel *model = source->modelPSF; // Model for this source
            float xIn, yIn;             // Coordinates of source
            xIn = model->params->data.F32[PM_PAR_XPOS] - inImage->col0;
            yIn = model->params->data.F32[PM_PAR_YPOS] - inImage->row0;
            int xGrid, yGrid;           // Grid coordinates for local map
            if (!pswarpMapGridSetGrid(sourceGrid, xIn + 0.5, yIn + 0.5, &xGrid, &yGrid)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to get grid coordinates for source at %f,%f\n",
                        xIn, yIn);
                psFree(outSources);
                psFree(sourceGrid);
                return false;
            }
            if (xGrid < 0 || xGrid >= sourceGrid->nXpts || yGrid < 0 || yGrid >= sourceGrid->nYpts) {
                // It's not even on the grid
                continue;
            }

            pswarpMap *map = sourceGrid->maps[xGrid][yGrid]; // Locally linear transformation
            double xOut, yOut;          // Output coordinates
            if (!pswarpMapApply(&xOut, &yOut, map, xIn + 0.5, yIn + 0.5)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to transform coordinates for source at %f,%f\n",
                        xIn, yIn);
                psFree(outSources);
                psFree(sourceGrid);
                return false;
            }
            xOut += output->image->col0 - 0.5;
            yOut += output->image->row0 - 0.5;
            if (xOut < minX || xOut > maxX || yOut < minY || yOut > maxY) {
                // It's not in the output image
                continue;
            }

            // Generate the new source in the output frame
            pmSource *new = pmSourceAlloc(); // New source
            new->peak = pmPeakAlloc(xOut, yOut, source->peak->flux, PM_PEAK_LONE);
            new->peak->flux = source->peak->flux;
            new->type = PM_SOURCE_TYPE_STAR;
            new->psfMag = source->psfMag;
            new->errMag = source->errMag;
            new->sky = source->sky;
            new->skyErr = source->skyErr;
            new->pixWeight = source->pixWeight;
            new->modelPSF = pmModelAlloc(source->modelPSF->type);

#if 0
            // XXX Note that this will not set the correct axes
            pmPSF_AxesToModel(new->modelPSF->params->data.F32,
                              pmPSF_ModelToAxes(source->modelPSF->params->data.F32, 20.0));
#endif

            // Propagate the position erorrs
            float dxIn, dyIn;           // Errors in input coordinates
            dxIn = model->dparams->data.F32[PM_PAR_XPOS];
            dyIn = model->dparams->data.F32[PM_PAR_YPOS];

            float dxOut, dyOut;         // Errors in output coordinates
            dxOut = sqrt(PS_SQR(map->Xx * dxIn) + PS_SQR(map->Xy * dyIn));
            dyOut = sqrt(PS_SQR(map->Yx * dxIn) + PS_SQR(map->Yy * dyIn));

            new->modelPSF->params->data.F32[PM_PAR_XPOS] = xOut;
            new->modelPSF->params->data.F32[PM_PAR_YPOS] = yOut;
            new->modelPSF->dparams->data.F32[PM_PAR_XPOS] = dxOut;
            new->modelPSF->dparams->data.F32[PM_PAR_YPOS] = dyOut;

            psArrayAdd(outSources, SOURCE_ARRAY_BUFFER, new);
            psFree(new);                // Drop reference
        }
        psFree(sourceGrid);
        psFree(outSources);             // Drop reference
    }

    output->data_exists = true;
    output->parent->data_exists = true;
    output->parent->parent->data_exists = true;

    psLogMsg ("pswarp", 3, "warping analysis: %f sec\n", psTimerMark ("warp"));

    return true;
}
