#include "streaksremove.h"
#include "libgen.h"
#include "unistd.h"

static void sFileFree(sFile *sfile);
static void streakFilesFree(streakFiles *sf);
static void freeImages(streakFiles *);
static void setExciseValue(sFile *in);

static nebServer *ourNebServer = NULL;

// Assumptions about the file structure of non-raw files
// The 'image' for each file (image, mask weight) is contained in the first
// extension.


// open the files required for streaks procesing.
// if remove is true the calling program is streaksremove and the recovery files are outputs
// if false the recovery files are inputs
streakFiles *openFiles(pmConfig *config, bool remove, char *program_name)
{
    bool status;
    streakFiles *sf = psAlloc(sizeof(streakFiles));
    psMemSetDeallocator(sf, (psFreeFunc) streakFilesFree);
    memset(sf, 0, sizeof(*sf));

    sf->config = config;
    sf->program_name = basename(program_name);

    if (remove) {
        // remember pointer so that streaksExit can delete temps
        setStreakFiles(sf);
    }


    // error checking is done by sFileOpen. If a file can't be opened we just exit
    ippStage stage = psMetadataLookupS32(&status, config->arguments, "STAGE");

    sf->stage = stage;
    sf->extnum = 0;

    sf->class_id = psMetadataLookupStr(&status, config->arguments, "CLASS_ID");

    sf->inImage = sFileOpen(config, stage, "INPUT", NULL, true, false);
    sf->nHDU = sf->inImage->nHDU;

    // don't need to free inputBasename see basename(3)
    // The names of the temporary and recovery files are taken from the input
    char *inputBasename = basename(sf->inImage->name);
    sf->outImage = sFileOpen(config, stage, "OUTPUT", inputBasename, true, true);

    if (remove) {
        // XXX: the recovery file should be required if stage is raw and the replace flag is set
        // that is if the input raw image is to be destroyed
        sf->recImage = sFileOpen(config, stage, "RECOVERY", inputBasename, false, false);
    } else {
        sf->recImage = sFileOpen(config, stage, "RECOVERY.IMAGE", NULL, true, false);
    }

    sf->inMask = sFileOpen(config, stage, "INPUT.MASK", NULL, false, false);
    if (sf->inMask && (sf->stage != IPP_STAGE_RAW)) {
        inputBasename = basename(sf->inMask->name);
        sf->outMask = sFileOpen(config, stage, "OUTPUT", inputBasename, true, true);
        if (remove) {
            sf->recMask = sFileOpen(config, stage, "RECOVERY", inputBasename, false, false);
        } else {
            sf->recMask = sFileOpen(config, stage, "RECOVERY.MASK", NULL, true, false);
        }
    }
    // If we're in streaksremove and operating on a chip stage image
    // if we are passed a mask it is camera stage mask and we also
    // need to destreak the chip level mask file as well
    // If it doesn't exist, we didn't have a camera mask
    if (remove && sf->inMask && (stage == IPP_STAGE_CHIP)) {
        sf->inChMask = sFileOpen(config, stage, "INPUT.CHMASK", NULL, false, false);
        if (sf->inChMask) {
            inputBasename = basename(sf->inChMask->name);
            sf->outChMask = sFileOpen(config, stage, "OUTPUT", inputBasename, true, true);
            sf->recChMask = sFileOpen(config, stage, "RECOVERY", inputBasename, false, false);
        }
    }

    sf->inWeight = sFileOpen(config, stage, "INPUT.WEIGHT", NULL, false, false);
    if (sf->inWeight) {
        inputBasename = basename(sf->inWeight->name);
        sf->outWeight = sFileOpen(config, stage, "OUTPUT", inputBasename, true, true);
        if (remove) {
            sf->recWeight = sFileOpen(config, stage, "RECOVERY", inputBasename, false, false);
        } else {
            sf->recWeight = sFileOpen(config, stage, "RECOVERY.WEIGHT", NULL, true, false);
        }
    }
    if (remove) {
        sf->inSources = sFileOpen(config, stage, "SOURCES", NULL, false, false);
        if (sf->inSources) {
            inputBasename = basename(sf->inSources->name);
            sf->outSources = sFileOpen(config, stage, "OUTPUT", inputBasename, true, true);
        }
    }

    psElemType tileType;                // Type corresponding to "long"
    if (sizeof(long) == sizeof(psS64)) {
        tileType = PS_TYPE_S64;
    } else if (sizeof(long) == sizeof(psS32)) {
        tileType = PS_TYPE_S32;
    } else {
        psAbort("can't map (long) type to a psLib type");
    }

    sf->tiles = psVectorAlloc(3, tileType); // Tile sizes
    if (tileType == PS_TYPE_S64) {
        sf->tiles->data.S64[0] = 0;
        sf->tiles->data.S64[1] = 1;
        sf->tiles->data.S64[2] = 1;
    } else {
        sf->tiles->data.S32[0] = 0;
        sf->tiles->data.S32[1] = 1;
        sf->tiles->data.S32[2] = 1;
    }

    sf->transparentStreaks = psMetadataLookupF64(&status, config->arguments, "TRANSPARENT_STREAKS");

    return sf;
}

static void
streakFilesFree(streakFiles *sf)
{
    freeImages(sf);
    psFree(sf->stats);
    psFree(sf->diffedPixels);
    psFree(sf->tiles);
    psFree(sf->view);
    psFree(sf->astrom);
}


