#include "streaksremove.h"

static pmConfig *parseArguments(int argc, char **argv);
static bool readAndCopyToOutput(streakFiles *sf, bool exciseAll);
static StreakPixels * getStreakPixels(streakFiles *sfiles, Streaks *streaks);
static void exciseNonWarpedPixels(streakFiles *sfiles, double newMaskValue);
static bool warpedPixel(streakFiles *sfiles, PixelPos *cellCoord);
static void excisePixel(streakFiles *sfiles, unsigned int x, unsigned int y, bool streak, double newMaskValue);
static void writeImages(streakFiles *sf, bool exciseImageCube);
static bool replicateOutputs(streakFiles *sfiles);

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

    psLibInit(NULL);
    psTimerStart("STREAKSREMOVE");

    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;
    }
    psU8 poorWarp = (double) psMetadataLookupU8(&status, masks, "POOR.WARP");
    if (!status) {
        psError(PM_ERR_CONFIG, false, "failed to lookup mask value for POOR.WARP in recipes\n");
        return PS_EXIT_CONFIG_ERROR;
    }
    // we're setting pixels with any mask bits execpt POOR.WARP to NAN
    psU8 maskMask = ~poorWarp;

    // Does true work here?
    streakFiles *sfiles = openFiles(config, true);

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

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

    int totalPixels = 0;
    int totalMaskedPixels = 0;
    int totalNANdPixels = 0;
    int totalNANdWeights = 0;

    // Iterate through each component of the input (there is only one except for raw images)
    do {
        psTimerStart("PROCESS_COMPONENT");

        // 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, false)) {
            // 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;
        }

        int maskedPixels = 0;
        int nandPixels = 0;
        int nandWeights = 0;

        psImage *image = sfiles->outImage->image;
        psImage *mask = sfiles->outMask->image;
        psImage *weight= sfiles->outWeight->image;

        for (int y=0; y < image->numRows; y++) {
            for (int x=0; x < image->numCols; x++) {
                // these gets are not necessary, we could just set the pixels to nan
                // but I want to get the counts
                double imageVal  = psImageGet(image, x, y);
                double weightVal = psImageGet(weight, x, y);
                psU8 maskVal   = psImageGet(mask, x, y);
                if (maskVal & maskMask) {
                    ++maskedPixels;
                    if (!isnan(imageVal)) {
                        ++nandPixels;
                        psImageSet(image, x, y, NAN);
                    }
                    if (!isnan(weightVal)) {
                        ++nandWeights;
                        psImageSet(weight, x, y, NAN);
                    }
                }
            }
        }
        totalPixels += sfiles->inImage->numRows * sfiles->inImage->numCols;
        totalMaskedPixels += maskedPixels;
        totalNANdPixels += nandPixels;
        totalNANdWeights += nandWeights;

        // write out the destreaked temporary images and the recovery images
        writeImages(sfiles, false);

        printf("time to process component %d: %f\n", sfiles->extnum, psTimerClear("PROCESS_COMPONENT"));
    } while (streakFilesNextExtension(sfiles));

    printf("pixels:        %10ld\n", totalPixels);
    printf("masked pixels: %10ld %4.2f%%\n", totalMaskedPixels, 100. * totalMaskedPixels / totalPixels);
    printf("nand pixels:   %10ld %4.2f%%\n", totalNANdPixels, 100. * totalNANdPixels / totalPixels);
    printf("nand weights:  %10ld %4.2f%%\n", totalNANdWeights, 100. * totalNANdWeights / totalPixels);

    psTimerStart("CLOSE_IMAGES");
    // close all files
    closeImages(sfiles);
    printf("time to close images: %f\n", psTimerClear("CLOSE_IMAGES"));

#ifdef NOTYET
    if (!replicateOutputs(sfiles)) {
        psError(PS_ERR_UNKNOWN, false, "failed to replicate output files");
        psErrorStackPrint(stderr, "");
        exit(PS_EXIT_UNKNOWN_ERROR);
    }

    if (psMetadataLookupBool(&status, config->arguments, "REPLACE")) {
        //     swap the instances for the input and output
        //     Note this is a database operation. No file I/O is performed
        if (!swapOutputsToInputs(sfiles)) {
            psError(PS_ERR_UNKNOWN, false, "failed to swap files");

            // XXX: Now what? I guess swapOutputsToInputs will need to undo anything that
            // it has done and give a detailed report of what happened

            psErrorStackPrint(stderr, "");
            exit(PS_EXIT_UNKNOWN_ERROR);
        }

        if (psMetadataLookupBool(&status, config->arguments, "REMOVE")) {
            //      delete the temporary storage objects (which now points to the original image(s)
            if (!deleteTemps(sfiles)) {
                psError(PS_ERR_UNKNOWN, false, "failed to delete temporary files");
                // XXX: Now what? At this point the output files have been swapped, so we can't
                // repeat the operation.

                // Returning error status here is problematic. The inputs have been streak removed
                // but they're still lying around
                // Maybe just print an error message and
                // let other system tools clean up
                psErrorStackPrint(stderr, "");
                exit(PS_EXIT_UNKNOWN_ERROR);
            }
        }
    }
