/*
 * Realfftf
 * Originally (C) Copyright 1993 Philib Van Baren
 * Cleaned up and refactored to be re-entrant by
 *     Matti 'ccr' Hamalainen <ccr@tnsp.org>, 2013
 *
 * Note: Output is BIT-REVERSED! so you must use the BitReversed to
 *       get legible output, (i.e. Real_i = buffer[ BitReversed[i] ]
 *                                 Imag_i = buffer[ BitReversed[i]+1 ] )
 *       Input is in normal order.
 */
#include "dmfft.h"
#include <math.h>


#define DM_PW(a, b, c) (sqrt((a) * (a) + (b) * (b)) * (c))


int dmInitializeFFT(DMFFTContext * ctx, const int fftlen)
{
    if (ctx == NULL)
        return DMERR_NULLPTR;

    if (fftlen <= 16)
        return DMERR_INVALID_ARGS;

    // Allocate necessary tables
    ctx->npoints = fftlen / 2;
    ctx->sinTable = (DMFFTType *) dmMalloc(fftlen * sizeof(DMFFTType));
    if (ctx->sinTable == NULL)
        return DMERR_MALLOC;

    // Bitreversion buffer
    ctx->breversed = (int *) dmMalloc(ctx->npoints * sizeof(int));
    if (ctx->breversed == NULL)
    {
        dmFree(ctx->sinTable);
        return DMERR_MALLOC;
    }

    // Calculate data
    for (int i = 0; i < ctx->npoints; i++)
    {
        int mask, tmp;
        for (tmp = 0, mask = ctx->npoints / 2; mask > 0; mask >>= 1)
        {
            tmp = (tmp >> 1) + ((i & mask) ? ctx->npoints : 0);
        }

        ctx->breversed[i] = tmp;
    }

    for (int i = 0; i < ctx->npoints; i++)
    {
        ctx->sinTable[ctx->breversed[i]] = -sin((2.0f * DM_PI * i) / fftlen);
        ctx->sinTable[ctx->breversed[i] + 1] =
            -cos((2.0f * DM_PI * i) / fftlen);
    }

    return DMERR_OK;
}


void dmEndFFT(DMFFTContext * ctx)
{
    if (ctx != NULL)
    {
        dmFree(ctx->breversed);
        dmFree(ctx->sinTable);
        dmMemset(ctx, 0, sizeof(DMFFTContext));
    }
}


int dmRealFFT(DMFFTContext * ctx, DMFFTType * buffer)
{
    if (ctx == NULL || buffer == NULL)
        return DMERR_NULLPTR;

    const int *br1 = ctx->breversed + 1;
    const int *br2 = ctx->breversed + ctx->npoints - 1;
    int rounds = ctx->npoints / 2;
    DMFFTType *endptr1 = buffer + ctx->npoints * 2;

    while (rounds > 0)
    {
        DMFFTType *sptr = ctx->sinTable;
        DMFFTType *A = buffer,
                  *B = buffer + rounds * 2;

        while (A < endptr1)
        {
            DMFFTType qsin = *sptr;
            DMFFTType qcos = *(sptr + 1);
            DMFFTType *endptr2 = B;

            while (A < endptr2)
            {
                DMFFTType v1 = (*B) * qcos + (*(B + 1)) * qsin;
                DMFFTType v2 = (*B) * qsin - (*(B + 1)) * qcos;
                *B = (*A + v1) * 0.5;
                *(A++) = *(B++) - v1;
                *B = (*A - v2) * 0.5;
                *(A++) = *(B++) + v2;
            }
            A = B;
            B += rounds * 2;
            sptr += 2;
        }

        rounds >>= 1;
    }

    // Massage output to get the output for a real input sequence.
    while (br1 <= br2)
    {
        const DMFFTType vsin = ctx->sinTable[*br1];
        const DMFFTType vcos = ctx->sinTable[*br1 + 1];
        DMFFTType *A = buffer + *br1,
                  *B = buffer + *br2,
                  HRplus, HRminus,
                  HIplus, HIminus,
                  tmp1, tmp2;

        HRminus  = *A - *B;
        HRplus   = HRminus + (*B * 2);

        HIminus  = *(A + 1) - *(B + 1);
        HIplus   = HIminus + (*(B + 1) * 2);

        tmp1     = (vsin * HRminus) - (vcos * HIplus);
        tmp2     = (vcos * HRminus) + (vsin * HIplus);

        *A       = (HRplus + tmp1) * 0.5f;
        *B       = *A - tmp1;

        *(A + 1) = (HIminus + tmp2) * 0.5f;
        *(B + 1) = (*(A + 1)) - HIminus;

        br1++;
        br2--;
    }

    // Handle DC bin separately
    buffer[0] += buffer[1];
    buffer[1] = 0;

    return DMERR_OK;
}


