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

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <getopt.h>
#include <string.h>
#include "pslib.h"
#include "stac.h"

void help(const char *programName
    )
{
    fprintf (stderr, "shift: shift an image, given the transformation\n"
             "Usage: %s [-h] [-v] [-s NX NY] IN OUT\n"
             "where\n"
             "\t-h           Help (this info)\n"
             "\t-v           Increase verbosity level\n"
             "\t-s NX NY     Size of output image\n"
             "\tIN           Input image, which has associated .map file\n"
             "\tOUT          Output image\n",
             programName
        );
}


int main(int argc, char **argv)
{
    // Set trace levels
    (void)psTraceSetLevel(".", 0);
    (void)psTraceSetLevel("stac.checkMemory", 10);
    (void)psTraceSetLevel("stac.read", 10);
    (void)psTraceSetLevel("stac.invertMaps", 10);
    (void)psTraceSetLevel("stac.transform", 10);
    (void)psTraceSetLevel("stac.size", 10);
    (void)psTraceSetLevel("shift", 10);

    // Set logging level
    psLogSetLevel(9);

    // Parameters with default values
    int outnx = 0, outny = 0;           // Size of output image
    int verbose = 0;                    // Verbosity level
    float bad = 0.0;                    // Bad level

    // Parse command line
    const char *programName = argv[0];  // Program name
    /* Variables for getopt */
    int opt;   /* Option, from getopt */
    extern char *optarg;   /* Argument accompanying switch */
    extern int optind;   /* getopt variables */

    /* Parse command-line arguments using getopt */
    while ((opt = getopt(argc, argv, "hvb:s:")) != -1) {
        switch (opt) {
          case 'h':
            help(programName);
            exit(EXIT_SUCCESS);
          case 'v':
            verbose++;
            break;
          case 's':
            if ((sscanf(argv[optind-1], "%d", &outnx) != 1) || (sscanf(argv[optind++], "%d", &outny) != 1)) {
                // Note: incrementing optind, so I can read more than one parameter.
                help(programName);
                exit(EXIT_FAILURE);
            }
            break;
          case 'b':
            if (sscanf(optarg, "%f", &bad) != 1) {
                help(programName);
                exit(EXIT_FAILURE);
            }
            break;
          default:
            help(programName);
            exit(EXIT_FAILURE);
        }
    }

    /* Get the remaining, mandatory values */
    argc -= optind;
    argv += optind;
    if (argc < 2) {
        help(programName);
        exit(EXIT_FAILURE);
    }
    const char *inName = argv[0];       // Input filename
    const char *outName = argv[1];      // Output filename

    psFits *imageFile = psFitsOpen(inName, "r");
    psRegion imageRegion = {0, 0, 0, 0}; // Region of image to read
    psMetadata *header = psFitsReadHeader(NULL, imageFile); // FITS header
    psImage *image = psFitsReadImage(imageFile, imageRegion, 0);
    if (image == NULL) {
        psErrorStackPrint(stderr, "Fatal error: Unable to read %s\n", inName);
        exit(EXIT_FAILURE);
    }
    psTrace("shift", 4, "Image %s is %dx%d\n", inName, image->numCols, image->numRows);
    // Convert to 32-bit floating point, in necessary
    if (image->type.type != PS_TYPE_F32) {
        psTrace("shift", 5, "Converting %s to floating point in memory....", inName);
        psImage *temp = psImageCopy(NULL, image, PS_TYPE_F32);
        psFree(image);
        image = temp;
    }
    psFitsClose(imageFile);

    // Generate masks
    psImage *mask = psImageAlloc(image->numCols, image->numRows, PS_TYPE_U8);
    for (int y = 0; y < image->numRows; y++) {
        for (int x = 0; x < image->numCols; x++) {
            if (image->data.F32[y][x] <= bad) {
                mask->data.U8[y][x] = 1;
            } else {
                mask->data.U8[y][x] = 0;
            }
        }
    }

    // Read map
    char mapName[MAXCHAR];              // Name of map file
    sprintf(mapName, "%s.map", inName);
    psPlaneTransform *map = stacReadMap(mapName);

    // Functions work on array of images, so:
    psArray *images = psArrayAlloc(1);  // An array of one
    psArray *maps = psArrayAlloc(1);    // An array of one
    psArray *masks = psArrayAlloc(1);   // An array of one
    images->data[0] = image;
    maps->data[0] = map;
    masks->data[0] = mask;

    // Get size
    if (outnx == 0 || outny == 0) {
        psVector *xSize = psVectorAlloc(1, PS_TYPE_S32); // A vector of one
        psVector *ySize = psVectorAlloc(1, PS_TYPE_S32); // A vector of one
        xSize->data.S32[0] = image->numCols;
        ySize->data.S32[0] = image->numRows;
        stacSize(&outnx, &outny, NULL, NULL, xSize, ySize, maps);
        psFree(xSize);
        psFree(ySize);
    }

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

    // Transform inputs and errors
    psArray *transformed = NULL;
    (void)stacTransform(&transformed, NULL, images, inverseMaps, NULL, masks, NULL, NULL, NULL, outnx, outny);

    // Update FITS header appropriately
    psMetadataAddS32(header, PS_LIST_TAIL, "NAXIS1", PS_META_REPLACE, "Number of pixels in x",
                     (int)((psImage*)transformed->data[0])->numCols);
    psMetadataAddS32(header, PS_LIST_TAIL, "NAXIS2", PS_META_REPLACE, "Number of pixels in y",
                     (int)((psImage*)transformed->data[0])->numRows);
    psMetadataAddS32(header, PS_LIST_TAIL, "BITPIX", PS_META_REPLACE, "Bits per pixel", -32);

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

    // Free everything I've used
    psFree(images);
    psFree(maps);
    psFree(masks);
    psFree(inverseMaps);
    psFree(transformed);

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

}
