#include "streaksremove.h"

static pmConfig *parseArguments(int, char **);
static bool readAndCopyToOutput(streakFiles *sf, bool restoreImageCube);
static void writeImages(streakFiles *sf, bool exciseImageCube);
static bool readAndCopyToOutput(streakFiles *sf, bool restoreImageCube);

char *streaksProgram = "streaksreplace";

int main(int argc, char *argv[])
{
    bool status;

    psLibInit(NULL);

    pmConfig *config = parseArguments(argc, argv);
    if (!config) {
        psError(PS_ERR_UNKNOWN, false, "failed to parse arguments\n");
        return PS_EXIT_CONFIG_ERROR;
    }

    psMetadata *masks = psMetadataLookupMetadata(&status, config->recipes, "MASKS");
    if (!status) {
        psError(PM_ERR_CONFIG, false, "failed to lookup MASKS in recipes\n");
        return PS_EXIT_CONFIG_ERROR;
    }
    // TODO: perhaps we should save this in the image header and get it from there?
    double maskStreak = (double) psMetadataLookupU8(&status, masks, "STREAK");
    if (!status) {
        psError(PM_ERR_CONFIG, false, "failed to lookup mask value for STREAK in recipes\n");
        return PS_EXIT_CONFIG_ERROR;
    }

    streakFiles *sfiles = openFiles(config, false, argv[0]);

    if (sfiles->stage == IPP_STAGE_RAW) {
        // copy PHU to output files
        copyPHU(sfiles, false, false);

        // advance to the first image extension
        if (!streakFilesNextExtension(sfiles)) {
            psErrorStackPrint(stderr, "failed to advance to first extension of input images");
            exit (PS_EXIT_PROG_ERROR);
        }
    }

    // Iterate through components of input files
    do {
        bool restoreImageCube = false;

        // read the images and copy the data from the inputs to the outputs
        // This also sets up the astrometry and initializes the recovery image(s)

        if (!readAndCopyToOutput(sfiles, restoreImageCube)) {
            // false return value indicates that this extension is not an image and
            // the contents has already been copied to the output file.
            // we don't need to process this extesion for streaks.
            continue;
        }

        // If pixel in recovery image is not NAN, then it contains an excised pixel
        // that has been stored in the recovery image
        for (int y=0; y < sfiles->inImage->numRows; y++) {
            for (int x=0; x < sfiles->inImage->numCols; x++) {
                double recValue = psImageGet(sfiles->recImage->image, x, y);
                if (!isnan(recValue)) {
                    psImageSet(sfiles->inImage->image, x, y, recValue);
                }
                if (sfiles->inMask) {
                    double recMaskValue = psImageGet(sfiles->recMask->image, x, y);
                    if (recValue != maskStreak) {
                        psImageSet(sfiles->inMask->image, x, y, recMaskValue);
                    }
                }
                if (sfiles->inWeight) {
                    double recWeightValue = psImageGet(sfiles->recWeight->image, x, y);
                    if (recWeightValue != NAN) {
                        psImageSet(sfiles->inWeight->image, x, y, recWeightValue);
                    }
                }
            }
        }

        // write out the re-streaked images
        writeImages(sfiles, restoreImageCube);

    } while (streakFilesNextExtension(sfiles));

    // close all files
    closeImages(sfiles);

    if (!replicateOutputs(sfiles)) {
        psErrorStackPrint(stderr, "failed to replicate output files");
        exit(PS_EXIT_UNKNOWN_ERROR);
    }

//    nebServerFree(ourNebServer);
    psFree(config);
    pmConceptsDone();
    pmConfigDone();
    psLibFinalize();

    //  PAU
    fprintf(stderr, "Found %d leaks at %s\n", psMemCheckLeaks (0, NULL, NULL, false), "streaksreplace");

    return 0;
}
static void usage(void)
{
    fprintf(stderr, "USAGE: streaksreplace -stage raw|chip|warp|diff -image IMAGE.fits -tmproot directory\n\n");
    fprintf(stderr, "Optional arguments:\n");
    fprintf(stderr, "\t-mask MASK.fits mask file to re-streak\n");
    fprintf(stderr, "\t-recmask RECOVERYMASK.fits: recovery mask file (required with -mask)\n");
    fprintf(stderr, "\t-weight WEIGHT.fits: weight file to re-streak\n");
    fprintf(stderr, "\t-recweight RECOVERYWEIGHT.fits: recovery weight file (required with -weight)\n");
    fprintf(stderr, "\t-replace: replace the input images with the output\n");
    fprintf(stderr, "\t-remove: remove the original image after processing (requires -replace)\n");

    exit(2);
}

