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

#include <stdio.h>
#include "pslib.h"
#include "stac.h"
#include "stacConfig.h"
#include <sys/time.h>

int main(int argc, char **argv)
{
    double startTime;
    startTime = getTime();

#if 0
    psMemAllocateCallbackSet(stacMemPrint);
    psMemAllocateCallbackSetID(185);
    psMemFreeCallbackSet(stacMemPrint);
    psMemFreeCallbackSetID(185);
#endif

    // Set trace levels
    (void)psTraceSetLevel(".",0);
    (void)psTraceSetLevel("stac",10);
    (void)psTraceSetLevel("stac.checkMemory",10);
    (void)psTraceSetLevel("stac.config",10);
    (void)psTraceSetLevel("stac.read",10);
    (void)psTraceSetLevel("stac.invertMaps",10);
    (void)psTraceSetLevel("stac.errors",10);
    (void)psTraceSetLevel("stac.transform",10);
    (void)psTraceSetLevel("stac.combine",10);
    (void)psTraceSetLevel("stac.rejection",10);
    (void)psTraceSetLevel("stac.time",10);
    (void)psTraceSetLevel("stac.area",10);
    (void)psTraceSetLevel("stac.size",10);
    (void)psTraceSetLevel("stac.scales",7);

    // Set logging level
    psLogSetLevel(9);

    // Parse command line
    stacConfig *config = stacParseConfig(argc, argv);

    // Read input files
    psArray *headers = NULL;            // Array of headers
    psArray *inputs = stacReadImages(&headers, config->inputs);

    // Generate masks
    psArray *masks = psArrayAlloc(inputs->n);
    for (int i = 0; i < inputs->n; i++) {
        psImage *image = inputs->data[i]; // Image for which to get mask
        psImage *mask = psImageAlloc(image->numCols, image->numRows, PS_TYPE_U8);
        for (int y = 0; y < image->numRows; y++) {
            for (int x = 0; x < image->numCols; x++) {
                if (image->data.F32[y][x] <= config->bad) {
                    mask->data.U8[y][x] = 1;
                } else {
                    mask->data.U8[y][x] = 0;
                }
            }
        }
        masks->data[i] = mask;
    }

    // Read maps
    psArray *maps = stacReadMaps(config->inputs);

    psTrace("stac.time", 1, "I/O completed at %f seconds\n", getTime() - startTime);

    // Get size, if not input
    if (config->outnx == 0 || config->outny == 0) {
        psVector *xSize = psVectorAlloc(inputs->n, PS_TYPE_S32); // Sizes of images in x
        psVector *ySize = psVectorAlloc(inputs->n, PS_TYPE_S32); // Sizes of images in y
        for (int i = 0; i < inputs->n; i++) {
            psImage *image = inputs->data[i]; // The i-th image
            xSize->data.S32[i] = image->numCols;
            ySize->data.S32[i] = image->numRows;
        }
        stacSize(&config->outnx, &config->outny, &config->xMapDiff, &config->yMapDiff, xSize, ySize, maps);
        psFree(xSize);
        psFree(ySize);
    }

    // Fix up the headers
    for (int i = 0; i < headers->n; i++) {
        psMetadata *header = headers->data[i];
        psMetadataAddS32(header, PS_LIST_TAIL, "NAXIS1", PS_META_REPLACE, "Number of pixels in x",
                         config->outnx);
        psMetadataAddS32(header, PS_LIST_TAIL, "NAXIS2", PS_META_REPLACE, "Number of pixels in y",
                         (int)config->outny);
        psMetadataAddS32(header, PS_LIST_TAIL, "BITPIX", PS_META_REPLACE, "Bits per pixel", -32);
#if 0
        bool mdok = true;               // Result of MD lookup
        float crpix1 = psMetadataLookupF32(&mdok, header, "CRPIX1");
        if (mdok && !isnan(crpix1)) {
            psMetadataAddF32(header, PS_LIST_TAIL, "CRPIX1", PS_META_REPLACE,
                             "WCS Coordinate reference pixel", crpix1 -
#endif
    }

    // Invert maps
    psArray *inverseMaps = stacInvertMaps(maps, config->outnx, config->outny);

    // Generate errors
    psArray *errors = stacErrorImages(inputs, config->gain, config->readnoise);
#ifdef TESTING
    // Write error images out to check
    for (int i = 0; i < inputs->n; i++) {
        char errName[MAXCHAR];          // Filename of error image
        sprintf(errName,"%s.err", (char*)config->inputs->data[i]);
        psFits *errorFile = psFitsOpen(errName, "w");
        if (!psFitsWriteImage(errorFile, NULL, errors->data[i], 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", errName);
        }
        psTrace("stac", 1, "Error image written to %s\n", errName);
        psFitsClose(errorFile);
    }
#endif

    // Transform inputs and errors
    psArray *transformed = NULL;
    psArray *transformedErrors = NULL;
    (void)stacTransform(&transformed, &transformedErrors, inputs, inverseMaps, errors, masks, NULL, NULL, NULL,
                        config->outnx, config->outny);
    psTrace("stac.time",1,"Transformation completed at %f seconds\n", getTime() - startTime);

#ifdef TESTING
    // Write transformed images out to check
    for (int i = 0; i < inputs->n; i++) {
        char shiftName[MAXCHAR];        // Filename of shift image
        char errName[MAXCHAR];          // Filename of error image
        sprintf(shiftName,"%s.shift1", (char*)config->inputs->data[i]);
        sprintf(errName,"%s.shifterr1",(char*)config->inputs->data[i]);
        psFits *shiftFile = psFitsOpen(shiftName, "w");
        psFits *errFile = psFitsOpen(errName, "w");
        psImage *trans = transformed->data[i]; // Transformed image
        psImage *transErr = transformedErrors->data[i]; // Transformed error image
        if (!psFitsWriteImage(shiftFile, NULL, trans, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", shiftName);
        }
        psTrace("stac", 1, "Shifted image written to %s\n", shiftName);
        if (!psFitsWriteImage(errFile, NULL, transErr, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", errName);
        }
        psTrace("stac", 1, "Shifted error image written to %s\n", errName);
        psFitsClose(shiftFile);
        psFitsClose(errFile);
    }
#endif

    psVector *scales = NULL;            // Relative scales between images
    psVector *offsets = NULL;           // Offsets between images
    (void)stacScales(&scales, &offsets, transformed, config->starFile, config->starMapFile, config->xMapDiff,
                     config->yMapDiff, config->aper);

#ifdef TESTING
    for (int i = 0; i < scales->n; i++) {
        psTrace("stac", 3, "Image %d: scale %f, offset %f\n", i, scales->data.F32[i], offsets->data.F32[i]);
    }
#endif

    // Rescale the images
    (void)stacRescale(transformed, transformedErrors, NULL, scales, offsets);
    // Set the saturation and bad values
    psVector *saturated = psVectorAlloc(transformed->n, PS_TYPE_F32); // Saturation limits
    psVector *bad = psVectorAlloc(transformed->n, PS_TYPE_F32); // Bad limits
    float minSaturated = INFINITY;      // The minimum saturation level
    float maxBad = -INFINITY;           // The maximum bad level
    for (int i = 0; i < transformed->n; i++) {
        saturated->data.F32[i] = (config->saturated - offsets->data.F32[i]) / scales->data.F32[i];
        if (saturated->data.F32[i] < minSaturated) {
            minSaturated = saturated->data.F32[i];
        }
        bad->data.F32[i] = (config->bad - offsets->data.F32[i]) / scales->data.F32[i];
        if (bad->data.F32[i] > maxBad) {
            maxBad = bad->data.F32[i];
        }
    }
    psTrace("stac", 1, "Saturation level: %f\n", minSaturated);
    psTrace("stac", 1, "Bad level: %f\n", maxBad);

    // Save shifted images
    if (config->saveShifts) {
        psImage *image = NULL;          // Copy of image, to remove NAN for output
        for (int i = 0; i < transformed->n; i++) {
            char shiftName[MAXCHAR];    // Filename of shifted image
            sprintf(shiftName, "%s.shift", (const char*)config->inputs->data[i]);
            psFits *shiftFile = psFitsOpen(shiftName, "w");
            image = psImageCopy(NULL, transformed->data[i], PS_TYPE_F32);
            (void)psImageClipNaN(image, 0.0);
            if (!psFitsWriteImage(shiftFile, headers->data[i], image, 0, NULL)) {
                psErrorStackPrint(stderr, "Unable to write image: %s\n", shiftName);
            }
            psTrace("stac", 1, "Shifted image %d written to %s\n", i, shiftName);
            psFitsClose(shiftFile);
            psFree(image);
        }
    }

    // Combine with rejection
    psArray *rejected = NULL;
    psImage *combined = NULL;
    (void)stacCombine(&combined, &rejected, transformed, transformedErrors, config->nReject, NULL, saturated,
                      bad, config->reject);

    psTrace("stac.time",1,"First combination completed at %f seconds\n", getTime() - startTime);


#ifdef TESTING
    // Write rejection images out to check
    for (int i = 0; i < rejected->n; i++) {
        char rejName[MAXCHAR];  // Filename of rejection image
        sprintf(rejName, "%s.shiftrej", (char*)config->inputs->data[i]);

        psFits *rejFile = psFitsOpen(rejName, "w");
        if (!psFitsWriteImage(rejFile, NULL, rejected->data[i], 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", rejName);
        }
        psTrace("stac", 1, "Rejection image written to %s\n", rejName);
        psFitsClose(rejFile);
    }

    // Write out pre-combined image
    char preName[MAXCHAR];              // Filename of precombined image
    sprintf(preName, "%s.pre", config->output);
    psFits *preFile = psFitsOpen(preName, "w");
    if (!psFitsWriteImage(preFile, NULL, combined, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write image: %s\n", preName);
    }
    psTrace("stac", 1, "Pre-combined image written to %s\n", preName);
    psFitsClose(preFile);
#endif

    // Get regions of interest in the source frame
    psArray *regions = psArrayAlloc(inputs->n); // Array of images denoting regions of interest
    for (int i = 0; i < inputs->n; i++) {
        regions->data[i] = stacAreaOfInterest(rejected->data[i], inverseMaps->data[i],
                                              ((psImage*)(inputs->data[i]))->numCols,
                                              ((psImage*)(inputs->data[i]))->numRows);
#ifdef TESTING
        char regionName[MAXCHAR];       // Filename of region image
        sprintf(regionName,"%s.region",(char*)config->inputs->data[i]);
        psFits *regionFile = psFitsOpen(regionName, "w");
        if (!psFitsWriteImage(regionFile, NULL, regions->data[i], 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", regionName);
        }
        psTrace("stac", 1, "Region image written to %s\n", regionName);
        psFitsClose(regionFile);
#endif
    }

    // Transform rejected pixels to source frame
    psArray *rejectedSource = stacRejection(inputs, rejected, regions, maps, inverseMaps, config->frac,
                                            config->grad, config->inputs);

    // Get regions of interest in the output frame
    psImage *combineRegion = psImageAlloc(config->outnx, config->outny, PS_TYPE_U8);
    psImageInit(combineRegion, 0);
    for (int i = 0; i < inputs->n; i++) {
        psImage *region = stacAreaOfInterest(rejectedSource->data[i], maps->data[i], config->outnx,
                                             config->outny);
        psBinaryOp(combineRegion, combineRegion, "+", region);
        psFree(region);

        // Do OR of masks
        psImage *mask = masks->data[i]; // Mask of input image (bad columns etc)
        psImage *cr = rejectedSource->data[i]; // Mask of CRs
        for (int y = 0; y < mask->numRows; y++) {
            for (int x = 0; x < mask->numCols; x++) {
                if (cr->data.U8[y][x] > 0) {
                    mask->data.U8[y][x] = 1;
                }
            }
        }
    }

    // Redo transformation with the masks and scales/offsets
    (void)stacTransform(&transformed, &transformedErrors, inputs, inverseMaps, errors, masks,
                        combineRegion, scales, offsets, config->outnx, config->outny);

#ifdef TESTING
    // Write transformed images out to check
    for (int i = 0; i < inputs->n; i++) {
        char shiftName[MAXCHAR];        // Filename of shift image
        char errName[MAXCHAR];          // Filename of error image
        sprintf(shiftName,"%s.shift2", (char*)config->inputs->data[i]);
        sprintf(errName,"%s.shifterr2",(char*)config->inputs->data[i]);
        psFits *shiftFile = psFitsOpen(shiftName, "w");
        psFits *errFile = psFitsOpen(errName, "w");
        psImage *trans = transformed->data[i]; // Transformed image
        psImage *transErr = transformedErrors->data[i]; // Transformed error image
        if (!psFitsWriteImage(shiftFile, NULL, trans, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", shiftName);
        }
        psTrace("stac", 1, "Shifted image written to %s\n", shiftName);
        if (!psFitsWriteImage(errFile, NULL, transErr, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", errName);
        }
        psTrace("stac", 1, "Shifted error image written to %s\n", errName);
        psFitsClose(shiftFile);
        psFitsClose(errFile);
    }
#endif

    // Combine the newly-transformed CR-free images, no rejection
    psFree(rejected);
    rejected = NULL;
    (void)stacCombine(&combined, &rejected, transformed, transformedErrors, 0, combineRegion, saturated,
                      bad, config->reject);

    // Write output image
    psFits *outFile = psFitsOpen(config->output, "w");
    int numPix = psImageClipNaN(combined, 0.0);
    if (numPix > 0) {
        psTrace("stac", 3, "Clipping %d NaN pixels to zero.\n", numPix);
    }
    if (!psFitsWriteImage(outFile, headers->data[0], combined, 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write image: %s\n", config->output);
    }
    psTrace("stac", 1, "Combined image written to %s\n", config->output);
    psFitsClose(outFile);

    // Free everything I've used
    stacConfigFree(config);
    psFree(scales);
    psFree(offsets);
    psFree(combineRegion);
    psFree(regions);
    psFree(inputs);
    psFree(headers);
    psFree(maps);
    psFree(inverseMaps);
    psFree(errors);
    psFree(masks);
    psFree(transformedErrors);
    psFree(transformed);
    psFree(rejectedSource);
    psFree(combined);
    psFree(saturated);
    psFree(bad);

    psTrace("stac.time",1,"Final combination completed at %f seconds\n", getTime() - startTime);

#ifdef TESTING
    // Check memory for leaks, corruption
    stacCheckMemory();
#endif

    exit(EXIT_SUCCESS);
}
