/** @file  psVectorFFT.c
 *
 *  @brief Contains FFT transform related functions for psVector
 *
 *  @author Robert DeSonia, MHPCC
 *
 *  @version $Revision: 1.29 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2004-11-04 22:14:17 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 */

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

#include "psVectorFFT.h"
#include "psError.h"
#include "psMemory.h"
#include "psLogMsg.h"
#include "psImageExtraction.h"

#include "psDataManipErrors.h"

#define P_FFTW_PLAN_RIGOR FFTW_ESTIMATE

static psBool p_fftwWisdomImported = false;

psVector* psVectorFFT(psVector* out, const psVector* in, psFFTFlags direction)
{
    psU32 numElements;
    psElemType type;
    fftwf_plan plan;

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

    type = in->type.type;

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

    numElements = in->n;

    out = psVectorCopy(out, in, PS_TYPE_C32);
    out->n = numElements;

    if ((direction & PS_FFT_FORWARD) != 0) {
        plan = fftwf_plan_dft_1d(numElements,
                                 (fftwf_complex *) out->data.C32,
                                 (fftwf_complex *) out->data.C32, FFTW_FORWARD, P_FFTW_PLAN_RIGOR);
    } else if ((direction & PS_FFT_REVERSE) != 0) {
        plan = fftwf_plan_dft_1d(numElements,
                                 (fftwf_complex *) out->data.C32,
                                 (fftwf_complex *) out->data.C32, FFTW_BACKWARD, P_FFTW_PLAN_RIGOR);
    } else {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                PS_ERRORTEXT_psVectorFFT_DIRECTION_NOTSET);
        psFree(out);
        return NULL;
    }

    /* check if a plan exists now */
    if (plan == NULL) {
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psVectorFFT_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) != 0) {
        for (psS32 i = 0; i < numElements; i++) {
            out->data.F32[i] = out->data.C32[i];
        }
        out->type.type = PS_TYPE_F32;
        out->data.V = psRealloc(out->data.V,PSELEMTYPE_SIZEOF(PS_TYPE_F32)*out->nalloc);
    }

    return out;
}

psVector* psVectorReal(psVector* out, const psVector* in)
{
    psElemType type;
    psU32 numElements;

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

    type = in->type.type;
    numElements = in->n;

    /* if not a complex number, this is logically just a copy */
    if (!PS_IS_PSELEMTYPE_COMPLEX(type)) {
        // Warn user, as this is probably not expected
        psLogMsg(__func__, PS_LOG_WARN, "Real portion of a non-Complex type called called for. "
                 "Just a vector copy was performed.");
        out = psVectorRecycle(out, numElements, type);
        out->n = numElements;
        memcpy(out->data.V, in->data.V, numElements * PSELEMTYPE_SIZEOF(type));
        return out;
    }

    if (type == PS_TYPE_C32) {
        psF32* outVec;
        psC32* inVec = in->data.C32;

        out = psVectorRecycle(out, numElements, PS_TYPE_F32);
        out->n = numElements;
        outVec = out->data.F32;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = crealf(inVec[i]);
        }
    } else if (type == PS_TYPE_C64) {
        psF64* outVec;
        psC64* inVec = in->data.C64;

        out = psVectorRecycle(out, numElements, PS_TYPE_F64);
        out->n = numElements;
        outVec = out->data.F64;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = creal(inVec[i]);
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psVectorFFT_TYPE_UNSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;
}

psVector* psVectorImaginary(psVector* out, const psVector* in)
{
    psElemType type;
    psU32 numElements;

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

    type = in->type.type;
    numElements = in->n;

    /* if not a complex number, this is logically just zeroed image of same size */
    if (!PS_IS_PSELEMTYPE_COMPLEX(type)) {
        // Warn user, as this is probably not expected
        psLogMsg(__func__, PS_LOG_WARN, "Imaginary portion of a non-Complex type called for. "
                 "A zeroed vector was returned.");
        out = psVectorRecycle(out, numElements, type);
        out->n = numElements;
        memset(out->data.V, 0, PSELEMTYPE_SIZEOF(type) * numElements);
        return out;
    }

    if (type == PS_TYPE_C32) {
        psF32* outVec;
        psC32* inVec = in->data.C32;

        out = psVectorRecycle(out, numElements, PS_TYPE_F32);
        out->n = numElements;
        outVec = out->data.F32;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = cimagf(inVec[i]);
        }
    } else if (type == PS_TYPE_C64) {
        psF64* outVec;
        psC64* inVec = in->data.C64;

        out = psVectorRecycle(out, numElements, PS_TYPE_F64);
        out->n = numElements;
        outVec = out->data.F64;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = cimag(inVec[i]);
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psVectorFFT_TYPE_UNSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;
}