// getNebServer()
//
// it gets created the first time a nebulous file is accessed.
// if config is passed in we are to create it and exit if not able.
// if config is null we return NULL indicating that there are no nebulous files in use
// The purpose of this is to avoid passing extra arguments between all of the functions.
// but it is probably a bad idea.
//
static nebServer *getNebServer(pmConfig *config)
{
    if (ourNebServer) {
        return ourNebServer;
    }

    if (!config) {
        return NULL;
    }

    psString serverURI = getenv("NEB_SERVER");

    // XXX: Note: all of the flags to psError should be true, but I don't think nebclient
    // uses psError
    if (!serverURI) {
        bool status;
        serverURI = psMetadataLookupStr(&status, config->site, "NEB_SERVER");
        if (!status) {
            psError(PM_ERR_CONFIG, true, "failed to lookup config value for NEB_SERVER.");
            streaksExit("", PS_EXIT_CONFIG_ERROR);
        }
    }

    if (!serverURI) {
        psError(PM_ERR_CONFIG, true, "Could not determine nebulous server URI.");

        streaksExit("", PS_EXIT_CONFIG_ERROR);
    }

    ourNebServer = nebServerAlloc(serverURI);
    if (!ourNebServer) {
        psError(PM_ERR_SYS, true, "failed to create a nebServer object from %s.", serverURI);
        streaksExit("", PS_EXIT_CONFIG_ERROR);
    }

    return ourNebServer;
}

// figure out if a nebulous instance is a non-destreaked file
static bool
nebFileIsDestreaked(sFile *sfile)
{
    if (!sfile->resolved_name) {
        psError(PS_ERR_PROGRAMMING, true, "resolved name is null");
        return false;
    }
    psFits *fits = psFitsOpen(sfile->resolved_name, "r");
    if (!fits) {
        psError(PS_ERR_IO, true, "failed open %s", sfile->name);
        // can't tell if it is a destreaked file
        return false;
    }
    psMetadata *header = psFitsReadHeader(NULL, fits);
    if (!header) {
        psError(PS_ERR_IO, true, "failed to read header for: %s", sfile->name);
        return false;
    }
    bool mdok;
    bool isDestreaked = psMetadataLookupBool(&mdok, header, "PSDESTRK");
    if (mdok && isDestreaked) {
        return true;
    } else {
        psError(PS_ERR_IO, false, "output file already exists and may not be de-streaked: %s", sfile->name);
        return false;
    }
}

static psString
resolveFilename(pmConfig *config, sFile *sfile, bool create, bool checkIfDestreaked)
{
    sfile->inNebulous = IN_NEBULOUS(sfile->name);

    if (sfile->inNebulous) {
        // make sure we have our neb server connection
        nebServer *server = getNebServer(config);
        if (create) {
            // delete the existing file, since there may be more than one
            // instance. It will get created below in pmConfigConvertFilename
            if ((sfile->resolved_name = nebFind(server, sfile->name)) != NULL) {
                if (checkIfDestreaked && !nebFileIsDestreaked(sfile)) {
                    psError(PS_ERR_IO, false, "attempting to delete file that has not been destreaked %s", sfile->name);
                    return NULL;
                }
                nebFree(sfile->resolved_name);
                sfile->resolved_name = NULL;
                nebDelete(server, sfile->name);
            }
        }
    }

    return pmConfigConvertFilename(sfile->name, config, create, create);
}

