/*  @file  psImageConvolve.c
 *
 *  @brief Contains FFT transform related functions for psImage.
 *
 *  @author Robert DeSonia, MHPCC
 *
 *  @version $Revision: 1.11 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2005-01-05 21:16:14 $
 *
 *  Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
 */

#include <string.h>

#include "psImageConvolve.h"
#include "psImageFFT.h"
#include "psImageExtraction.h"
#include "psMatrixVectorArithmetic.h"
#include "psMemory.h"
#include "psLogMsg.h"
#include "psError.h"
#include "psImageIO.h"

#include "psImageErrors.h"

#define FOURIER_PADDING 32 /* padding amount in every side of the image for fourier convolution */

static void freeKernel(psKernel* ptr);

psKernel* psKernelAlloc(psS32 xMin, psS32 xMax, psS32 yMin, psS32 yMax)
{
    psKernel* result;
    psS32 numRows;
    psS32 numCols;

    // following is explicitly spelled out in the SDRS as a requirement
    if (yMin > yMax) {
        psLogMsg(__func__, PS_LOG_WARN,
                 "Specified yMin, %d, was greater than yMax, %d.  Values swapped.",
                 yMin, yMax);

        psS32 temp = yMin;
        yMin = yMax;
        yMax = temp;
    }

    // following is explicitly spelled out in the SDRS as a requirement
    if (xMin > xMax) {
        psLogMsg(__func__, PS_LOG_WARN,
                 "Specified xMin, %d, was greater than xMax, %d.  Values swapped.",
                 xMin, xMax);

        psS32 temp = xMin;
        xMin = xMax;
        xMax = temp;
    }

    numRows = yMax - yMin + 1;
    numCols = xMax - xMin + 1;

    result = psAlloc(sizeof(psKernel));
    result->xMin = xMin;
    result->xMax = xMax;
    result->yMin = yMin;
    result->yMax = yMax;
    result->image = psImageAlloc(numCols,numRows,PS_TYPE_KERNEL);
    memset(result->image->rawDataBuffer,0,numCols*numRows*PSELEMTYPE_SIZEOF(PS_TYPE_KERNEL));
    result->p_kernelRows = psAlloc(sizeof(psKernelType*)*numRows);

    psKernelType** kernelRows = result->p_kernelRows;
    psKernelType** imageRows = result->image->data.PS_TYPE_KERNEL_DATA;
    for (psS32 i = 0; i < numRows; i++) {
        kernelRows[i] = imageRows[i] - xMin;
    }
    result->kernel = kernelRows - yMin;

    p_psMemSetDeallocator(result,(psFreeFcn)freeKernel);

    return result;
}

void freeKernel(psKernel* ptr)
{
    if (ptr != NULL) {
        psFree(ptr->image);
        psFree(ptr->p_kernelRows);
    }
}

psKernel* psKernelGenerate(const psVector* tShifts,
                           const psVector* xShifts,
                           const psVector* yShifts,
                           psBool relative)
{
    psS32 lastX;
    psS32 lastY;
    psS32 lastT;
    psS32 x;
    psS32 y;
    psS32 t;
    psS32 xMin = 0;
    psS32 xMax = 0;
    psS32 yMin = 0;
    psS32 yMax = 0;
    psS32 length = 0;
    psKernelType normalizeTime = 1.0;  // fraction of total time for each shift clock
    psKernel* result = NULL;
    psKernelType** kernel = NULL;

    // got non-NULL vectors?
    if (tShifts == NULL || xShifts == NULL || yShifts == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psImageConvolve_SHIFT_NULL);
        return NULL;
    }

    // types match?
    if (xShifts->type.type != yShifts->type.type ||
            tShifts->type.type != xShifts->type.type) {
        char* typeXStr;
        char* typeYStr;
        char* typeTStr;
        PS_TYPE_NAME(typeXStr,xShifts->type.type);
        PS_TYPE_NAME(typeYStr,yShifts->type.type);
        PS_TYPE_NAME(typeTStr,tShifts->type.type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImageConvolve_SHIFT_TYPE_MISMATCH,
                typeTStr, typeXStr, typeYStr);
        return NULL;
    }

    // sizes match?
    length = xShifts->n;
    if (length != yShifts->n ||
            length != tShifts->n) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true,
                "Shift vectors can not be of different sizes.");
        return NULL;
    }

    // if no shifts, the kernel is just a 1 at 0,0
    if (length < 1) {
        result = psKernelAlloc(0,0,0,0);
        result->kernel[0][0] = 1;
        return result;
    }

    #define KERNEL_GENERATE_CASE(TYPE) \