psVector* psVectorComplex(psVector* out, const psVector* real, const psVector* imag)
{
    psElemType type;
    psU32 numElements;

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

    type = real->type.type;
    if (real->n < imag->n) {
        numElements = real->n;
    } else {
        numElements = imag->n;
    }

    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_psVectorFFT_REAL_IMAG_TYPE_MISMATCH,
                typeStrReal,typeStrImag);
        psFree(out);
        return NULL;
    }

    if (type == PS_TYPE_F32) {
        psC32* outVec;
        psF32* realVec = real->data.F32;
        psF32* imagVec = imag->data.F32;

        out = psVectorRecycle(out, numElements, PS_TYPE_C32);
        out->n = numElements;
        outVec = out->data.C32;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = realVec[i] + I * imagVec[i];
        }
    } else if (type == PS_TYPE_F64) {
        psC64* outVec;
        psF64* realVec = real->data.F64;
        psF64* imagVec = imag->data.F64;

        out = psVectorRecycle(out, numElements, PS_TYPE_C64);
        out->n = numElements;
        outVec = out->data.C64;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = realVec[i] + I * imagVec[i];
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psVectorFFT_NONREAL_NOTSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;
}

psVector* psVectorConjugate(psVector* out, const psVector* in)
{
    psElemType type;
    psU32 numElements;

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

    type = in->type.type;
    numElements = in->n;

    /* if not a complex number, this is logically just a image copy */
    if (!PS_IS_PSELEMTYPE_COMPLEX(type)) {
        // Warn user, as this is probably not expected
        psLogMsg(__func__, PS_LOG_WARN, "Complex Conjugate of a non-Complex type called for. "
                 "Vector copy was performed instead.");

        out = psVectorRecycle(out, numElements, type);
        out->n = numElements;
        memcpy(out->data.V, in->data.V, PSELEMTYPE_SIZEOF(type) * numElements);
        return out;
    }

    if (type == PS_TYPE_C32) {
        psC32* outVec;
        psC32* inVec = in->data.C32;

        out = psVectorRecycle(out, numElements, PS_TYPE_C32);
        out->n = numElements;
        outVec = out->data.C32;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = crealf(inVec[i]) - I * cimagf(inVec[i]);
        }
    } else if (type == PS_TYPE_C64) {
        psC64* outVec;
        psC64* inVec = in->data.C64;

        out = psVectorRecycle(out, numElements, PS_TYPE_C64);
        out->n = numElements;
        outVec = out->data.C64;

        for (psU32 i = 0; i < numElements; i++) {
            outVec[i] = creal(inVec[i]) - I * cimag(inVec[i]);
        }
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psVectorFFT_NONCOMPLEX_NOTSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;
}

psVector* psVectorPowerSpectrum(psVector* out, const psVector* in)
{
    psElemType type;
    psU32 outNumElements;
    psU32 inNumElements;
    psU32 inHalfNumElements;
    psU32 inNumElementsSquared;

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

    type = in->type.type;
    inNumElements = in->n;
    inNumElementsSquared = inNumElements * inNumElements;
    inHalfNumElements = inNumElements / 2;
    outNumElements = inHalfNumElements + 1;

    if (type == PS_TYPE_C32) {
        psF32* outVec;
        psC32* inVec = in->data.C32;
        psF32 inAbs1;
        psF32 inAbs2;

        out = psVectorRecycle(out, outNumElements, PS_TYPE_F32);
        out->n = outNumElements;
        outVec = out->data.F32;

        // from ADD: P_0 = |C_0|^2/N^2
        inAbs1 = cabsf(inVec[0]);
        outVec[0] = inAbs1 * inAbs1 / inNumElementsSquared;

        // from ADD: P_j = (|C_j|^2+|C_N-j|^2)/N^2, where j = 1,2,...,(N/2-1)
        for (psU32 i = 1; i < inHalfNumElements; i++) {
            inAbs1 = cabsf(inVec[i]);
            inAbs2 = cabsf(inVec[inNumElements - i]);
            outVec[i] = (inAbs1 * inAbs1 + inAbs2 * inAbs2) / inNumElementsSquared;
        }

        // from ADD: P_N/2 = |C_N/2|^2/N^2
        inAbs1 = cabsf(inVec[inHalfNumElements]);
        outVec[inHalfNumElements] = inAbs1 * inAbs1 / inNumElementsSquared;
    } else if (type == PS_TYPE_C64) {
        psF64* outVec;
        psC64* inVec = in->data.C64;
        psF64 inAbs1;
        psF64 inAbs2;

        out = psVectorRecycle(out, outNumElements, PS_TYPE_F64);
        out->n = outNumElements;
        outVec = out->data.F64;

        // from ADD: P_0 = |C_0|^2/N^2
        inAbs1 = cabs(inVec[0]);
        outVec[0] = inAbs1 * inAbs1 / inNumElementsSquared;

        // from ADD: P_j = (|C_j|^2+|C_N-j|^2)/N^2, where j = 1,2,...,(N/2-1)
        for (psU32 i = 1; i < inHalfNumElements; i++) {
            inAbs1 = cabs(inVec[i]);
            inAbs2 = cabs(inVec[inNumElements - i]);
            outVec[i] = (inAbs1 * inAbs1 + inAbs2 * inAbs2) / inNumElementsSquared;
        }

        // from ADD: P_N/2 = |C_N/2|^2/N^2
        inAbs1 = cabs(inVec[inHalfNumElements]);
        outVec[inHalfNumElements] = inAbs1 * inAbs1 / inNumElementsSquared;
    } else {
        char* typeStr;
        PS_TYPE_NAME(typeStr,type);
        psError(PS_ERR_BAD_PARAMETER_TYPE, true,
                PS_ERRORTEXT_psVectorFFT_NONCOMPLEX_NOTSUPPORTED,
                typeStr);
        psFree(out);
        return NULL;
    }

    return out;

}