static pmConfig *parseArguments(int argc, char **argv)
{
    pmConfig *config = pmConfigRead(&argc, argv, NULL);
    if (!config) {
        streaksExit("failed to parse arguments", PS_EXIT_CONFIG_ERROR);
    }

    int argnum;
    ippStage stage = IPP_STAGE_NONE;
    if ((argnum = psArgumentGet(argc, argv, "-stage"))) {
        psArgumentRemove(argnum, &argc, argv);
        stage = parseStage(argv[argnum]);
        psMetadataAddS32(config->arguments, PS_LIST_TAIL, "STAGE", 0,
                "pipeline stage for streak removal", stage);
        psArgumentRemove(argnum, &argc, argv);
    } else {
        psError(PS_ERR_UNKNOWN, true, "-stage is required\n");
        usage();
    }

    // If we are asked to replace inputs with the outputs and input image is in nebulous
    // we require mask, weight, and tmproot to all be
    bool nebulousImage = false;
    if ((argnum = psArgumentGet(argc, argv, "-image"))) {
        // for raw and chip level images we use psFits for I/O and ...
        nebulousImage = IN_NEBULOUS(argv[argnum+1]);
        if (CHIP_LEVEL_INPUT(stage)) {
            psArgumentRemove(argnum, &argc, argv);
            psMetadataAddStr(config->arguments, PS_LIST_TAIL, "INPUT", 0,
                    "name of input image", argv[argnum]);
            psArgumentRemove(argnum, &argc, argv);
        } else  {
            // ... for warped images we use pmFPAfiles
            if (!pmConfigFileSetsMD(config->arguments, &argc, argv, "INPUT", "-image", NULL)) { ;
                psError(PS_ERR_UNKNOWN, false, "failed to process -image");
                usage();
            }
        }
    } else {
        psError(PS_ERR_UNKNOWN, true, "-image is required\n");
        usage();
    }

    bool gotReplace = false;
    if ((argnum = psArgumentGet(argc, argv, "-replace"))) {
        gotReplace = true;
        psArgumentRemove(argnum, &argc, argv);
        psMetadataAddBool(config->arguments, PS_LIST_TAIL, "REPLACE", 0, "replace input files",
            true);
    }

    bool gotMask = false;
    if ((argnum = psArgumentGet(argc, argv, "-mask"))) {
        psArgumentRemove(argnum, &argc, argv);
        bool nebulousMask = IN_NEBULOUS(argv[argnum]);
        if (nebulousMask != nebulousImage) {
            psError(PS_ERR_UNKNOWN, true, "mask image must be %snebulous path with %s image path\n",
                nebulousImage ? "" : "non ", nebulousImage ? "nebulous" : "regular");
            usage();
        }
        gotMask = true;
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "INPUT.MASK", 0,
                "name of input mask image", argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    }

    bool gotWeight = false;
    if ((argnum = psArgumentGet(argc, argv, "-weight"))) {
        psArgumentRemove(argnum, &argc, argv);
        bool nebulousWeight = IN_NEBULOUS(argv[argnum]);
        if (nebulousWeight != nebulousImage) {
            psError(PS_ERR_UNKNOWN, true, "weight image must have %snebulous path with %s image path\n",
                nebulousImage ? "" : "non ", nebulousImage ? "nebulous" : "regular");
            usage();
        }
        gotWeight = true;
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "INPUT.WEIGHT", 0,
                "name of input weight image", argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    }

    if ((argnum = psArgumentGet(argc, argv, "-tmproot"))) {
        psArgumentRemove(argnum, &argc, argv);
        bool nebulousTmp = IN_NEBULOUS(argv[argnum]);
        if (gotReplace && (nebulousTmp != nebulousImage)) {
            psError(PS_ERR_UNKNOWN, true, "tmproot must have %snebulous path with %s image path with -replace\n",
                nebulousImage ? "" : "non ", nebulousImage ? "nebulous" : "regular");
            usage();
        }
        psString dir = pathToDirectory(argv[argnum]);
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "OUTPUT", 0, "directory for temporary files",
            dir);
        psArgumentRemove(argnum, &argc, argv);
    } else {
        psError(PS_ERR_UNKNOWN, true, "-tmproot is required\n");
        usage();
    }

    if ((argnum = psArgumentGet(argc, argv, "-recimage"))) {
        psArgumentRemove(argnum, &argc, argv);
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "RECOVERY.IMAGE", 0, "recovery image file name",
            argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    } else {
        psError(PS_ERR_UNKNOWN, true, "-recimage is required\n");
        usage();
    }
    if ((argnum = psArgumentGet(argc, argv, "-recmask"))) {
        psArgumentRemove(argnum, &argc, argv);
        bool nebulousMask = IN_NEBULOUS(argv[argnum]);
        if (nebulousMask != nebulousImage) {
            psError(PS_ERR_UNKNOWN, true, "mask image must be %snebulous path with %s image path\n",
                nebulousImage ? "" : "non ", nebulousImage ? "nebulous" : "regular");
            usage();
        }
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "INPUT.MASK", 0,
                "name of input mask image", argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    } else if (gotMask) {
        psError(PS_ERR_UNKNOWN, true, "-recmask is required with -mask\n");
        usage();
    }

    if ((argnum = psArgumentGet(argc, argv, "-recweight"))) {
        psArgumentRemove(argnum, &argc, argv);
        bool nebulousWeight = IN_NEBULOUS(argv[argnum]);
        if (nebulousWeight != nebulousImage) {
            psError(PS_ERR_UNKNOWN, true, "weight image must have %snebulous path with %s image path\n",
                nebulousImage ? "" : "non ", nebulousImage ? "nebulous" : "regular");
            usage();
        }
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "INPUT.WEIGHT", 0,
                "name of input weight image", argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    } else if (gotWeight) {
        psError(PS_ERR_UNKNOWN, true, "-recweight is required with -weight\n");
        usage();
    }

    if ((argnum = psArgumentGet(argc, argv, "-remove"))) {
        if (!gotReplace) {
            psError(PS_ERR_UNKNOWN, true, "-replace is required with -remove\n");
            usage();
        }
        psArgumentRemove(argnum, &argc, argv);
        psMetadataAddBool(config->arguments, PS_LIST_TAIL, "REMOVE", 0, "remove original files",
            true);
    }

    if (argc != 1) {
        psString unexpectedArguments = NULL;
        for (int i=1; i<argc; i++) {
            psStringAppend(&unexpectedArguments, "%s ", argv[i]);
        }
        psError(PS_ERR_UNKNOWN, true, "unexpected arguments: %s", unexpectedArguments);
        psFree(unexpectedArguments);
        usage();
    }

    return config;
}

