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

#include <stdio.h>
#include <assert.h>
#include "pslib.h"
#include "stac.h"


#define MAX(x,y) ((x) > (y) ? (x) : (y))
#define NUM_GRID 20

psArray *stacInvertMaps(const psArray *maps, // Array of maps to invert
                        int outnx, int outny // Size of output image
    )
{
    int nMaps = maps->n;                // Number of maps
    psArray *inverted = psArrayAlloc(nMaps); // Array of inverted maps for output

    psTrace("stac.invertMaps", 1, "Inverting maps....\n");

    // Coordinates for the transformations
    psPlane *inCoord = psAlloc(sizeof(psPlane));
    psPlane *outCoord = psAlloc(sizeof(psPlane));

    for (int mapNum = 0; mapNum < nMaps; mapNum++) {

        psPlaneTransform *oldMap = (psPlaneTransform*)maps->data[mapNum]; // Uninverted map
        // Check input
        assert(oldMap->x->nX == oldMap->x->nY && oldMap->y->nX == oldMap->y->nY &&
               oldMap->x->nX == oldMap->y->nX);
        int order = oldMap->x->nX;      // Polynomial order
        psTrace("stac.invertMaps", 4, "Generating order %d polynomial inverse transformation.\n", order);
        psPlaneTransform *newMap = psPlaneTransformAlloc(order, order); // Inverted map

        // Create fake polynomial to use in evaluation
        psPolynomial2D *fakePoly = psPolynomial2DAlloc(PS_POLYNOMIAL_ORD, order, order);
        for (int i = 0; i <= order; i++) {
            for (int j = 0; j <= order; j++) {
                fakePoly->coeff[i][j] = 1.0; // Set all coeffecients to 1
                fakePoly->mask[i][j] = 1; // Mask all coefficients; unmask to evaluate
            }
        }

        // A grid of xin,yin --> xout,yout
        psVector *xIn = psVectorAlloc(NUM_GRID * NUM_GRID, PS_TYPE_F32);
        psVector *yIn = psVectorAlloc(NUM_GRID * NUM_GRID, PS_TYPE_F32);
        psVector *xOut = psVectorAlloc(NUM_GRID * NUM_GRID, PS_TYPE_F32);
        psVector *yOut = psVectorAlloc(NUM_GRID * NUM_GRID, PS_TYPE_F32);

        // Create grid of points
        for (int yint = 0; yint < NUM_GRID; yint++) {
            inCoord->y = (float)(yint * outny) / (float)(NUM_GRID - 1);
            for (int xint = 0; xint < NUM_GRID; xint++) {
                inCoord->x = (float)(xint * outnx) / (float)(NUM_GRID - 1);

                (void)psPlaneTransformApply(outCoord, oldMap, inCoord);

                xOut->data.F32[yint*NUM_GRID + xint] = inCoord->x;
                yOut->data.F32[yint*NUM_GRID + xint] = inCoord->y;
                xIn->data.F32[yint*NUM_GRID + xint] = outCoord->x;
                yIn->data.F32[yint*NUM_GRID + xint] = outCoord->y;
            }
        }

        // Initialise the matrix and vectors
        int nCoeff = (order + 1) * (order + 2) / 2; // Number of polynomial coefficients
        psImage *matrix = psImageAlloc(nCoeff, nCoeff, PS_TYPE_F64); // Matrix for solution
        psVector *xVector = psVectorAlloc(nCoeff, PS_TYPE_F64); // Vector for solution in x
        psVector *yVector = psVectorAlloc(nCoeff, PS_TYPE_F64); // Vector for solution in y
        for (int i = 0; i < nCoeff; i++) {
            for (int j = 0; j < nCoeff; j++) {
                matrix->data.F64[i][j] = 0.0;
            }
            xVector->data.F64[i] = 0.0;
            yVector->data.F64[i] = 0.0;
        }

        // Iterate over the grid points
        for (int g = 0; g < NUM_GRID*NUM_GRID; g++) {

            // Iterate over the polynomial coefficients, accumulating the matrix and vectors
            for (int i = 0, ijIndex = 0; i <= order; i++) {
                for (int j = 0; j <= order - i; j++, ijIndex++) {

                    fakePoly->mask[i][j] = 0;
                    double ijPoly = psPolynomial2DEval(fakePoly, xIn->data.F32[g], yIn->data.F32[g]);
                    fakePoly->mask[i][j] = 1;

                    for (int m = 0, mnIndex = 0; m <= order; m++) {
                        for (int n = 0; n <= order - m; n++, mnIndex++) {

                            fakePoly->mask[m][n] = 0;
                            double mnPoly = psPolynomial2DEval(fakePoly, xIn->data.F32[g], yIn->data.F32[g]);
                            fakePoly->mask[m][n] = 1;

                            matrix->data.F64[ijIndex][mnIndex] += ijPoly * mnPoly;
                        }
                    }

                    xVector->data.F64[ijIndex] += ijPoly * (double)xOut->data.F32[g];
                    yVector->data.F64[ijIndex] += ijPoly * (double)yOut->data.F32[g];
                }
            } // Iterating over coefficients
        } // Iterating over grid points

        // Solution via LU Decomposition
        psVector *permutation = psVectorAlloc(nCoeff, PS_TYPE_F64); // Permutation vector for LU Decomposition
        psImage *luMatrix = psMatrixLUD(NULL, &permutation, matrix); // LU decomposed matrix
        psVector *xSolution = psMatrixLUSolve(NULL, luMatrix, xVector, permutation); // Solution in x
        psVector *ySolution = psMatrixLUSolve(NULL, luMatrix, yVector, permutation); // Solution in y

        // Stuff coefficients into transformation
        for (int i = 0, ijIndex = 0; i <= order; i++) {
            for (int j = 0; j <= order - i; j++, ijIndex++) {
                newMap->x->coeff[i][j] = xSolution->data.F64[ijIndex];
                newMap->y->coeff[i][j] = ySolution->data.F64[ijIndex];
            }
        }
        inverted->data[mapNum] = newMap;

#ifdef TESTING
        // Print x coefficients
        psTrace("stac.invertMaps", 7, "x' = \n");
        for (int i = 0; i <= order; i++) {
            for (int j = 0; j <= order - i; j++) {
                psTrace("stac.invertMaps", 7, "      %f x^%d y^%d\n", newMap->x->coeff[i][j], i, j);
            }
        }
        // Print y coefficients
        psTrace("stac.invertMaps", 7, "y' = \n");
        for (int i = 0; i <= order; i++) {
            for (int j = 0; j <= order - i; j++) {
                psTrace("stac.invertMaps", 7, "      %f x^%d y^%d\n", newMap->y->coeff[i][j], i, j);
            }
        }
#endif

#ifdef TESTING
        // Go forward then backward
        double x = 123.4;
        double y = 432.1;
        psPlane *oldCoords = psAlloc(sizeof(psPlane));
        oldCoords->x = x;
        oldCoords->y = y;
        psPlane *newCoords = psPlaneTransformApply(NULL, oldMap, oldCoords);
        psTrace("stac.invertMaps.test", 5, "%f,%f --> %f,%f\n", x, y, newCoords->x, newCoords->y);
        (void)psPlaneTransformApply(oldCoords, newMap, newCoords);
        psTrace("stac.invertMaps.test", 5, "--------> %f,%f\n", oldCoords->x, oldCoords->y);
        psFree(newCoords);
        psFree(oldCoords);
#endif

        psFree(permutation);
        psFree(luMatrix);
        psFree(matrix);
        psFree(xVector);
        psFree(yVector);
        psFree(xSolution);
        psFree(ySolution);
        psFree(fakePoly);
        psFree(xIn);
        psFree(yIn);
        psFree(xOut);
        psFree(yOut);
    }

    psFree(inCoord);
    psFree(outCoord);


#if 0
        // Can't handle higher order than linear yet
        assert(((psPlaneTransform*)maps->data[i])->x->nX == 2 &&
               ((psPlaneTransform*)maps->data[i])->x->nY == 2 &&
               ((psPlaneTransform*)maps->data[i])->y->nX == 2 &&
               ((psPlaneTransform*)maps->data[i])->y->nY == 2);
        psPlaneTransform *newMap = psPlaneTransformAlloc(2, 2); // Inverted map
        psPlaneTransform *oldMap = (psPlaneTransform*)maps->data[i]; // Uninverted map

        // Now, simply do a 2x2 matrix inversion

        double a = oldMap->x->coeff[1][0];
        double b = oldMap->x->coeff[0][1];
        double c = oldMap->y->coeff[1][0];
        double d = oldMap->y->coeff[0][1];
        double e = oldMap->x->coeff[0][0];
        double f = oldMap->y->coeff[0][0];

        double invDet = 1.0 / (a * d - b * c); // Inverse of the determinant

        // Not entirely sure why this works, but it appears to do so.......................................
        newMap->x->coeff[1][0] = invDet * a;
        newMap->x->coeff[0][1] = - invDet * b;
        newMap->y->coeff[1][0] = - invDet * c;
        newMap->y->coeff[0][1] = invDet * d;

        newMap->x->coeff[0][0] = - invDet * (d * e + c * f);
        newMap->y->coeff[0][0] = - invDet * (b * e + a * f);
#endif


    return inverted;
}