#endif  // REPLACE, REMOVE
    printf("time to run streaksrelease: %f\n", psTimerClear("STREAKSREMOVE"));

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

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


    return 0;
}

static void usage()
{
    psErrorStackPrint(stderr, "\n");
    fprintf(stderr, "USAGE: streaksrelease -stage raw|chip|warp|diff -image IMAGE.fits -outroot directory\n\n");
    fprintf(stderr, "Other arguments:\n");
    fprintf(stderr, "\t-class_id class_id: class_id (required for raw and chip stages)\n");
    fprintf(stderr, "\t-mask MASK.fits mask file to de-streak\n");
    fprintf(stderr, "\t-weight WEIGHT.fits: weight file to de-streak\n");
    fprintf(stderr, "\t-chip_mask MASK.fits: name of mask for chip stage (camera stage mask is passed tih -mask)\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 ((argnum = psArgumentGet(argc, argv, "-class_id"))) {
        psArgumentRemove(argnum, &argc, argv);
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "CLASS_ID", 0,
                "class_id", argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    } else if ((stage == IPP_STAGE_RAW) || (stage == IPP_STAGE_CHIP)) {
        psError(PS_ERR_UNKNOWN, true, "-class_id is required\n");
        usage();
    }

    if ((argnum = psArgumentGet(argc, argv, "-image"))) {
        // for raw and chip level images we use psFits for I/O and ...
        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();
    }

    if ((argnum = psArgumentGet(argc, argv, "-mask"))) {
        psArgumentRemove(argnum, &argc, argv);
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "INPUT.MASK", 0,
                "name of input mask image", argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    }

    if ((argnum = psArgumentGet(argc, argv, "-chip_mask"))) {
        if (stage != IPP_STAGE_CHIP) {
            psError(PS_ERR_UNKNOWN, true, "-chip_mask only allowed at chip stage");
            usage();
        }
        psArgumentRemove(argnum, &argc, argv);
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "INPUT.CHMASK", 0,
                "name of input chip mask image", argv[argnum]);
        psArgumentRemove(argnum, &argc, argv);
    }

    if ((argnum = psArgumentGet(argc, argv, "-weight"))) {
        psArgumentRemove(argnum, &argc, argv);
        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, "-outroot"))) {
        psArgumentRemove(argnum, &argc, argv);
        psString dir = pathToDirectory(argv[argnum]);
        psMetadataAddStr(config->arguments, PS_LIST_TAIL, "OUTPUT", 0, "directory for output files",
            dir);
        psArgumentRemove(argnum, &argc, argv);
    } else {
        psError(PS_ERR_UNKNOWN, true, "-outroot is required\n");
        usage();
    }

    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 exciseAll)
{
    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);
    }
    sf->outImage->header = (psMetadata*) psMemIncrRefCounter(sf->inImage->header);

    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, sf->recImage, sf->inImage, sf->extnum, exciseAll);
    }  else {
        // Image cubes are handeled specially
    }

    // set up the compression parameters
#ifdef STREAKS_COMPRESS_OUTPUT
    // compression of the image pixels is disabled for now. Some consortium members
    // have problems reading them
    copyFitsOptions(sf->outImage, sf->recImage, sf->inImage);

    // XXX: TODO: can we derive these values from the input header?
    // psFitsCompressionGet(sf->inImage->image) gives compression none
    // perhaps we should just use the definition of COMP_IMG in the configuration 
    psFitsSetCompression(sf->outImage->fits, PS_FITS_COMPRESS_RICE, sf->tiles, 8, 0, 0);
    if (sf->recImage) {
        psFitsSetCompression(sf->recImage->fits, PS_FITS_COMPRESS_RICE, sf->tiles, 8, 0, 0);
    }