case PS_TYPE_##TYPE: { \
        ps##TYPE *tShiftData = tShifts->data.TYPE; \
        ps##TYPE *xShiftData = xShifts->data.TYPE; \
        ps##TYPE *yShiftData = yShifts->data.TYPE; \
        lastX = xShiftData[length-1]; \
        lastY = yShiftData[length-1]; \
        lastT = tShiftData[length-1]; \
        \
        for (int lcv = 0; lcv < length; lcv++) { \
            x = lastX - xShiftData[lcv]; \
            y = lastY - yShiftData[lcv]; \
            \
            if (x < xMin) { \
                xMin = x; \
            } else if (x > xMax) { \
                xMax = x; \
            } \
            if (y < yMin) { \
                yMin = y; \
            } else if (y > yMax) { \
                yMax = y; \
            } \
        } \
        \
        normalizeTime = 1.0 / (psKernelType)(tShiftData[length-1]); \
        result = psKernelAlloc(xMin,xMax,yMin,yMax); \
        kernel = result->kernel; \
        \
        psS32 prevT = 0; \
        for (int i = 0; i < length; i++) { \
            t = tShiftData[i] - prevT; \
            x = lastX - xShiftData[i]; \
            y = lastY - yShiftData[i]; \
            \
            kernel[y][x] += (psKernelType)t / (psKernelType)lastT; \
            prevT = tShiftData[i]; \
        } \
        break; \
    }

    #define RELATIVE_KERNEL_GENERATE_CASE(TYPE) \
case PS_TYPE_##TYPE: { \
        ps##TYPE *tShiftData = tShifts->data.TYPE; \
        ps##TYPE *xShiftData = xShifts->data.TYPE; \
        ps##TYPE *yShiftData = yShifts->data.TYPE; \
        \
        x = 0; \
        y = 0; \
        t = 0; \
        \
        for (int lcv = length-1; lcv >= 0; lcv--) { \
            t += tShiftData[lcv]; \
            \
            if (x < xMin) { \
                xMin = x; \
            } else if (x > xMax) { \
                xMax = x; \
            } \
            if (y < yMin) { \
                yMin = y; \
            } else if (y > yMax) { \
                yMax = y; \
            } \
            x -= xShiftData[lcv]; \
            y -= yShiftData[lcv]; \
            \
        } \
        result = psKernelAlloc(xMin,xMax,yMin,yMax); \
        kernel = result->kernel; \
        \
        normalizeTime = 1.0 / (psKernelType)t; \
        x = 0; \
        y = 0; \
        for (psS32 i = length-1; i >= 0; i--) { \
            kernel[y][x] += (psKernelType)(tShiftData[i]) * normalizeTime; \
            x -= xShiftData[i]; \
            y -= yShiftData[i]; \
            \
        } \
        break; \
    }

    if (relative) {
        switch (xShifts->type.type) {
            RELATIVE_KERNEL_GENERATE_CASE(U8);
            RELATIVE_KERNEL_GENERATE_CASE(U16);
            RELATIVE_KERNEL_GENERATE_CASE(U32);
            RELATIVE_KERNEL_GENERATE_CASE(U64);
            RELATIVE_KERNEL_GENERATE_CASE(S8);
            RELATIVE_KERNEL_GENERATE_CASE(S16);
            RELATIVE_KERNEL_GENERATE_CASE(S32);
            RELATIVE_KERNEL_GENERATE_CASE(S64);
            RELATIVE_KERNEL_GENERATE_CASE(F32);
            RELATIVE_KERNEL_GENERATE_CASE(F64);
            RELATIVE_KERNEL_GENERATE_CASE(C32);
            RELATIVE_KERNEL_GENERATE_CASE(C64);

        default: {
                char* typeStr;
                PS_TYPE_NAME(typeStr,xShifts->type.type);
                psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                        PS_ERRORTEXT_psImage_IMAGE_TYPE_UNSUPPORTED,
                        typeStr);
            }
        }
    } else {
        switch (xShifts->type.type) {
            KERNEL_GENERATE_CASE(U8);
            KERNEL_GENERATE_CASE(U16);
            KERNEL_GENERATE_CASE(U32);
            KERNEL_GENERATE_CASE(U64);
            KERNEL_GENERATE_CASE(S8);
            KERNEL_GENERATE_CASE(S16);
            KERNEL_GENERATE_CASE(S32);
            KERNEL_GENERATE_CASE(S64);
            KERNEL_GENERATE_CASE(F32);
            KERNEL_GENERATE_CASE(F64);
            KERNEL_GENERATE_CASE(C32);
            KERNEL_GENERATE_CASE(C64);

        default: {
                char* typeStr;
                PS_TYPE_NAME(typeStr,xShifts->type.type);
                psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                        PS_ERRORTEXT_psImage_IMAGE_TYPE_UNSUPPORTED,
                        typeStr);
            }
        }
    }

    return result;
}

