#include <stdio.h>
#include <math.h>
#include <pslib.h>

#define IMAGE_SIZE 4096                 // Size of output image
#define INTERPOLATION PS_INTERPOLATE_BICUBE // Interpolation kernel
#define SHIFT_X 0.1                     // Number of pixels to shift in x
#define SHIFT_Y 0.2                     // Number of pixels to shift in y
#define SCALE (2.0 * M_LN2)             // Scale difference
#define MASK_BAD 0xF0                   // Mask for bad pixels
#define MASK_POOR 0x0F                  // Mask for poor pixels
#define POOR_FRAC 0.1                   // Poor fraction
#define ROTATION (30 / 180 * M_PI)      // Rotation

// ( xOut ) =  scale * ( cos  -sin ) ( xIn ) + ( xShift )
// ( yOut )            ( sin   cos ) ( yIn )   ( yShift )


static inline void transformIn(float *xIn, float *yIn, // Coordinates in input image
                               int xOut, int yOut // Coordinates in output image
    )
{
    float xShift = (xOut - SHIFT_X) / SCALE;
    float yShift = (yOut - SHIFT_Y) / SCALE;

    *xIn = xShift * cos(ROTATION) + yShift * sin(ROTATION);
    *yIn = - xShift * sin(ROTATION) + yShift * cos(ROTATION);
    return;
}

static inline void transformOut(float *xOut, float *yOut, // Coordinates in output image
                                int xIn, int yIn // Coordinates in input image
    )
{
    *xOut = SCALE * (xIn * cos(ROTATION) - yIn * sin(ROTATION)) + SHIFT_X;
    *yOut = SCALE * (xIn * sin(ROTATION) + yIn * cos(ROTATION)) + SHIFT_Y;
    return;
}

static psImage *readImage(const char *name)
{
    psFits *fits = psFitsOpen(name, "r");
    psImage *image = psFitsReadImage(fits, psRegionSet(0,0,0,0), 0);
    psFitsClose(fits);
    return image;
}

void writeImage(const psImage *image,   // Image to write
                const char *name        // Name of FITS file
    )
{
    psFits *fits = psFitsOpen(name, "w");
    psFitsWriteImage(fits, NULL, image, 0, NULL);
    psFitsClose(fits);
}

