/// @file  psImageFFT.c
///
/// @brief Contains FFT transform related functions for psImage.
///
/// @author Paul Price, IfA
/// @author Robert DeSonia, MHPCC
///
/// @version $Revision: 1.23 $ $Name: not supported by cvs2svn $
/// @date $Date: 2007-02-09 00:22:55 $
///
/// Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
///

#ifdef HAVE_CONFIG_H
# include "config.h"
#endif

#include <unistd.h>
#include <string.h>
#include <complex.h>
#include <fftw3.h>

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


#define FFTW_PLAN_RIGOR FFTW_ESTIMATE   // How rigorous the FFTW planning is

static psBool fftwWisdomImported = false; // Has the FFTW wisdom been imported yet?

bool psImageForwardFFT(psImage **real, psImage **imag, const psImage *in)
{
    PS_ASSERT_IMAGE_NON_NULL(in, false);
    PS_ASSERT_IMAGE_TYPE(in, PS_TYPE_F32, false);
    PS_ASSERT_PTR_NON_NULL(real, false);
    PS_ASSERT_PTR_NON_NULL(imag, false);

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

    int numCols = in->numCols;          // Number of columns
    int numRows = in->numRows;          // Number of rows

    psF32 *input;                       // Input data, for FFTW
    if (!in->parent) {
        // No parent --- can just use the data buffer
        input = psMemIncrRefCounter(in->p_rawDataBuffer);
    } else {
        // Need to copy the data
        input = psAlloc(numCols * numRows);
        for (int y = 0; y < numRows; y++) {
            memcpy(&input[y * numRows], in->data.F32[y], numCols);
        }
    }

    // Do the FFT
    fftwf_complex *out = fftwf_malloc((numCols/2 + 1) * numRows * sizeof(fftwf_complex)); // Output data
    fftwf_plan plan = fftwf_plan_dft_r2c_2d(numRows, numCols, input, out, FFTW_PLAN_RIGOR);
    fftwf_execute(plan);
    fftwf_destroy_plan(plan);
    psFree(input);

    // Pull the real and imaginary parts out
    numCols = numCols/2 + 1;            // x dimension is halved by FFTW
    *real = psImageRecycle(*real, numCols, numRows, PS_TYPE_F32);
    *imag = psImageRecycle(*imag, numCols, numRows, PS_TYPE_F32);
    for (int y = 0, index = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; index++, x++) {
#if !defined(FFTW_NO_Complex) && defined(_Complex_I) && defined(complex) && defined(I)
            // C99 complex support
            (*real)->data.F32[y][x] = creal(out[index]);
            (*imag)->data.F32[y][x] = cimag(out[index]);
#else
            // FFTW's backup complex support
            (*real)->data.F32[y][x] = out[index][0];
            (*imag)->data.F32[y][x] = out[index][1];
#endif
        }
    }
    fftwf_free(out);

    return true;
}

bool psImageBackwardFFT(psImage **out, const psImage *real, const psImage *imag, int origCols)
{
    PS_ASSERT_IMAGE_NON_NULL(real, false);
    PS_ASSERT_IMAGE_TYPE(real, PS_TYPE_F32, false);
    PS_ASSERT_IMAGE_NON_NULL(imag, false);
    PS_ASSERT_IMAGE_TYPE(imag, PS_TYPE_F32, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(real, imag, false);
    PS_ASSERT_PTR_NON_NULL(out, false);

    int numCols = real->numCols;        // Number of columns
    int numRows = real->numRows;        // Number of rows

    // Because of the way FFT r2c and c2r work, need the number of columns in the target to be:
    // 2 * numCols = 2*(numCols/2 + 1) = origCols % 2 ? origCols + 1 : origCols + 2
    if (2 * numCols != ((origCols % 2) ? origCols + 1 : origCols + 2)) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, false,
                "Number of columns in FFT-ed images (%d) should be origCols (%d) / 2 + 1 = %d",
                numCols, origCols, origCols/2 + 1);
        return false;
    }

    if (*out && (*out)->parent) {
        // It has a parent, so we can't write directly into the buffer.
        // It had better be the correct size and type, because we don't want to resize a child image
        PS_ASSERT_IMAGE_SIZE(*out, origCols, numRows, false);
        PS_ASSERT_IMAGE_TYPE(*out, PS_TYPE_F32, false);
    }


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

    // Stuff the real and imaginary parts in
    psF32 *target = psAlloc(2 * numCols * numRows * PSELEMTYPE_SIZEOF(PS_TYPE_F32)); // Target for FFTW
    fftwf_complex *in = fftwf_malloc(numCols * numRows * sizeof(fftwf_complex)); // Input data

    for (int y = 0, index = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; index++, x++) {
#if !defined(FFTW_NO_Complex) && defined(_Complex_I) && defined(complex) && defined(I)
            // C99 complex support
            in[index] = real->data.F32[y][x] + imag->data.F32[y][x] * I;
#else
            // FFTW's backup complex support
            in[index][0] = real->data.F32[y][x];
            in[index][1] = imag->data.F32[y][x];
#endif
        }
    }

    // Do the FFT
    fftwf_plan plan = fftwf_plan_dft_c2r_2d(numRows, origCols, in, target, FFTW_PLAN_RIGOR);
    fftwf_execute(plan);
    fftwf_destroy_plan(plan);
    fftwf_free(in);

    // Copy the target pixels into the output
    if (!(*out) || !(*out)->parent) {
        *out = psImageRecycle(*out, origCols, numRows, PS_TYPE_F32);
        memcpy((*out)->p_rawDataBuffer, target, numRows * origCols * PSELEMTYPE_SIZEOF(PS_TYPE_F32));
    } else {
        for (int y = 0, index = 0; y < numRows; y++, index += origCols) {
            memcpy(&(*out)->data.F32[y][0], &target[index], origCols * PSELEMTYPE_SIZEOF(PS_TYPE_F32));
        }
    }

    psFree(target);

    return true;
}

