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

#include "pmWarpMap.h"

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Allocators and deallocators
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

static void warpMapFree(pswarpMap *map)
{
    // No operations required; this function exists merely to identify a pmWarpMap
    return;
}

pswarpMap *pswarpMapAlloc(void)
{
    pswarpMap *map = (pswarpMap *) psAlloc (sizeof(pswarpMap));
    psMemSetDeallocator(map, (psFreeFunc)warpMapFree);

    map->Xo = map->Xx = map->Xy = NAN;
    map->Yo = map->Yx = map->Yy = NAN;
    map->xo = map->yo = NAN;

    return map;
}


static void warpMapGridFree(pswarpMapGrid *grid)
{
    if (!grid->maps) {
        // Nothing to free
        return;
    }

    for (int i = 0; i < grid->nXpts; i++) {
        for (int j = 0; j < grid->nYpts; j++) {
            psFree(grid->maps[i][j]);
        }
        psFree(grid->maps[i]);
    }
    psFree (grid->maps);
    return;
}

pswarpMapGrid *pswarpMapGridAlloc(int nXpts, int nYpts)
{
    PS_ASSERT_INT_POSITIVE(nXpts, NULL);
    PS_ASSERT_INT_POSITIVE(nYpts, NULL);

    pswarpMapGrid *grid = (pswarpMapGrid *) psAlloc(sizeof(pswarpMapGrid));
    psMemSetDeallocator(grid, (psFreeFunc)warpMapGridFree);

    grid->maps = psAlloc(nXpts*sizeof(void **));
    for (int i = 0; i < nXpts; i++) {
        grid->maps[i] = psAlloc(nYpts*sizeof(void *));
        for (int j = 0; j < nYpts; j++) {
            grid->maps[i][j] = pswarpMapAlloc();
        }
    }
    grid->nXpts = nXpts;
    grid->nYpts = nYpts;

    grid->nXpix = 0;
    grid->nYpix = 0;

    return grid;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Operations
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// XXX for the moment, ignore readout->cell->chip offsets
pswarpMapGrid *pswarpMapGridFromImage(pmSkyCell *dest, const pmReadout *src, int nXpix, int nYpix)
{
    PS_ASSERT_PTR_NON_NULL(dest, NULL);
    PS_ASSERT_PTR_NON_NULL(src, NULL);
    PS_ASSERT_IMAGE_NON_NULL(src->image, NULL);
    PS_ASSERT_INT_NON_NEGATIVE(nXpix, NULL);
    PS_ASSERT_INT_NON_NEGATIVE(nYpix, NULL);

    // split the difference of the remainder
    int xMin = 0.5 * (src->image->numCols % nXpix); // Minimum in x
    int yMin = 0.5 * (src->image->numRows % nYpix); // Minimum in y

    int nXpts = src->image->numCols / nXpix; // Number of points in x
    if (src->image->numCols % nXpix == 0) {
        nXpts++;
    }
    int nYpts = src->image->numRows / nYpix; // Number of points in y
    if (src->image->numRows % nYpix == 0) {
        nYpts++;
    }

    pmWarpMapGrid *grid = pmWarpMapGridAlloc(nXpts, nYpts);
    grid->nXpix = nXpix;
    grid->nYpix = nYpix;
    grid->xMin = xMin;
    grid->yMin = yMin;

    for (int ni = 0, i = xMin; ni < nXpts; i += nXpix, ni++) {
        for (int nj = 0, j = yMin; nj < nYpts; j += nYpix, nj++) {
            pmWarpMapSetLocalModel(grid->maps[ni][nj], dest, src, i, j);
        }
    }

    return grid;
}

bool pmWarpMapGridSetGrid(const pmWarpMapGrid *grid, int ix, int iy, int *gridX, int *gridY)
{
    *gridX = (ix - grid->xMin + 0.5*grid->nXpix) / grid->nXpix;
    *gridY = (iy - grid->yMin + 0.5*grid->nYpix) / grid->nYpix;

    return true;
}

bool pmWarpMapGridNextGrid(const pmWarpMapGrid *grid, int gridX, int gridY, int *nextX, int *nextY)
{
    *nextX = gridX*grid->nXpix + grid->xMin + 0.5*grid->nXpix;
    *nextY = gridY*grid->nYpix + grid->yMin + 0.5*grid->nYpix;

    return true;
}

// measure the max error accumulated in appling one grid point to its neighbors
double pmWarpMapGridMaxError(const pmWarpMapGrid *grid)
{
    double maxError = 0;                // Maximum error

    for (int i = 0; i < grid->nXpts - 1; i++) {
        for (int j = 0; j < grid->nYpts - 1; j++) {
            // measure the output coordinates for the next grid position using the current grid map
            // compare with the coordinates measured using the next grid map
            double xRaw, yRaw;          // Raw position
            double xRef, yRef;          // Reference position

            pmWarpMapApply(&xRaw, &yRaw, grid->maps[i][j], grid->maps[i][j]->xo + grid->nXpix,
                           grid->maps[i][j]->yo);
            pmWarpMapApply(&xRef, &yRef, grid->maps[i+1][j], grid->maps[i][j]->xo + grid->nXpix,
                           grid->maps[i][j]->yo);

            double posError = hypot(xRaw-xRef, yRaw-yRef);
            maxError = PS_MAX(maxError, posError);
        }
    }

    return maxError;
}

bool pmWarpMapApply(double *outX, double *outY, const pmWarpMap *map, double inX, double inY)
{
    *outX = map->Xo + map->Xx*inX + map->Xy*inY;
    *outY = map->Yo + map->Yx*inX + map->Yy*inY;

    return true;
}

bool pmWarpMapSetLocalModel(pmWarpMap *map, const pmSkyCell *dest, const pmReadout *src, int ix, int iy)
{
    PS_ASSERT_PTR_NON_NULL(map, false);
    PS_ASSERT_PTR_NON_NULL(dest, false);
    PS_ASSERT_PTR_NON_NULL(dest->toSky, false);
    PS_ASSERT_PTR_NON_NULL(dest->fromTPA, false);
    PS_ASSERT_PTR_NON_NULL(src, false);
    pmCell *cell = src->parent;         // Parent cell of source
    PS_ASSERT_PTR_NON_NULL(cell, false);
    pmChip *chip = cell->parent;        // Parent chip of source
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_PTR_NON_NULL(chip->toFPA, false);
    pmFPA *fpa = chip->parent;          // Parent FPA of source
    PS_ASSERT_PTR_NON_NULL(fpa, false);
    PS_ASSERT_PTR_NON_NULL(fpa->toTPA, false);
    PS_ASSERT_PTR_NON_NULL(fpa->toSky, false);
    PS_ASSERT_INT_NON_NEGATIVE(ix, false);
    PS_ASSERT_INT_NON_NEGATIVE(iy, false);

    // These will get allocated in the course of the calculations; may as well allocate them here
    // XXX save these as static for speed?
    // Alternatively, pass in a buffer containing these pre-allocated.
    psPlane *offset = psPlaneAlloc();   // Position offset from the main position of interest (ix,iy)
    psPlane *FP = psPlaneAlloc();       // Position in the focal plane
    psPlane *TP = psPlaneAlloc();       // Position in the tangent plane
    psSphere *sky = psSphereAlloc();    // Position on the sky

    // XXX need to include readout->cell->chip offsets
    // --> Look up CELL.X0, CELL.Y0 in concepts
    // --> Use readout->col0, readout->row0

    // V(0,0) position
    offset->x = ix;
    offset->y = iy;
    psPlaneTransformApply(FP, chip->toFPA, offset);
    psPlaneTransformApply (TP, fpa->toTPA, FP);
    psDeproject(sky, TP, fpa->toSky);
    psProject(TP, sky, dest->toSky);
    psPlane *V00 = psPlaneTransformApply(NULL, dest->fromTPA, TP); // Coordinates at the position of interest

    // V(1,0) position
    offset->x = ix + 1;
    offset->y = iy;
    psPlaneTransformApply(FP, chipSrc->toFPA, offset);
    psPlaneTransformApply(TP, fpaSrc->toTPA, FP);
    psDeproject(sky, TP, fpaSrc->toSky);
    psProject(TP, sky, dest->toSky);
    psPlane *V10 = psPlaneTransformApply(NULL, dest->fromTPA, TP); // Coordinates at position + (1,0)

    // V(0,1) position
    offset->x = ix;
    offset->y = iy + 1;
    psPlaneTransformApply(FP, chipSrc->toFPA, offset);
    psPlaneTransformApply(TP, fpaSrc->toTPA, FP);
    psDeproject(sky, TP, fpaSrc->toSky);
    psProject(TP, sky, fpaDest->toSky);
    psPlane *V01 = psPlaneTransformApply(NULL, dest->fromTPA, TP); // Coordinates at position + (0,1)

    psFree(offset);
    psFree(FP);
    psFree(TP);
    psFree(sky);

    // Generate local linear transformation
    map->Xx = V10->x - V00->x;
    map->Xy = V01->x - V00->x;
    map->Xo = V00->x - map->Xx * ix - map->Xy * iy;

    map->Yx = V10->y - V00->y;
    map->Yy = V01->y - V00->y;
    map->Yo = V00->y - map->Yx * ix - map->Yy * iy;

    map->xo = ix;
    map->yo = iy;

    psFree(V00);
    psFree(V10);
    psFree(V01);

    return true;
}