int main(int argc, char *argv[])
{
    psImage *inBright = readImage("bright.fits");
    psImage *inFaint = readImage("faint.fits");
    psImage *inWeight = readImage("weight.fits");
    psImage *inMask = readImage("mask.fits");

    psImage *outBright = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_F32);
    psImage *outFaint = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_F32);
    psImage *outWeight = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_F32);
    psImage *outMask = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_MASK);

    int inNumCols = inBright->numCols, inNumRows = inBright->numRows; // Size of input image

    // Transform forward
    {
        psImageInterpolateOptions *interp1 = psImageInterpolateOptionsAlloc(INTERPOLATION, inBright,
                                                                            NULL, inMask, 0xFF, NAN, NAN,
                                                                            MASK_BAD, MASK_POOR, POOR_FRAC);
        psImageInterpolateOptions *interp2 = psImageInterpolateOptionsAlloc(INTERPOLATION, inFaint,
                                                                            inWeight, inMask, 0xFF, NAN, NAN,
                                                                            MASK_BAD, MASK_POOR, POOR_FRAC);

        for (int yOut = 0; yOut < IMAGE_SIZE; yOut++) {
            for (int xOut = 0; xOut < IMAGE_SIZE; xOut++) {
                float xIn, yIn;
                transformIn(&xIn, &yIn, xOut + 0.5, yOut + 0.5);
                xIn += 0.5;
                yIn += 0.5;

                if (xIn < 0 || xIn >= inNumCols || yIn < 0 || yIn >= inNumRows) {
                    outBright->data.F32[yOut][xOut] = NAN;
                    outFaint->data.F32[yOut][xOut] = NAN;
                    outWeight->data.F32[yOut][xOut] = NAN;
                    outMask->data.PS_TYPE_MASK_DATA[yOut][xOut] = MASK_BAD;
                    continue;
                }

                double brightValue, faintValue, weightValue;
                psMaskType maskValue = inMask->data.PS_TYPE_MASK_DATA[(int)yIn][(int)xIn]; // Value of mask
                psImageInterpolate(&brightValue, NULL, NULL, xIn, yIn, interp1);
                psImageInterpolate(&faintValue, &weightValue, &maskValue, xIn, yIn, interp2);

                outBright->data.F32[yOut][xOut] = brightValue / SCALE;
                outFaint->data.F32[yOut][xOut] = faintValue / SCALE;
                outWeight->data.F32[yOut][xOut] = weightValue / PS_SQR(SCALE);
                outMask->data.PS_TYPE_MASK_DATA[yOut][xOut] = maskValue;
            }
        }

        psFree(interp1);
        psFree(interp2);
    }

    int outNumCols = outBright->numCols, outNumRows = inBright->numRows; // Size of output image

    // Transform back
    {
        psImageInterpolateOptions *interp1 = psImageInterpolateOptionsAlloc(INTERPOLATION, outBright,
                                                                            NULL, outMask, 0xFF, NAN, NAN,
                                                                            MASK_BAD, MASK_POOR, POOR_FRAC);
        psImageInterpolateOptions *interp2 = psImageInterpolateOptionsAlloc(INTERPOLATION, outFaint,
                                                                            outWeight, outMask, 0xFF,
                                                                            NAN, NAN, MASK_BAD, MASK_POOR,
                                                                            POOR_FRAC);

        for (int yIn = 0; yIn < inNumRows; yIn++) {
            for (int xIn = 0; xIn < inNumCols; xIn++) {
                float xOut, yOut;
                transformOut(&xOut, &yOut, xIn + 0.5, yIn + 0.5);
                xOut += 0.5;
                yOut += 0.5;

                if (xOut < 0 || xOut >= outNumCols || yOut < 0 || yOut >= outNumRows) {
                    inBright->data.F32[yIn][xIn] = NAN;
                    inFaint->data.F32[yIn][xIn] = NAN;
                    inWeight->data.F32[yIn][xIn] = NAN;
                    inMask->data.PS_TYPE_MASK_DATA[yIn][xIn] = MASK_BAD;
                    continue;
                }

                double brightValue, faintValue, weightValue;
                psMaskType maskValue = inMask->data.PS_TYPE_MASK_DATA[(int)yOut][(int)xOut]; // Value of mask
                psImageInterpolate(&brightValue, NULL, NULL, xOut, yOut, interp1);
                psImageInterpolate(&faintValue, &weightValue, &maskValue, xOut, yOut, interp2);

                inBright->data.F32[yIn][xIn] = brightValue * SCALE;
                inFaint->data.F32[yIn][xIn] = faintValue * SCALE;
                inWeight->data.F32[yIn][xIn] = weightValue * PS_SQR(SCALE);
                inMask->data.PS_TYPE_MASK_DATA[yIn][xIn] = maskValue;
            }
        }

        psFree(interp1);
        psFree(interp2);
    }

    writeImage(inBright, "bright.shift.fits");
    writeImage(inFaint, "faint.shift.fits");
    writeImage(inWeight, "weight.shift.fits");
    writeImage(inMask, "mask.shift.fits");

    psFree(inBright);
    psFree(inFaint);
    psFree(inWeight);
    psFree(inMask);

    psFree(outBright);
    psFree(outFaint);
    psFree(outWeight);
    psFree(outMask);

#if 0

    psArray *stars = psVectorsReadFromFile("fake_stars.dat", "%d %d");
    psVector *x = stars->data[0], *y = stars->data[1]; // Star coordinates
    FILE *starFile = fopen("fake_stars.shift.dat", "w");
    int num = x->n;                     // Number of stars
    for (int i = 0; i < num; i++) {
        float xOut, yOut;
        transformOut(&xOut, &yOut, x->data.S32[i], y->data.S32[i]);
        fprintf(starFile, "%f %f\n", xOut, yOut);
    }
    fclose(starFile);
#endif

    return PS_EXIT_SUCCESS;
}