psImage* psImagePowerSpectrum(psImage *out, const psImage *in)
{
    PS_ASSERT_IMAGE_NON_NULL(in, NULL);
    PS_ASSERT_IMAGE_TYPE(in, PS_TYPE_F32, NULL);

    psImage *real = NULL;              // Real component of FFT
    psImage *imag = NULL;              // Imaginary component of FFT

    if (!psImageForwardFFT(&real, &imag, in)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to perform forward FFT.");
        return NULL;
    }

    int numCols = real->numCols;        // Number of columns
    int numRows = real->numRows;        // Number of rows

    float norm = 1.0 / PS_SQR(in->numCols) / PS_SQR(in->numRows);
    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            // Power spectrum is the square of the complex modulus
            real->data.F32[y][x] = norm * (PS_SQR(real->data.F32[y][x]) + PS_SQR(imag->data.F32[y][x]));
        }
    }
    psFree(imag);

    return real;
}



bool psImageComplexMultiply(psImage **outReal, psImage **outImag,
                            const psImage *in1Real, const psImage *in1Imag,
                            const psImage *in2Real, const psImage *in2Imag)
{
    PS_ASSERT_IMAGE_NON_NULL(in1Real, false);
    PS_ASSERT_IMAGE_NON_NULL(in1Imag, false);
    PS_ASSERT_IMAGE_NON_NULL(in2Real, false);
    PS_ASSERT_IMAGE_NON_NULL(in2Imag, false);
    PS_ASSERT_IMAGE_TYPE(in1Real, PS_TYPE_F32, false);
    PS_ASSERT_IMAGE_TYPE(in1Imag, PS_TYPE_F32, false);
    PS_ASSERT_IMAGE_TYPE(in2Real, PS_TYPE_F32, false);
    PS_ASSERT_IMAGE_TYPE(in2Imag, PS_TYPE_F32, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(in1Imag, in1Real, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(in2Real, in1Real, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(in2Imag, in1Real, false);
    PS_ASSERT_PTR_NON_NULL(outReal, false);
    PS_ASSERT_PTR_NON_NULL(outImag, false);
    if (*outReal) {
        PS_ASSERT_IMAGE_NON_NULL(*outReal, false);
        PS_ASSERT_IMAGE_NON_NULL(*outImag, false);
    } else if (*outImag) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                "If the output real part is provided, the output imaginary part must also be provided.");
        return false;
    }

    int numRows = in1Real->numRows;     // Number of rows
    int numCols = in1Real->numCols;     // Number of columns

    // Need to worry if the outputs are children
    psImage *targetReal, *targetImag;   // Target real and imaginary parts

    if ((*outReal)->parent) {
        // It had better be the correct size and type, because we don't want to resize a child image
        PS_ASSERT_IMAGE_TYPE(*outReal, PS_TYPE_F32, false);
        PS_ASSERT_IMAGE_SIZE(*outReal, numCols, numRows, false);
        targetReal = psImageAlloc(numCols, numRows, PS_TYPE_F32);
    } else {
        *outReal = psImageRecycle(*outReal, numCols, numRows, PS_TYPE_F32);
        targetReal = psMemIncrRefCounter(*outReal);
    }
    if ((*outImag)->parent) {
        // It had better be the correct size and type, because we don't want to resize a child image
        if ((*outReal)->numCols != numCols || (*outReal)->numRows != numRows ||
            (*outImag)->type.type != PS_TYPE_F32) {
            // Plug potential memory leak --- need to free targetReal if there's a problem
            psFree(targetReal);
        }
        PS_ASSERT_IMAGE_TYPE(*outReal, PS_TYPE_F32, false);
        PS_ASSERT_IMAGE_SIZE(*outReal, numCols, numRows, false);
        targetImag = psImageAlloc(numCols, numRows, PS_TYPE_F32);
    } else {
        *outImag = psImageRecycle(*outImag, numCols, numRows, PS_TYPE_F32);
        targetImag = psMemIncrRefCounter(*outImag);
    }

    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            // (a + bi) * (c + di) = (ac - bd) + (bc + ad)i
            float real = in1Real->data.F32[y][x] * in2Real->data.F32[y][x] -
                in1Imag->data.F32[y][x] * in2Imag->data.F32[y][x];
            float imag = in1Imag->data.F32[y][x] * in2Real->data.F32[y][x] +
                in1Real->data.F32[y][x] * in2Imag->data.F32[y][x];
            targetReal->data.F32[y][x] = real;
            targetImag->data.F32[y][x] = imag;
        }
    }

    if ((*outReal)->parent) {
        *outReal = psImageCopy(*outReal, targetReal, PS_TYPE_F32);
    }
    if ((*outImag)->parent) {
        *outImag = psImageCopy(*outImag, targetImag, PS_TYPE_F32);
    }

    psFree(targetReal);
    psFree(targetImag);

    return true;
}
