#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);

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

    // Values to set for masked pixels
    psU32 maskStreak = 0;           // for the image and weight (usually NAN, MAXINT for integer images)
    psU32 maskMask = 0;             // value looked up for MASK.STREAK 

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

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

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

        // now that we've read the input files, lookup the mask values that we read
        if (maskStreak == 0) {
            strkGetMaskValues(sfiles, &maskStreak, &maskMask);
        }

        setMaskedToNAN(sfiles, maskMask, true);

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

    // check the weight and mask files for extra extensions that might be in files
    // (covariance matrix for example)
    copyExtraExtensions(sfiles);



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

    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, false);
        
        // astrom struct is only needed for raw stage, and only the chip to cell parameters are used
        sf->astrom = streakSetAstrometry(sf->astrom, sf->stage, NULL, NULL, false,
            sf->inImage->header, sf->inImage->numCols, sf->inImage->numRows);
    }
    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 if (sf->inImage->imagecube) {
        // Image cubes should have been excised in the destreaking process
        streaksExit("unexpected imagecube found", PS_EXIT_CONFIG_ERROR);
    }  else {
        return false;
    }

    // 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);
        if (sf->outMask) {
            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);
                }
            }

            copyFitsOptions(sf->outMask, sf->recMask, sf->inMask, sf->tiles);
            if (sf->outChMask) {
                copyFitsOptions(sf->outChMask, sf->recChMask, sf->inMask, sf->tiles);
            }
        }
    }

    if (sf->inWeight) {
        readImage(sf->inWeight, sf->extnum, sf->stage, false);
        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);

        copyFitsOptions(sf->outWeight, sf->recWeight, sf->inWeight, sf->tiles);
    }
    // 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);
    }
}