psImage* psImageConvolve(psImage* out, const psImage* in, const psKernel* kernel, psBool direct)
{
    if (in == NULL) {
        psFree(out);
        return NULL;
    }

    if (kernel == NULL) {
        psError(PS_ERR_BAD_PARAMETER_NULL, true,
                PS_ERRORTEXT_psImageConvolve_KERNEL_NULL);
        psFree(out);
        return NULL;
    }
    psS32 xMin = kernel->xMin;
    psS32 xMax = kernel->xMax;
    psS32 yMin = kernel->yMin;
    psS32 yMax = kernel->yMax;
    psKernelType** kData = kernel->kernel;

    // make the output image to the proper size and type
    psS32 numRows = in->numRows;
    psS32 numCols = in->numCols;



    if (direct) {
        // spatial convolution

        #define SPATIAL_CONVOLVE_CASE(TYPE) \
    case PS_TYPE_##TYPE: { \
            ps##TYPE** inData = in->data.TYPE; \
            out = psImageRecycle(out, numCols, numRows, PS_TYPE_##TYPE); \
            for (psS32 row=0;row<numRows;row++) { \
                ps##TYPE* outRow = out->data.TYPE[row]; \
                for (psS32 col=0;col<numCols;col++) { \
                    ps##TYPE pixel = 0.0; \
                    for (psS32 kRow = yMin; kRow < yMax; kRow++) { \
                        if (row-kRow >= 0 && row-kRow < numRows) { \
                            for (psS32 kCol = xMin; kCol < xMax; kCol++) { \
                                if (col-kCol >= 0 && col-kCol < numCols) { \
                                    pixel += kData[kRow][kCol] * inData[row-kRow][col-kCol]; \
                                } \
                            } \
                        } \
                    } \
                    outRow[col] = pixel; \
                } \
            } \
        } \
        break;

        switch (in->type.type) {
            SPATIAL_CONVOLVE_CASE(U8)
            SPATIAL_CONVOLVE_CASE(U16)
            SPATIAL_CONVOLVE_CASE(U32)
            SPATIAL_CONVOLVE_CASE(U64)
            SPATIAL_CONVOLVE_CASE(S8)
            SPATIAL_CONVOLVE_CASE(S16)
            SPATIAL_CONVOLVE_CASE(S32)
            SPATIAL_CONVOLVE_CASE(S64)
            SPATIAL_CONVOLVE_CASE(F32)
            SPATIAL_CONVOLVE_CASE(F64)
            SPATIAL_CONVOLVE_CASE(C32)
            SPATIAL_CONVOLVE_CASE(C64)

        default: {
                char* typeStr;
                PS_TYPE_NAME(typeStr,in->type.type);
                psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                        PS_ERRORTEXT_psImage_IMAGE_TYPE_UNSUPPORTED,
                        typeStr);
                psFree(out);
                return NULL;

            }
        }


    } else {
        // fourier convolution
        psS32 paddedCols = numCols+2*FOURIER_PADDING;
        psS32 paddedRows = numRows+2*FOURIER_PADDING;

        // check to see if kernel is smaller, otherwise padding it up will fail.
        psS32 kRows = kernel->image->numRows;
        psS32 kCols = kernel->image->numCols;
        if (kRows >= numRows || kCols >= numCols) {
            psError(PS_ERR_BAD_PARAMETER_SIZE, true,
                    PS_ERRORTEXT_psImageConvolve_KERNEL_TOO_LARGE,
                    kCols,kRows,
                    numCols, numRows);
            psFree(out);
            return NULL;
        }

        // pad the image
        psImage* paddedImage = psImageAlloc(paddedCols,paddedRows,in->type.type);
        psS32 elementSize = PSELEMTYPE_SIZEOF(in->type.type);

        // zero out padded area on top and bottom
        memset(paddedImage->data.U8[0],0,FOURIER_PADDING*paddedCols*elementSize);
        memset(paddedImage->data.U8[FOURIER_PADDING+numRows-1],0,FOURIER_PADDING*paddedCols*elementSize);

        // fill in the image-containing rows.
        psS32 sidePaddingSize = FOURIER_PADDING*elementSize;
        psS32 imageRowSize = numCols*elementSize;
        psU8* paddedData = paddedImage->data.U8[FOURIER_PADDING];
        for (psS32 row=0;row<numRows;row++) {
            // zero out padded area on left edge.
            memset(paddedData,0,sidePaddingSize);
            paddedData += sidePaddingSize;
            memcpy(paddedData,in->data.U8[row],imageRowSize);
            paddedData += imageRowSize;
            // zero out padded area on right edge.
            memset(paddedData,0,sidePaddingSize);
            paddedData += sidePaddingSize;
        }

        psImageWriteSection(paddedImage,0,0,0,NULL,0,"paddedImage.fits");

        // pad the kernel to the same size of paddedImage
        psImage* paddedKernel = psImageAlloc(paddedCols,paddedRows,PS_TYPE_KERNEL);
        memset(paddedKernel->data.U8[0],0,sizeof(psKernelType)*numCols*numRows); // zero-out image
        psS32 yMax = kernel->yMax;
        psS32 xMax = kernel->xMax;
        for (psS32 row = kernel->yMin; row <= yMax;row++) {
            psS32 padRow = row;
            if (padRow < 0) {
                padRow += paddedRows;
            }
            psKernelType* padData = paddedKernel->data.PS_TYPE_KERNEL_DATA[padRow];
            psKernelType* kernelRow = kernel->kernel[row];
            for (psS32 col = kernel->xMin; col <= xMax; col++) {
                if (col < 0) {
                    padData[col+paddedCols] = kernelRow[col];
                } else {
                    padData[col] = kernelRow[col];
                }
            }
        }

        psImageWriteSection(paddedKernel,0,0,0,NULL,0,"paddedKernel.fits");

        psImage* kernelFourier = psImageFFT(NULL, paddedKernel, PS_FFT_FORWARD);
        if (kernelFourier == NULL) {
            psError(PS_ERR_UNKNOWN, false,
                    PS_ERRORTEXT_psImageConvolve_KERNEL_FFT_FAILED);
            psFree(out);
            return NULL;
        }

        psImage* inFourier = psImageFFT(NULL, paddedImage, PS_FFT_FORWARD);
        if (inFourier == NULL) {
            psError(PS_ERR_UNKNOWN, false,
                    PS_ERRORTEXT_psImageConvolve_FFT_FAILED);
            psFree(out);
            return NULL;
        }

        // convolution in fourier domain is just a pixel-wise multiplication
        psBinaryOp(inFourier,inFourier,"*",kernelFourier);

        psImage* complexOut = psImageFFT(NULL, inFourier,
                                         PS_FFT_REVERSE);
        if (complexOut == NULL) {
            psError(PS_ERR_UNKNOWN, false,
                    PS_ERRORTEXT_psImageConvolve_FFT_FAILED);
            psFree(out);
            return NULL;
        }

        {
            psImage* tmp = psImageReal(NULL,complexOut);
            psImageWriteSection(tmp,0,0,0,NULL,0,"complexOut.fits");
            psFree(tmp);
        }

        // subset out the padded area now.
        psImage* complexOutSansPad = psImageSubset(complexOut,
                                     FOURIER_PADDING,FOURIER_PADDING,
                                     FOURIER_PADDING+numCols,FOURIER_PADDING+numRows);

        out = psImageRecycle(out,numCols,numRows,PS_TYPE_F32);
        float factor = 1.0f/numCols/numRows;
        for (psS32 row = 0; row < numRows; row++) {
            psF32* outRow = out->data.F32[row];
            psC32* resultRow = complexOutSansPad->data.C32[row];
            for (psS32 col = 0; col < numCols; col++) {
                outRow[col] = crealf(resultRow[col])*factor;
            }
        }

        psFree(complexOut); // frees complexOutSansPad, as it is a child.
        psFree(kernelFourier);
        psFree(inFourier);
        psFree(paddedImage);
        psFree(paddedKernel);

    }

    return out;
}