int dmConvertFFTtoFreqDomain(DMFFTContext *ctx, DMFFTType *buffer,
    DMFFTType *real, DMFFTType *imag)
{
    if (ctx == NULL || buffer == NULL ||
        real == NULL || imag == NULL)
        return DMERR_NULLPTR;

    // Convert from bitreversed form to normal frequency domain
    const int *breversed = ctx->breversed;

    for (int i = 1; i < ctx->npoints; i++)
    {
        real[i] = buffer[breversed[i]  ];
        imag[i] = buffer[breversed[i]+1];
    }

    // Make the ends meet
    real[0]            = buffer[0];
    imag[0]            = 0;
    real[ctx->npoints] = buffer[1];
    imag[ctx->npoints] = 0;

    return DMERR_OK;
}


int dmConvertFFTtoFreqAndPower(DMFFTContext *ctx, DMFFTType *buffer,
    DMFFTType *real, DMFFTType *imag, DMFFTType *power, const DMFFTType scale)
{
    if (ctx == NULL || buffer == NULL ||
        real == NULL || imag == NULL || power == NULL)
        return DMERR_NULLPTR;

    // Convert from bitreversed form to normal frequency domain
    const int *breversed = ctx->breversed;

    for (int i = 1; i < ctx->npoints; i++)
    {
        DMFFTType fre = real[i] = buffer[breversed[i]  ],
                  fim = imag[i] = buffer[breversed[i]+1];

        power[i] = DM_PW(fre, fim, scale);
    }

    // Make the ends meet
    real[0]             = buffer[0];
    imag[0]             = 0;
    power[0]            = DM_PW(buffer[0], 0, scale);

    real[ctx->npoints]  = buffer[1];
    imag[ctx->npoints]  = 0;
    power[ctx->npoints] = DM_PW(buffer[1], 0, scale);

    return DMERR_OK;
}


int dmConvertFFTtoPowerAndSum(DMFFTContext *ctx, DMFFTType *buffer,
    DMFFTType *power, const DMFFTType pscale, DMFFTType *sum, const DMFFTType sscale)
{
    if (ctx == NULL || buffer == NULL || power == NULL || sum == NULL)
        return DMERR_NULLPTR;

    // Convert from bitreversed form to normal frequency domain
    const int *breversed = ctx->breversed;
    DMFFTType psum = 0;

    for (int i = 1; i < ctx->npoints; i++)
    {
        DMFFTType fre = buffer[breversed[i]  ],
                  fim = buffer[breversed[i]+1],
                  fpw = sqrt(fre * fre + fim * fim);

        power[i]  = fpw * pscale;
        psum     += fpw * sscale;
    }

    // Make the ends meet
    power[0]            = DM_PW(buffer[0], 0, pscale);
    power[ctx->npoints] = DM_PW(buffer[1], 0, pscale);

    *sum = psum;

    return DMERR_OK;
}


int dmConvertFFTtoTimeDomain(DMFFTContext *ctx, DMFFTType *buffer, DMFFTType *tdom)
{
    int i, n;
    if (ctx == NULL || buffer == NULL || tdom == NULL)
        return DMERR_NULLPTR;

    for (n = i = 0; i < ctx->npoints; i++, n += 2)
    {
        tdom[n  ] = buffer[ctx->breversed[i]  ];
        tdom[n+1] = buffer[ctx->breversed[i]+1];
    }

    return DMERR_OK;
}