sFile *sFileOpen(pmConfig *config, ippStage stage, psString fileSelect,
    psString outputFilename, bool required, bool checkIfDestreaked)
{
    bool status;
    sFile *sfile = psAlloc(sizeof(sFile));
    memset(sfile, 0, sizeof(sFile));

    // We use psFits directly to read the image unless the stage is warp
    // or diff. In those cases we use the pmFPAfile functions to read the image file
    // to make managing the astrometry easy.
    // The reason we don't use pmFPAfile in all cases is that I was having trouble getting
    // all of the keywords in the raw image files written to the output destreaked files

    if (!outputFilename && !CHIP_LEVEL_INPUT(stage) && !strcmp(fileSelect, "INPUT")) {
        // stage is warp or diff AND fileSelect eq "INPUT"
        // get data from pmFPAfile.

        // we need to know what the nebulous and real filenames are so we steal
        // some code from pmFPAfileDefineFromArgs
        // XXX: create a pmFPAfile function that does this and use it there
        psArray *infiles = psMetadataLookupPtr(&status, config->arguments, "INPUT");
        if (!status) {
            psError(PS_ERR_PROGRAMMING, false, "INPUT not found in config->arguments");
            streaksExit("", PS_EXIT_PROG_ERROR);
        }
        if (infiles->n != 1) {
            psError(PS_ERR_IO, false, "Found n == %ld files in %s in arguments expencted 1\n",
                infiles->n, "INPUT");
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        // end of file name lookup code adapted from pmFPAfileDefineFromArgs
        //
        sfile->name = psStringCopy(infiles->data[0]);
        sfile->inNebulous = IN_NEBULOUS(sfile->name);

        // XXX: I should probably be using a different file rule for diffs, but I don't
        // have an input rule that takes a diff image.
        // Since they're skycells, this should be compatible
        sfile->pmfile = pmFPAfileDefineFromArgs(&status, config, "PPSUB.INPUT", "INPUT");
        if (!sfile->pmfile) {
            psError(PS_ERR_IO, false, "Failed to define file for %s", fileSelect);
            streaksExit("", PS_EXIT_DATA_ERROR);
        }

        pmFPAview *view = pmFPAviewAlloc(0);
        if (!pmFPAfileIOChecks(config, view, PM_FPA_BEFORE)) {
            psError(PS_ERR_UNKNOWN, false, "Failed to load input.");
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        psFree(view);
        sfile->resolved_name = psStringCopy(sfile->pmfile->filename);

        sfile->nHDU = psFitsGetSize(sfile->pmfile->fits);

        // copy header from fpu
        sfile->header = (psMetadata*) psMemIncrRefCounter(sfile->pmfile->fpa->hdu->header);

        return sfile;
    }

    // For all other files we use using psFits for i/o

    psString name = psMetadataLookupStr(&status, config->arguments, fileSelect);
    if (!status || !name) {
        sFileFree(sfile);
        if (required) {
            psError(PS_ERR_IO, false, "Failed to lookup name for %s", fileSelect);
            streaksExit("", 1);
        }
        return NULL;
    }

    // if outputFilename is not null name it contains the "directory" (perhaps with a prefix like SR_)
    // and outputFilename is the basename name of the file (or nebulous key)
    // and the file is to be opened for writing
    if (outputFilename) {
        sfile->write = true;
        psStringAppend(&sfile->name, "%s%s", name, outputFilename);
        sfile->resolved_name = resolveFilename(config, sfile, true, checkIfDestreaked);
        if (!sfile->resolved_name) {
            psError(PS_ERR_IO, false, "Failed to resolve filename for %s", sfile->name);
            sFileFree(sfile);
            streaksExit("", 1);
        }
        sfile->fits = psFitsOpen(sfile->resolved_name, "w");
        if (sfile->fits) {
            sfile->fits->options = psFitsOptionsAlloc();
        }
    } else {
        sfile->name = psStringCopy(name);
        sfile->resolved_name = resolveFilename(config, sfile, false, false);
        if (!sfile->resolved_name) {
            psError(PS_ERR_IO, false, "Failed to resolve name for %s", sfile->name);
            sFileFree(sfile);
            streaksExit("", 1);
        }
        sfile->fits = psFitsOpen(sfile->resolved_name, "r");
        if (sfile->fits) {
            sfile->nHDU = psFitsGetSize(sfile->fits);
        }
    }

    if (!sfile->fits) {
        psError(PS_ERR_IO, false, "failed to open fits file %s for %s",
                    sfile->resolved_name, outputFilename ? "writing" : "reading");
        sFileFree(sfile);
        streaksExit("", 1);
    }

    return sfile;
}


static bool
moveExt(sFile *sfile, int extnum)
{
    bool status = psFitsMoveExtNum(sfile->fits, extnum, false);
    if (!status) {
        psError(PS_ERR_IO, false,
            "failed to move to extension %d for %s", extnum, sfile->resolved_name);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    return true;
}

bool
streakFilesNextExtension(streakFiles *sf)
{
    // unless stage is raw when we get here we're done
    if (sf->stage != IPP_STAGE_RAW) {
        if (sf->view) {
            sf->view->readout = -1;
            pmFPAfileIOChecks(sf->config, sf->view, PM_FPA_AFTER);
            sf->view->cell = -1;
            pmFPAfileIOChecks(sf->config, sf->view, PM_FPA_AFTER);
            sf->view->chip = -1;
            pmFPAfileIOChecks(sf->config, sf->view, PM_FPA_AFTER);
        }
        return false;
    }

    ++sf->extnum;
    if (sf->nHDU == 1) {
        // return true the first time through
        return (sf->extnum == 1);
    }
    if (sf->extnum < sf->nHDU) {
        moveExt(sf->inImage, sf->extnum);
        // No mask and weight images for raw stage
        return true;
    } else {
        // we're all done
        return false;
    }
}

void
addDestreakKeyword(psMetadata *header, bool value)
{
    psMetadataAddBool(header, PS_LIST_TAIL, "PSDESTRK", PS_META_REPLACE,
        "Have streaks been removed from image?", value);
}

void
addRecoveryKeyword(psMetadata *header)
{
    psMetadataAddBool(header, PS_LIST_TAIL, "PSRECOVR", PS_META_REPLACE,
        "Does this image contain excised streak pixels?", true);
}

void
copyPHU(streakFiles *sfiles, bool remove, bool destreak)
{
    psAssert(sfiles->stage == IPP_STAGE_RAW, "copyPHU should only be used for raw stage");

    psMetadata *imageHeader = psFitsReadHeader(NULL, sfiles->inImage->fits);
    if (!imageHeader) {
        psError(PS_ERR_IO, false, "failed to read primary header from %s\n",
            sfiles->inImage->resolved_name);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    psMetadata *recHeader = NULL;
    if (remove && sfiles->recImage) {
       recHeader = psMetadataCopy(NULL, imageHeader);
       addRecoveryKeyword(recHeader);
    }

    // add keyword indicating that streaks have been removed
    addDestreakKeyword(imageHeader, destreak);

    if (!psFitsWriteBlank(sfiles->outImage->fits, imageHeader, NULL)) {
        psError(PS_ERR_IO, false, "failed to write primary header to %s",
            sfiles->outImage->resolved_name);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    if (recHeader && !psFitsWriteBlank(sfiles->recImage->fits, recHeader, NULL)) {
        psError(PS_ERR_IO, false, "failed to write primary header to %s",
            sfiles->recImage->resolved_name);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    psFree(recHeader);
    recHeader = NULL;
    psFree(imageHeader);

    sFile *inMask = sfiles->inMask;
    if (inMask && sfiles->outMask) {
        psMetadata *maskHeader = psFitsReadHeader(NULL, inMask->fits);
        if (!maskHeader) {
            psError(PS_ERR_IO, false, "failed to read primary header from %s\n",
                sfiles->inMask->resolved_name);
            streaksExit("", 1);
        }
        if (remove && sfiles->recMask) {
            recHeader = psMetadataCopy(NULL, maskHeader);
            // add keyword indicating that this is the recovery image
           addRecoveryKeyword(recHeader);
        }
        // add keyword indicating that streaks have been removed
        addDestreakKeyword(maskHeader, destreak);
        if (!psFitsWriteBlank(sfiles->outMask->fits, maskHeader, NULL)) {
            psError(PS_ERR_IO, false, "failed to write primary header to %s",
                sfiles->outMask->resolved_name);
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        if (recHeader && !psFitsWriteBlank(sfiles->recMask->fits, recHeader, NULL)) {
            psError(PS_ERR_IO, false, "failed to write primary header to %s",
                sfiles->recMask->resolved_name);
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        psFree(recHeader);
        recHeader = NULL;
        psFree(maskHeader);
    }
    sFile *inWeight = sfiles->inWeight;
    if (inWeight) {
        psMetadata *weightHeader = psFitsReadHeader(NULL, inWeight->fits);
        if (!weightHeader) {
            psError(PS_ERR_IO, false, "failed to read primary header from %s\n",
                sfiles->inWeight->resolved_name);
            streaksExit("", 1);
        }
        if (remove && sfiles->recWeight) {
            recHeader = psMetadataCopy(NULL, weightHeader);
            // add keyword indicating that this is a recovery image
           addRecoveryKeyword(recHeader);
        }

        // add keyword indicating that streaks have been removed
        addDestreakKeyword(weightHeader, destreak);
        if (!psFitsWriteBlank(sfiles->outWeight->fits, weightHeader, NULL)) {
            psError(PS_ERR_IO, false, "failed to write primary header to %s",
                sfiles->outWeight->resolved_name);
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        if (recHeader && !psFitsWriteBlank(sfiles->recWeight->fits, recHeader, NULL)) {
            psError(PS_ERR_IO, false, "failed to write primary header to %s",
                sfiles->recWeight->resolved_name);
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        psFree(weightHeader);
        psFree(recHeader);
    }
}

// Determine whether the current header for this file is an image or not
// Find a cleaner way to do this
static bool
isImage(sFile *in)
{
    bool status;
    psString exttype = psMetadataLookupStr(&status, in->header, "EXTTYPE");
    if (exttype && !strcmp(exttype, "IMAGE")) {
        return true;
    }
    // examine current header and determine if it is an image
    psString xtension = psMetadataLookupStr(&status, in->header, "XTENSION");
    if (xtension) {
        if (!strcmp(xtension, "IMAGE")) {
            return true;
        } else if (!strcmp(xtension, "BINTABLE")) {
            psString ztension = psMetadataLookupStr(&status, in->header, "ZTENSION");
            if (ztension && !strcmp(ztension, "IMAGE")) {
                return true;
            }
        }
    } else if (psMetadataLookupBool(&status, in->header, "ZIMAGE")) {
        return true;
    } else if (in->nHDU == 1) {
        // no extensions in the file, can just return true? For now require
        // at least one dimension
        int naxis =  psMetadataLookupS32(&status, in->header, "NAXIS");
        return naxis > 0;
    }
    return false;
}

void
readImageFrom_pmFile(streakFiles *sf)
{
    // XXX: This function assumes that it is only used for a single
    // chip single cell FPA (i.e. a skycell)
    assert(sf->inImage->pmfile->fpa->chips->n == 1);

    pmFPAview *view = sf->view;
    if (view) {
        // The view was allocated and the chip set by setupAstrometry()
        assert(view->chip == 0);
    } else {
        sf->view = pmFPAviewAlloc(0);
        view = sf->view;
        sf->chip = pmFPAviewNextChip(view, sf->inImage->pmfile->fpa, 1);
        assert(sf->chip != NULL);
    }
    view->cell = 0;
    sf->cell =  pmFPAviewThisCell(view, sf->inImage->pmfile->fpa);

    if (!sf->cell) {
        streaksExit("no cells found in chip", PS_EXIT_PROG_ERROR);
    }
    if (sf->cell->readouts->n != 1) {
        psError(PS_ERR_PROGRAMMING, true, "unexpected number of readouts found: %ld", sf->cell->readouts->n);
        streaksExit("", PS_EXIT_PROG_ERROR);
    }
    view->readout = 0;
    pmReadout *readout = pmFPAviewThisReadout(view, sf->inImage->pmfile->fpa);
    if (!readout) {
        streaksExit("readout 0 not found", PS_EXIT_PROG_ERROR);
    }

    sf->inImage->image = (psImage*) psMemIncrRefCounter(readout->image);
    sf->inImage->numCols = readout->image->numCols;
    sf->inImage->numRows = readout->image->numRows;
    setExciseValue(sf->inImage);
}

static void
setDataExtent(ippStage stage, sFile *in, bool rawImage)
{
    if (rawImage) {
        psString datasec = psMetadataLookupStr(NULL, in->header, "DATASEC");
        if (!datasec) {
            psError(PS_ERR_IO, false, "failed to find DATASEC in header");
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        int xmin, xmax, ymin, ymax;
        sscanf(datasec, "[%d:%d,%d:%d]", &xmin, &xmax, &ymin, &ymax);

        // I've inadvertantly introduced an implicit assumption that cell coordinates
        // are the same as image coordinates. This is true for GPC1
        // Test for this condition not being true
        if (xmin != 1) {
            psError(PS_ERR_IO, true, "can't handle datasec with minimum x != 1: %s", datasec);
            streaksExit("", PS_EXIT_PROG_ERROR);
        }
        if (ymin != 1) {
            psError(PS_ERR_IO, true, "can't handle datasec with minimum y != 1: %s", datasec);
            streaksExit("", PS_EXIT_PROG_ERROR);
        }

        // These values are different than the size of the actual image which includes
        // the bias area.
        in->numCols = xmax - xmin + 1;
        in->numRows = ymax - ymin + 1;
    } else {
        in->numCols = in->image->numCols;
        in->numRows = in->image->numRows;
    }
}

void
readImage(sFile *in, int extnum, ippStage stage, bool isMask)
{
    psRegion region = {0, 0, 0, 0};

    if (in->header) psFree(in->header);
    if (in->image) {
        psFree(in->image);
        in->image = NULL;
    }
    if (in->imagecube) {
        psFree(in->imagecube);
        in->imagecube = NULL;
    }

    in->header = psFitsReadHeader(NULL, in->fits);
    if (!in->header) {
        psError(PS_ERR_IO, false, "failed to read header from %s extnum: %d",
            in->resolved_name, extnum);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }

    bool status;
    in->extname = psMetadataLookupStr(&status, in->header, "EXTNAME");

    if (!isImage(in)) {
        psLogMsg("streaksremove", PS_LOG_INFO, "%s contains no image\n", in->name);
        return;
    }

    in->numZPlanes = psMetadataLookupS32(NULL, in->header, "NAXIS3");

    if (in->numZPlanes == 0) {
        in->image = psFitsReadImage(in->fits, region, 0);
        if (!in->image) {
            psError(PS_ERR_IO, false, "failed to read image from %s extnum: %d",
                in->resolved_name, extnum);
            streaksExit("", PS_EXIT_DATA_ERROR);
        }
        setExciseValue(in);
    }  else {
        if (stage != IPP_STAGE_RAW) {
            psError(PS_ERR_PROGRAMMING, true, "unexpected image cube found in non-raw image");
            streaksExit("", PS_EXIT_PROG_ERROR);
        }
        in->imagecube = psFitsReadImageCube(in->fits, region);
        if (!in->imagecube) {
            psError(PS_ERR_IO, false, "failed to read image cube from %s extnum: %d",
                in->resolved_name, extnum);
            streaksExit("", PS_EXIT_DATA_ERROR);
        }

        // Ensure input is of the expected type
        psDataType expected = isMask ? PS_TYPE_IMAGE_MASK : PS_TYPE_F32; // Expected type for image
        for (int i = 0; i < in->imagecube->n; i++) {
            psImage *image = in->imagecube->data[i]; // Image of interest
            if (image->type.type != expected) {
                psImage *temp = psImageCopy(NULL, image, expected);
                psFree(image);
                in->imagecube->data[i] = temp;
            }
        }
    }
    setDataExtent(stage, in, (stage == IPP_STAGE_RAW) && !isMask);
}

static void
setFitsOptions(sFile *sfile, int bitpix, float bscale, float bzero, psFitsCompressionType compType,
    psVector *tiles)
{
    if (!sfile || !sfile->write) {
        return;
    }
    if (sfile->fits->options) {
        psFree(sfile->fits->options);
    }
    sfile->fits->options = psFitsOptionsAlloc();
    sfile->fits->options->scaling = PS_FITS_SCALE_MANUAL;
    sfile->fits->options->fuzz = false;
    sfile->fits->options->bitpix = bitpix;
    sfile->fits->options->bscale = bscale;
    sfile->fits->options->bzero = bzero;

    psFitsSetCompression(sfile->fits, compType, tiles, 8, 0, 0);
}

void
copyFitsOptions(sFile *out, sFile *rec, sFile *in, psVector *tiles)
{
    bool mdok;
    psString compTypeStr = psMetadataLookupStr(&mdok, in->header, "ZCMPTYPE");
    psFitsCompressionType compType = psFitsCompressionTypeFromString(compTypeStr);
    if (compType == PS_FITS_COMPRESS_NONE) {
        return;
    }
    // Get current BITPIX, BSCALE, BZERO, EXTNAME
    // Probably not necessary to look the numerical values up in this
    // way, but guards against changes to psLib and cfitsio FITS
    // handling.
    psMetadataItem *bitpixItem = psMetadataLookup(in->header, "BITPIX");
    psAssert(bitpixItem, "Every FITS image should have BITPIX");
    int bitpix = psMetadataItemParseS32(bitpixItem);
    psMetadataItem *bscaleItem = psMetadataLookup(in->header, "BSCALE");

    float bscale;
    if (!bscaleItem) {
        psWarning("BSCALE isn't set; defaulting to unity");
        bscale = 1.0;
    } else {
        bscale = psMetadataItemParseF32(bscaleItem);
    }
    psMetadataItem *bzeroItem = psMetadataLookup(in->header, "BZERO");
    float bzero;
    if (!bzeroItem) {
        psWarning("BZERO isn't set; defaulting to zero");
        bzero = 0.0;
    } else {
        bzero = psMetadataItemParseF32(bzeroItem);
    }

#ifdef STREAKS_COMPRESS_OUTPUT
    // printf("%d %f %f\n", bitpix, bscale, bzero);
    setFitsOptions(out, bitpix, bscale, bzero, compType, tiles);
    setFitsOptions(rec, bitpix, bscale, bzero, compType, tiles);
#endif
}


void
copyTable(sFile *out, sFile *in, int extnum)
{
    bool mdok;
    psString xtension = psMetadataLookupStr(&mdok, in->header, "XTENSION");
    psString extname = psMetadataLookupStr(&mdok, in->header, "EXTNAME");

    if (!xtension) {
        psWarning("extnum %d has no image and xtension is not defined in %s",
            extnum, in->resolved_name);
        // XXX abort?
        return;
    }
    if (!extname) {
        psWarning("extnum %d has no image and extname not defined in %s",
            extnum, in->resolved_name);
        // XXX abort?
        return;
    }
    psFitsTable *table = psFitsReadTableNew(in->fits);
    if (!table) {
        psError(PS_ERR_UNKNOWN, false, "failed to read table in extension %d from in->resolved name", extnum);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }

    if (!psFitsWriteTableNew(out->fits, out->header, table, extname)) {
        psError(PS_ERR_UNKNOWN, false, "failed to copy table in extension %d", extnum);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    psFree(table);
}

static void
createNewImage(sFile *out, sFile *in, int extnum, double initValue)
{
    out->image = psImageRecycle(out->image, in->image->numCols, in->image->numRows, in->image->type.type);
    if (!out->image) {
        psError(PS_ERR_UNKNOWN, false, "failed to allocate image for extnsion %d", extnum);
        streaksExit("", PS_EXIT_UNKNOWN_ERROR);
    }
    psImageInit(out->image, initValue);
}

void
setupImageRefs(sFile *out, sFile *recoveryOut, sFile *in, int extnum, bool exciseAll)
{
    if (!in->image) {
        // Input file is likely corrupt. XXX: Shouldn't we have caught this earlier?
        psError(PS_ERR_UNKNOWN, true, "image is null for extension %d of %s file is likely corrupt", extnum, in->name);
        streaksExit("", PS_EXIT_UNKNOWN_ERROR);
    }
    if (!exciseAll) {
        // set output image to be the input image
        out->image = (psImage*) psMemIncrRefCounter(in->image);
        if (recoveryOut) {
            // create a recovery image initialized to NAN
            createNewImage(recoveryOut , in, extnum, NAN);
        }
    } else {
        // Excising all pixels in the input
        // set the output to an image all NANs
        createNewImage(out, in, extnum, NAN);
        if (recoveryOut) {
            // save the whole input as the recovery image
            recoveryOut->image = (psImage*) psMemIncrRefCounter(in->image);
        }
    }
}

void
writeImage(sFile *sfile, psString extname, int extnum)
{
    if (!sfile || !sfile->image) {
        return;
    }
    streaksVersionHeaderFull(sfile->header);

    if (!psFitsWriteImage(sfile->fits, sfile->header, sfile->image, 0, extname)) {
        psError(PS_ERR_IO, false, "failed to write image to %s extnum: %d",
            sfile->resolved_name, extnum);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    psFree(sfile->image);
    sfile->image = NULL;
    psFree(sfile->header);
    sfile->header = NULL;
}

void
writeImageCube(sFile *sfile, psArray *imagecube, psString extname, int extnum)
{
    if (!sfile || !imagecube) {
        return;
    }
    streaksVersionHeaderFull(sfile->header);
    if (!psFitsWriteImageCube(sfile->fits, sfile->header, imagecube, extname)) {
        psError(PS_ERR_IO, false, "failed to write image to %s extnum: %d",
            sfile->resolved_name, extnum);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    psFree(sfile->header);
    sfile->header = NULL;
}

void
sFileFree(sFile *sfile)
{
    if (!sfile) {
        return;
    }
    psFree(sfile->name);
    psFree(sfile->resolved_name);
    psFree(sfile->fits);
    psFree(sfile);
}

void
closeImage(sFile *sfile)
{
    if (!sfile) {
        return;
    }
    if (sfile->fits && !psFitsClose(sfile->fits)) {
        psError(PS_ERR_IO, false, "failed to close image to %s",
            sfile->resolved_name);
        streaksExit("", PS_EXIT_DATA_ERROR);
    }
    sfile->fits = NULL;
    psFree(sfile->header);
    sfile->header = NULL;
    psFree(sfile->image);
    sfile->image = NULL;
    psFree(sfile->imagecube);
    sfile->imagecube = NULL;
}

void
closeImages(streakFiles *sf)
{
    if (sf->stage == IPP_STAGE_RAW) {
        if (sf->view) {
            sf->view->readout = -1;
            pmFPAfileIOChecks(sf->config, sf->view, PM_FPA_AFTER);
            sf->view->cell = -1;
            pmFPAfileIOChecks(sf->config, sf->view, PM_FPA_AFTER);
            sf->view->chip = -1;
            pmFPAfileIOChecks(sf->config, sf->view, PM_FPA_AFTER);
        }
    }
    closeImage(sf->inImage);
    closeImage(sf->outImage);
    closeImage(sf->recImage);
    if (sf->inMask) {
        closeImage(sf->inMask);
        closeImage(sf->outMask);
        closeImage(sf->recMask);
    }
    if (sf->inChMask) {
        closeImage(sf->inChMask);
        closeImage(sf->outChMask);
        closeImage(sf->recChMask);
    }
    if (sf->inWeight) {
        closeImage(sf->inWeight);
        closeImage(sf->outWeight);
        closeImage(sf->recWeight);
    }
}

static void
freeImages(streakFiles *sf)
{
    sFileFree(sf->inImage);
    sFileFree(sf->outImage);
    sFileFree(sf->recImage);
    if (sf->inMask) {
        sFileFree(sf->inMask);
        sFileFree(sf->outMask);
        sFileFree(sf->recMask);
    }
    if (sf->inChMask) {
        sFileFree(sf->inChMask);
        sFileFree(sf->outChMask);
        sFileFree(sf->recChMask);
    }
    if (sf->inWeight) {
        sFileFree(sf->inWeight);
        sFileFree(sf->outWeight);
        sFileFree(sf->recWeight);
    }
    if (sf->inSources) {
        sFileFree(sf->inSources);
        sFileFree(sf->outSources);
    }
}

bool
replicate(sFile *outFile, sFile *inFile)
{
    if (!outFile->inNebulous) {
        return true;
    }
    nebServer *server = getNebServer(NULL);

    char *user_copies = nebGetXattr(server, inFile->name, "user.copies");
    bool free_user_copies = true;
    if (user_copies == NULL) {
        // input does not have replication requested
        return true;
    }
    if (!nebSetXattr(server, outFile->name, "user.copies", user_copies,  NEB_REPLACE)) {
        psError(PM_ERR_UNKNOWN, true, "nebSetXattr failed for %s\n%s", outFile->name, nebErr(server));
        return false;
    }
    if (!nebReplicate(server, outFile->name, NULL, NULL)) {
        psError(PM_ERR_UNKNOWN, true, "nebReplicate failed for %s\n%s", outFile->name, nebErr(server));
        return false;
    }
    if (free_user_copies) {
        nebFree(user_copies);
    }
    return true;
}

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

bool
replicateOutputs(streakFiles *sfiles)
{
    if (!replicate(sfiles->outImage, sfiles->inImage)) {
        psError(PM_ERR_SYS, false, "failed to replicate outImage.");
        return false;
    }

    if (sfiles->outMask) {
        if (!replicate(sfiles->outMask, sfiles->inMask)) {
            psError(PM_ERR_SYS, false, "failed to replicate outImage.");
            return false;
        }
    }
    if (sfiles->outChMask) {
        if (!replicate(sfiles->outChMask, sfiles->inChMask)) {
            psError(PM_ERR_SYS, false, "failed to replicate outImage.");
            return false;
        }
    }
    if (sfiles->outWeight) {
        if (!replicate(sfiles->outWeight, sfiles->inWeight)) {
            psError(PM_ERR_SYS, false, "failed to replicate outImage.");
            return false;
        }
    }

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

//      XXX: replicate the recovery images (if in nebulous)
//      perhaps whether we do that or not should be configurable.
//      Sounds like we need a recipe

    return true;
}

static bool
swapOutputToInput(sFile *in, sFile *out)
{

    if (in->inNebulous) {
        nebServer *server = getNebServer(NULL);

        if (!nebSwap(server, in->name, out->name)) {
            psError(PM_ERR_SYS, true, "failed to swap files for %s: %s.",
                in->name, nebErr(server));
            return false;
        }
    } else {
        psString command = NULL;

        // regular files. use mv creating a backup
        // NOTE: -b is a gnu option and may not always be available
        psStringAppend(&command, "mv -b --suffix=.bak %s %s", out->resolved_name, in->resolved_name);
        int status = system(command);
        if (status != 0) {
            psError(PM_ERR_SYS, true, "%s failed with status %d.",
                command, status);
            psFree(command);
            return false;
        }
        psFree(command);
        // XXX: shouldn't I be doing mv in->resolved_name.bak out->resolved_name
        // to complete the swap if !remove?
    }

    return true;
}

bool
swapOutputsToInputs(streakFiles *sfiles)
{
    if (sfiles->outMask) {
        if (!swapOutputToInput(sfiles->inMask, sfiles->outMask)) {
            psError(PM_ERR_SYS, false, "failed to swap instances for Mask.");
            return false;
        }
    }

    if (sfiles->outWeight) {
        if (!swapOutputToInput(sfiles->inWeight, sfiles->outWeight)) {
            psError(PM_ERR_SYS, false, "failed to swap instances for Weight.");
            return false;
        }
    }

    if (sfiles->outChMask) {
        if (!swapOutputToInput(sfiles->inChMask, sfiles->outChMask)) {
            psError(PM_ERR_SYS, false, "failed to swap instances for chip mask.");
            return false;
        }
    }

    if (sfiles->outSources) {
        if (!swapOutputToInput(sfiles->inSources, sfiles->outSources)) {
            psError(PM_ERR_SYS, false, "failed to swap instances for sources.");
            return false;
        }
    }

    if (!swapOutputToInput(sfiles->inImage, sfiles->outImage)) {
        psError(PM_ERR_SYS, false, "failed to swap instances for Image.");
        return false;
    }


    return true;
}

static bool
deleteFile(sFile *sfile)
{
    if (sfile->inNebulous) {
        nebServer *server = getNebServer(NULL);
        if (!nebDelete(server, sfile->name)) {
            psError(PM_ERR_SYS, false, "failed to delete %s\n%s.", sfile->name,
                nebErr(server));
            return false;
        }
    } else {
        if (unlink(sfile->resolved_name)) {
            psError(PM_ERR_SYS, false, "failed to delete %s\n%s.", sfile->resolved_name,
                strerror(errno));
            return false;
        }
    }
    return true;
}

bool
deleteTemps(streakFiles *sfiles)
{
    if (sfiles->outMask) {
        deleteFile(sfiles->outMask);
    }

    if (sfiles->outChMask) {
        deleteFile(sfiles->outChMask);
    }

    if (sfiles->outWeight) {
        deleteFile(sfiles->outWeight);
    }

    if (sfiles->outImage) {
        deleteFile(sfiles->outImage);
    }

    if (sfiles->outSources) {
        deleteFile(sfiles->outSources);
    }

    return true;
}

bool
isExciseValue(double value, double exciseValue)
{
    if (isnan(exciseValue)) {
        return isnan(value);
    } else {
        return value == exciseValue;
    }
}

void
setMaskedToNAN(streakFiles *sfiles, psU32 maskMask, bool printCounts)
{
        long maskedPixels = 0;
        long nandPixels = 0;
        long nandWeights = 0;

        psImage *image = sfiles->outImage->image;
        int imageType = image->type.type;
        double exciseValue = sfiles->inImage->exciseValue;

        psImage *mask = sfiles->inMask->image;
        psImage *weight = NULL;
        int weightType = 0;
        double weightExciseValue = NAN;
        if (sfiles->outWeight) {
            weight = sfiles->outWeight->image;
            weightType = weight->type.type;
            weightExciseValue = sfiles->inWeight->exciseValue;
        }

        if (printCounts) {
            psTimerStart("NAN_MASKED");
        }

        for (int y=0; y < sfiles->inImage->numRows; y++) {
            for (int x=0; x < sfiles->inImage->numCols; x++) {
                // these gets are not necessary, we could just set the pixels to nan
                // but I want to get the counts
                double imageVal = 0; // avoid compiler warning
                if (imageType == PS_TYPE_F32) {
                    imageVal  = image->data.F32[y][x];
                } else if (imageType == PS_TYPE_S16) {
                    imageVal = image->data.S16[y][x];
                } else if (imageType == PS_TYPE_U16) {
                    imageVal = image->data.U16[y][x];
                } else {
                    psError(PS_ERR_PROGRAMMING, true, "unexpected image type found: %d\n",
                                imageType);
                    streaksExit("", PS_EXIT_PROG_ERROR);
                }
                psU32 maskVal = 0;
                if (sfiles->stage == IPP_STAGE_RAW) {
                    unsigned int xChip, yChip;
                    cellToChipInt(&xChip, &yChip, sfiles->astrom, x, y);
                    maskVal = mask->data.PS_TYPE_IMAGE_MASK_DATA[yChip][xChip];
                } else {
                    maskVal = mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x];
                }
                if (maskVal & maskMask) {
                    ++maskedPixels;
                    if (!isExciseValue(imageVal, exciseValue)) {
                        ++nandPixels;
                        if (imageType == PS_TYPE_F32) {
                            image->data.F32[y][x] = exciseValue;
                        } else if (imageType == PS_TYPE_S16) {
                            image->data.S16[y][x] = exciseValue;
                        } else if (imageType == PS_TYPE_U16) {
                            image->data.U16[y][x] = exciseValue;
                        } else {
                            psError(PS_ERR_PROGRAMMING, true, "unexpected image type found: %d\n",
                                imageType);
                            streaksExit("", PS_EXIT_PROG_ERROR);
                        }
                    }
                    if (weight) {
                        if (weightType == PS_TYPE_F32) {
                            double weightVal = weight->data.F32[y][x];
                            if (!isnan(weightVal)) {
                                ++nandWeights;
                                weight->data.F32[y][x] = NAN;
                            }
                        } else if (weightType == PS_TYPE_F64) {
                            double weightVal = weight->data.F64[y][x];
                            if (!isnan(weightVal)) {
                                ++nandWeights;
                                weight->data.F64[y][x] = NAN;
                            }
                        } else if(weightType == PS_TYPE_S16) {
                            double weightVal = weight->data.S16[y][x];
                            if (!isExciseValue(weightVal, weightExciseValue)) {
                                ++nandWeights;
                                image->data.S16[y][x] = weightExciseValue;
                            }
                        } else if (weightType == PS_TYPE_U16) {
                            double weightVal = weight->data.U16[y][x];
                            if (!isExciseValue(weightVal, weightExciseValue)) {
                                ++nandWeights;
                                image->data.U16[y][x] = weightExciseValue;
                            }
                        } else {
                            psError(PS_ERR_PROGRAMMING, true, "unexpected image type found: %d\n",
                                    weightType);
                                streaksExit("", PS_EXIT_PROG_ERROR);
                        }
                    }
                }
            }
        }
        if (printCounts) {
            psLogMsg(sfiles->program_name, PS_LOG_INFO, "time to NAN mask pixels: %f\n", psTimerClear("NAN_MASKED"));
            long totalPixels = image->numRows * image->numCols;
            psLogMsg(sfiles->program_name, PS_LOG_INFO, "pixels:        %10ld\n", totalPixels);
            psLogMsg(sfiles->program_name, PS_LOG_INFO, "masked pixels: %10ld %4.2f%%\n", maskedPixels, 100. * maskedPixels / totalPixels);
            psLogMsg(sfiles->program_name, PS_LOG_INFO, "nand pixels:   %10ld %4.2f%%\n", nandPixels, 100. * nandPixels / totalPixels);
            if (weight) {
                psLogMsg(sfiles->program_name, PS_LOG_INFO, "nand weights:  %10ld %4.2f%%\n", nandWeights, 100. * nandWeights / totalPixels);
            }
        }
}

void
streaksNebulousCleanup(void)
{
    nebServerFree(ourNebServer);
    ourNebServer = NULL;
}
static void
setExciseValue(sFile *in)
{
    if (in->image->type.type == PS_TYPE_U16) {
        in->exciseValue = 65535;
        psMetadataAddU16(in->header, PS_LIST_TAIL, "BLANK", 0, "", 65535);
        psMetadataAddU16(in->header, PS_LIST_TAIL, "ZBLANK", 0, "", 65535);
    } else if (in->image->type.type == PS_TYPE_S16) {
        in->exciseValue = 32767;
        psMetadataAddU16(in->header, PS_LIST_TAIL, "BLANK", 0, "", 32767);
        psMetadataAddU16(in->header, PS_LIST_TAIL, "ZBLANK", 0, "", 32767);
    } else if (in->image->type.type == PS_TYPE_F32) {
        in->exciseValue = NAN;
    } else if (in->image->type.type == PS_TYPE_F64) {
        in->exciseValue = NAN;
    } else {
        psError(PS_ERR_PROGRAMMING, true, "unexpected image type found: %d\n", in->image->type.type);
        streaksExit("", PS_EXIT_PROG_ERROR);
    }
}

// Get the mask values that we need.
// If an input mask file is provided get them from there.
// Otherwise get them from the recipe MASKS
void
strkGetMaskValues(streakFiles *sfiles)
{
    if (sfiles->maskStreak) {
        // already done
        return;
    }
    if (!pmCensorGetMasks(sfiles->config, &sfiles->maskMask, &sfiles->maskStreak)) {
        streaksExit("failed to find mask values", PS_EXIT_CONFIG_ERROR);
    }
}

static void
doCopyExtraExtensions(streakFiles *sf, sFile *in, sFile *out)
{
    for (int extnum = sf->extnum; extnum < in->nHDU; extnum++) {
        moveExt(in, extnum);
        readImage(in, extnum, sf->stage, false);

        if (SFILE_IS_IMAGE(in)) {
            // Turn off compression (code adapted from pmFPAWrite)
            int bitpix = out->fits->options ? out->fits->options->bitpix : 0; // Desired bits per pixel
            psFitsCompression *compress = psFitsCompressionGet(out->fits); // Current compression options
            if (!psFitsSetCompression(out->fits, PS_FITS_COMPRESS_NONE, NULL, 0, 0, 0)) {
                psError(PM_ERR_UNKNOWN, false, "failed to turn off compression for extension %d\n", extnum);
                streaksExit("", PS_EXIT_UNKNOWN_ERROR);
            }
            if (out->fits->options) {
                out->fits->options->bitpix = 0;
            }
            if (!psFitsWriteImage(out->fits, in->header, in->image, 0, in->extname)) {
                psError(PM_ERR_UNKNOWN, false, "failed to write image for extension %d\n", extnum);
                streaksExit("", PS_EXIT_UNKNOWN_ERROR);
            }
            if (out->fits->options) {
                out->fits->options->bitpix = bitpix;
            }
            if (!psFitsCompressionApply(out->fits, compress)) {
                psError(PM_ERR_UNKNOWN, false, "failed to rest compression image for extension %d\n", extnum);
                streaksExit("", PS_EXIT_UNKNOWN_ERROR);
            }
            psFree(compress);
        } else {
            copyTable(out, in, extnum);
        }
    }
}

void
copyExtraExtensions(streakFiles *sf)
{
    // XXX: Bad assumption, the begining of the mask and weight files have the same structure as the image
    // So we are currently at the image portion
    sf->extnum = sf->inImage->nHDU;

    if (sf->inWeight && sf->outWeight) {
        doCopyExtraExtensions(sf, sf->inWeight, sf->outWeight);
    }
    if (sf->inMask && sf->outMask) {
        doCopyExtraExtensions(sf, sf->inMask, sf->outMask);
    }
    if (sf->inChMask && sf->outChMask) {
        doCopyExtraExtensions(sf, sf->inChMask, sf->outChMask);
    }
}
