/** @file  psImageFFT.c
 *
 *  @brief Contains FFT transform related functions for psImage.
 *
 *  @author Robert DeSonia, MHPCC
 *
 *  @version $Revision: 1.17 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2005-07-12 19:30:18 $
 *
 *  Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
 */
#include <unistd.h>
#include <string.h>
#include <complex.h>
#include <fftw3.h>

#include "psImageFFT.h"
#include "psError.h"
#include "psMemory.h"
#include "psLogMsg.h"
#include "psImageStructManip.h"

#include "psErrorText.h"

#define PS_FFTW_PLAN_RIGOR FFTW_ESTIMATE

static psBool p_fftwWisdomImported = false;

psImage* psImageFFT(psImage* out, const psImage* image, psFFTFlags direction)
{
    psU32 numCols;
    psU32 numRows;
    psElemType type;
    fftwf_plan plan;

    /* got good image data? */
    if (image == NULL) {
        psFree(out);
        return NULL;
    }

    if ( ((direction & PS_FFT_FORWARD) != 0) ) {
        if ((direction & PS_FFT_REVERSE) != 0) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                    PS_ERRORTEXT_psImageFFT_FORWARD_REVERSE);
            psFree(out);
            return NULL;
        }
        if ((direction & PS_FFT_REAL_RESULT) != 0) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                    PS_ERRORTEXT_psImageFFT_REAL_FORWARD_NOTSUPPORTED);
            psFree(out);
            return NULL;
        }
    } else if ((direction & PS_FFT_REVERSE) == 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                PS_ERRORTEXT_psImageFFT_NO_DIRECTION_OPTION);
        psFree(out);
        return NULL;
    }

    type = image->type.type;

    /* make sure the system-level wisdom information is imported. */
    if (!p_fftwWisdomImported) {
        fftwf_import_system_wisdom();
        p_fftwWisdomImported = true;
    }

    numRows = image->numRows;
    numCols = image->numCols;

    // n.b. FFTW can perform a in-place transform at the same rate or faster than out-of-place.
    psS32 sign = ((direction & PS_FFT_FORWARD) != 0) ? FFTW_FORWARD : FFTW_BACKWARD;

    fftwf_complex* outBuffer = fftwf_malloc(numRows*numCols*sizeof(fftwf_complex));
    p_psImageCopyToRawBuffer(outBuffer, image, PS_TYPE_C32);

    plan = fftwf_plan_dft_2d(numRows, numCols,
                             outBuffer,
                             outBuffer,
                             sign,
                             PS_FFTW_PLAN_RIGOR);

    /* check if a plan exists now -- if not, it is a real problem at this point */
    if (plan == NULL) {
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImageFFT_FFTW_PLAN_NULL);
        psFree(out);
        return NULL;
    }

    /* finally, call FFTW with the plan made above */
    fftwf_execute(plan);

    fftwf_destroy_plan(plan);

    if (direction & PS_FFT_REAL_RESULT) {
        // n.b., we do this instead of using fftwf_plan_dft_c2r because that
        // plan requires a half-image, which would require a image reordering
        // that is not as simple as performing a normal complex transform and
        // then taking the real part of the result.  If performance here
        // becomes an issue, the use of fftwf_plan_dft_r2c should be considered
        // as well as fftwf_plan_dft_c2r.
        out = psImageRecycle(out,numCols,numRows,PS_TYPE_F32);
        int index = 0;
        for (int row=0; row < numRows; row++) {
            psF32* outRow = out->data.F32[row];
            for (int col=0; col < numCols; col++) {
                outRow[col] = crealf(outBuffer[index++]); // take just the real part
            }
        }
    } else {
        out = psImageRecycle(out,numCols,numRows,PS_TYPE_C32);
        int index = 0;
        for (int row=0; row < numRows; row++) {
            psC32* outRow = out->data.C32[row];
            for (int col=0; col < numCols; col++) {
                outRow[col] = outBuffer[index++]; // take just the real part
            }
        }
        //        memcpy(out->rawDataBuffer, outBuffer, numRows*numCols*sizeof(fftwf_complex));
    }

    fftwf_free(outBuffer);

    return out;

}

psImage* psImageReal(psImage* out, const psImage* in)
{
    psElemType type;
    psU32 numCols;
    psU32 numRows;

    if (in == NULL) {
        psFree(out);
        return NULL;
    }

    type = in->type.type;
    numCols = in->numCols;
    numRows = in->numRows;

    /* if not a complex number, this is logically just a copy then */
    if (!PS_IS_PSELEMTYPE_COMPLEX(type)) {
        return psImageCopy(out, in, type);
    }

    if (type == PS_TYPE_C32) {
        psF32* outRow;
        psC32* inRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_F32);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.F32[row];
            inRow = in->data.C32[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = crealf(inRow[col]);
            }
        }
    } else if (type == PS_TYPE_C64) {
        psF64* outRow;
        psC64* inRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_F64);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.F64[row];
            inRow = in->data.C64[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = creal(inRow[col]);
            }
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImage_IMAGE_TYPE_UNSUPPORTED,
                typeStr);

        psFree(out);
        return NULL;
    }

    return out;
}