#endif

    if (sf->inMask) {
        readImage(sf->inMask, sf->extnum, sf->stage);
        sf->outMask->header = (psMetadata*) psMemIncrRefCounter(sf->inMask->header);
        if (sf->recMask) {
            sf->recMask->header = (psMetadata*) psMemIncrRefCounter(sf->inMask->header);
        }
        setupImageRefs(sf->outMask, sf->recMask, sf->inMask, sf->extnum, exciseAll);
        if (sf->outChMask) {
            sf->outChMask->header = (psMetadata *) psMemIncrRefCounter(sf->outMask->header);
            sf->outChMask->image = (psImage *) psMemIncrRefCounter(sf->outMask->image);
            if (sf->recChMask) {
                sf->recChMask->header = (psMetadata *) psMemIncrRefCounter(sf->recMask->header);
                sf->recChMask->image = (psImage *) psMemIncrRefCounter(sf->recMask->image);
            }
        }

#ifdef STREAKS_COMPRESS_OUTPUT
        // XXX: see note above
        copyFitsOptions(sf->outMask, sf->recMask, sf->inMask);
        psFitsSetCompression(sf->outMask->fits, PS_FITS_COMPRESS_PLIO, sf->tiles, 8, 0, 0);
        if (sf->recMask) {
            psFitsSetCompression(sf->recMask->fits, PS_FITS_COMPRESS_PLIO, sf->tiles, 8, 0, 0);
        }
        if (sf->outChMask) {
            copyFitsOptions(sf->outChMask, sf->recChMask, sf->inMask);
            psFitsSetCompression(sf->outChMask->fits, PS_FITS_COMPRESS_PLIO, sf->tiles, 8, 0, 0);
            if (sf->recChMask) {
                psFitsSetCompression(sf->recChMask->fits, PS_FITS_COMPRESS_PLIO, sf->tiles, 8, 0, 0);
            }
        }
#endif
    }

    if (sf->inWeight) {
        readImage(sf->inWeight, sf->extnum, sf->stage);
        sf->outWeight->header = (psMetadata*) psMemIncrRefCounter(sf->inWeight->header);
        if (sf->recWeight) {
            sf->recWeight->header = (psMetadata*) psMemIncrRefCounter(sf->inWeight->header);
        }
        setupImageRefs(sf->outWeight, sf->recWeight, sf->inWeight, sf->extnum, exciseAll);

#ifdef STREAKS_COMPRESS_OUTPUT
        copyFitsOptions(sf->outWeight, sf->recWeight, sf->inWeight);
        // XXX: see note above
        psFitsSetCompression(sf->outWeight->fits, PS_FITS_COMPRESS_RICE, sf->tiles, 8, 0, 0);
        if (sf->recWeight) {
            psFitsSetCompression(sf->recWeight->fits, PS_FITS_COMPRESS_RICE, sf->tiles, 8, 0, 0);
        }
#endif
    }
    // and for raw images, create sub images that represent the actual image
    // area (no overscan)

    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);
        writeImage(sf->recImage, extname, sf->extnum);
    } else {
        // we have an image cube
        double initValue;
        if (exciseImageCube) {
            // copy the entire input image to the recovery image
            writeImageCube(sf->recImage, sf->inImage->imagecube, extname, sf->extnum);
            initValue = NAN;
        } else {
            // otherwise write it to the output 
            writeImageCube(sf->outImage, sf->inImage->imagecube, extname, sf->extnum);
            initValue = 0;
        }

        // 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);
            }
        }
        // Assumption: there are no mask and weight images for video cells
        return;
    }
    if (sf->outMask) {
        writeImage(sf->outMask, extname, sf->extnum);
        writeImage(sf->recMask, extname, sf->extnum);
        if (sf->stage == IPP_STAGE_CHIP) {
            writeImage(sf->outChMask, extname, sf->extnum);
            writeImage(sf->recChMask, extname, sf->extnum);
        }
    }
    if (sf->outWeight) {
        writeImage(sf->outWeight, extname, sf->extnum);
        writeImage(sf->recWeight, extname, sf->extnum);
    }
}

// for any of the outputs that are stored in Nebulous set their extended attributes
// and replicate the files

static bool
replicateOutputs(streakFiles *sfiles)
{
    bool status = false;

    // XXX: TODO: need a nebGetXatrr function, but there isn't one
    // another option would be to take the number of copies to be
    // created as an option. That way the system could decide
    // whether to replicate anything other than raw Image files
    void *xattr = NULL;

    if (!replicate(sfiles->outImage, xattr)) {
        psError(PM_ERR_SYS, false, "failed to replicate outImage.");
        return false;
    }

#ifdef notyet
    // XXX: don't replicate mask and weight images until we can look up
    // the input's xattr. There may be a perl program that can getXattr
    if (sfiles->outMask) {
        // get xattr from input to see if we need to replicate
        if (!replicate(sfiles->outMask, xattr)) {
            psError(PM_ERR_SYS, false, "failed to replicate outImage.");
            return false;
        }
    }
    if (sfiles->outWeight) {
        // get xattr from input to see if we need to replicate
        if (!replicate(sfiles->outWeight, xattr)) {
            psError(PM_ERR_SYS, false, "failed to replicate outImage.");
            return false;
        }
    }
#endif

    return true;
}