// set the image for the output files to the contents of the corresponding input file.
static bool readAndCopyToOutput(streakFiles *sf, bool restoreImageCube)
{
    if (sf->inImage->pmfile) {
        // image data from pmFPAfile (diff or warp file)
        readImageFrom_pmFile(sf);
    } else {
        // image data directly from psFits
        readImage(sf->inImage, sf->extnum, sf->stage, false);
    }
    // read the recovery pixels
    readImage(sf->recImage, sf->extnum, sf->stage, false);

    sf->outImage->header = (psMetadata*) psMemIncrRefCounter(sf->inImage->header);
    addDestreakKeyword(sf->outImage->header, false);

    if (!SFILE_IS_IMAGE(sf->inImage)) {
        // fprintf(stderr, "skipping extnum: %d not an image\n", sf->extnum);
        // XXX: This extension is probably a video table.
        copyTable(sf->outImage, sf->inImage, sf->extnum);
        return false;
    }

    if (sf->inImage->image) {
        setupImageRefs(sf->outImage, NULL, sf->inImage, sf->extnum, false);
    }  else {
        // Image cubes are handeled specially
    }

    // set up the compression parameters
    copyFitsOptions(sf->outImage, sf->recImage, sf->inImage, sf->tiles);

    if (sf->inMask) {
        readImage(sf->inMask, sf->extnum, sf->stage, true);
        readImage(sf->recMask, sf->extnum, sf->stage, true);
        sf->outMask->header = (psMetadata*) psMemIncrRefCounter(sf->inMask->header);
        
        addDestreakKeyword(sf->outMask->header, false);

        // XXX: TODO
        setupImageRefs(sf->outMask, NULL, sf->inMask, sf->extnum, false);

        copyFitsOptions(sf->outMask, NULL, sf->inMask, sf->tiles);
    }

    if (sf->inWeight) {
        readImage(sf->inWeight, sf->extnum, sf->stage, false);
        readImage(sf->recWeight, sf->extnum, sf->stage, false);

        sf->outWeight->header = (psMetadata*) psMemIncrRefCounter(sf->inWeight->header);
        addDestreakKeyword(sf->outWeight->header, false);
        setupImageRefs(sf->outMask, NULL, sf->inMask, sf->extnum, false);

        copyFitsOptions(sf->outWeight, NULL, sf->inWeight, sf->tiles);
    }

    return true;
}

static void writeImages(streakFiles *sf, bool exciseImageCube)
{
    psString extname = NULL;
    if (sf->nHDU > 1) {
        bool mdok;
        extname = psMetadataLookupStr(&mdok, sf->inImage->header, "EXTNAME");
    }
    if (sf->inImage->numZPlanes == 0)  {
        // note exciseing complete images is handled in readAndCopyToOutput
        writeImage(sf->outImage, extname, sf->extnum);
    } else {
#ifdef TODO
        // we have an image cube
        double initValue;
        // otherwise write it to the output
        writeImageCube(sf->outImage, sf->recImage->imagecube, extname, sf->extnum);

        // borrow one of the images from the imagecube and set it to init value
        psImage *image = psArrayGet (sf->inImage->imagecube, 0);
        psMemIncrRefCounter(image);
        psImageInit(image, initValue);
        if (exciseImageCube) {
            sf->outImage->image = image;
            writeImage(sf->outImage, extname, sf->extnum);
        } else {
            // write zero valued image to reccovery
            if (sf->recImage) {
                sf->recImage->image = image;
                writeImage(sf->recImage, extname, sf->extnum);
            }
        }
#endif
        // Assumption: there are no mask and weight images for video cells because they only appear
        // at the raw stage. TODO: insure this
        return;
    }
    if (sf->outMask) {
        writeImage(sf->outMask, extname, sf->extnum);
    }
    if (sf->outWeight) {
        writeImage(sf->outWeight, extname, sf->extnum);
    }
}