psImage* psImageImaginary(psImage* out, const psImage* in)
{
    psElemType type;
    psU32 numCols;
    psU32 numRows;

    if (in == NULL) {
        psFree(out);
        return NULL;
    }

    type = in->type.type;
    numCols = in->numCols;
    numRows = in->numRows;

    /* if not a complex image type, this is logically just zeroed image of same size */
    if (!PS_IS_PSELEMTYPE_COMPLEX(type)) {
        out = psImageRecycle(out, numCols, numRows, type);
        memset(out->data.V[0], 0, PSELEMTYPE_SIZEOF(type) * numCols * numRows);
        return out;
    }

    if (type == PS_TYPE_C32) {
        psF32* outRow;
        psC32* inRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_F32);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.F32[row];
            inRow = in->data.C32[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = cimagf(inRow[col]);
            }
        }
    } else if (type == PS_TYPE_C64) {
        psF64* outRow;
        psC64* inRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_F64);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.F64[row];
            inRow = in->data.C64[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = cimag(inRow[col]);
            }
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImage_IMAGE_TYPE_UNSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;
}

psImage* psImageComplex(psImage* out, const psImage* real, const psImage* imag)
{
    psElemType type;
    psU32 numCols;
    psU32 numRows;

    if (real == NULL || imag == NULL) {
        psFree(out);
        return NULL;
    }

    type = real->type.type;
    numCols = real->numCols;
    numRows = real->numRows;

    if (imag->type.type != type) {
        char* typeStrReal;
        char* typeStrImag;
        PS_TYPE_NAME(typeStrReal,type);
        PS_TYPE_NAME(typeStrImag,imag->type.type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImageFFT_REAL_IMAG_TYPE_MISMATCH,
                typeStrReal,typeStrImag);
        psFree(out);
        return NULL;
    }

    if (imag->numCols != numCols || imag->numRows != numRows) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                PS_ERRORTEXT_psImageFFT_REAL_IMAG_SIZE_MISMATCH,
                numCols, numRows, imag->numCols, imag->numRows);
        psFree(out);
        return NULL;
    }

    if (type == PS_TYPE_F32) {
        psC32* outRow;
        psF32* realRow;
        psF32* imagRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_C32);

        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.C32[row];
            realRow = real->data.F32[row];
            imagRow = imag->data.F32[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = realRow[col] + I * imagRow[col];
            }
        }
    } else if (type == PS_TYPE_F64) {
        psC64* outRow;
        psF64* realRow;
        psF64* imagRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_C64);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.C64[row];
            realRow = real->data.F64[row];
            imagRow = imag->data.F64[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = realRow[col] + I * imagRow[col];
            }
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImageFFT_NONREAL_NOTSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }


    return out;
}

psImage* psImageConjugate(psImage* out, const psImage* in)
{
    psElemType type;
    psU32 numCols;
    psU32 numRows;

    if (in == NULL) {
        psFree(out);
        return NULL;
    }

    type = in->type.type;
    numCols = in->numCols;
    numRows = in->numRows;

    /* if not a complex image, this is logically just a image copy */
    if (!PS_IS_PSELEMTYPE_COMPLEX(type)) {
        return psImageCopy(out, in, type);
    }

    if (type == PS_TYPE_C32) {
        psC32* outRow;
        psC32* inRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_C32);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.C32[row];
            inRow = in->data.C32[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = crealf(inRow[col]) - I * cimagf(inRow[col]);
            }
        }
    } else if (type == PS_TYPE_C64) {
        psC64* outRow;
        psC64* inRow;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_C64);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.C64[row];
            inRow = in->data.C64[row];

            for (psU32 col = 0; col < numCols; col++) {
                outRow[col] = creal(inRow[col]) - I * cimag(inRow[col]);
            }
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImageFFT_NONCOMPLEX_NOTSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;
}

psImage* psImagePowerSpectrum(psImage* out, const psImage* in)
{
    psElemType type;
    psU32 numCols;
    psU32 numRows;

    if (in == NULL) {
        psFree(out);
        return NULL;
    }

    type = in->type.type;
    numCols = in->numCols;
    numRows = in->numRows;

    if (type == PS_TYPE_C32) {
        psF32* outRow;
        psC32* inRow;
        psF32 real;
        psF32 imag;
        psF32 fNumCols = numCols;
        psF32 fNumRows = numRows;
        psF32 numElementsSquared = fNumCols * fNumCols * fNumRows * fNumRows;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_F32);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.F32[row];
            inRow = in->data.C32[row];

            for (psU32 col = 0; col < numCols; col++) {
                real = crealf(inRow[col]);
                imag = cimagf(inRow[col]);
                outRow[col] = (real * real + imag * imag) / numElementsSquared;
            }
        }
    } else if (type == PS_TYPE_C64) {
        psF64* outRow;
        psC64* inRow;
        psF64 real;
        psF64 imag;
        psF64 numElementsSquared = numCols * numCols * numRows * numRows;

        out = psImageRecycle(out, numCols, numRows, PS_TYPE_F64);
        for (psU32 row = 0; row < numRows; row++) {
            outRow = out->data.F64[row];
            inRow = in->data.C64[row];

            for (psU32 col = 0; col < numCols; col++) {
                real = creal(inRow[col]);
                imag = cimag(inRow[col]);
                outRow[col] = (real * real + imag * imag) / numElementsSquared;
            }
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psImageFFT_NONCOMPLEX_NOTSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;

}
