/*
 * PFA FFT for transforms of size 2^k * 15.
 *
 * Parts taken from fft.c from libavcodec:
 * Copyright (c) 2008 Loren Merritt
 * Copyright (c) 2002 Fabrice Bellard
 * Partly based on libdjbfft by D. J. Bernstein
 *
 * To compile the test program, run:
 *   clang -Wall -O2 -fno-math-errno -fno-trapping-math -DTEST -o pfa_test pfa.c -lfftw3f
 * (or equivalent w/the C compiler of your choice)
 *
 */
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdio.h>
#include <malloc.h>
#ifdef TEST
#include <fftw3.h>
#endif

/* FFT computation */
#ifndef FFTCOMPLEX_T_DEFINED
typedef struct FFTComplex {
    float re, im;
} FFTComplex;
#define FFTCOMPLEX_T_DEFINED
#endif

typedef struct FFTContext {
    uint8_t nbits;
    uint16_t *revtab_fwd;
    uint16_t *revtab_inv;
} FFTContext;
void ff_fft_calc_c(FFTComplex *z, uint32_t nbits);

// complex c = a * b
#define CMUL(cre, cim, are, aim, bre, bim)          \
do {                                                 \
    cre = are * bre - aim * bim;                     \
    cim = are * bim + aim * bre;                     \
} while (0)
#define CMUL3(c, a, b) CMUL((c).re, (c).im, (a).re, (a).im, (b).re, (b).im)

#define DECLARE_ALIGNED(n,t,v) t __attribute__ ((aligned (n))) v
#define COSTABLE(size) DECLARE_ALIGNED(16, float, ff_cos_##size)[size/2]

static const float
C1  = -5.0000000000e-01,
C2  =  4.1666667908e-02, /* 0x3d2aaaab */
C3  = -1.3888889225e-03, /* 0xbab60b61 */
C4  =  2.4801587642e-05, /* 0x37d00d01 */
C5  = -2.7557314297e-07, /* 0xb493f27c */
C6  =  2.0875723372e-09, /* 0x310f74f6 */
C7  = -1.1359647598e-11; /* 0xad47d74e */

static inline float k_cosf(float x)
{
    float z,r;
    z  = x*x;
    r  = z*(C2+z*(C3+z*(C4+z*(C5+z*(C6+z*C7)))));
    return 1.0f + (z*r + C1*z);
}

/* cos(2*pi*x/n) for 0<=x<=n/4, followed by its reverse */
COSTABLE(16);
COSTABLE(32);
COSTABLE(64);
COSTABLE(128);
COSTABLE(256);
COSTABLE(512);
COSTABLE(1024);
COSTABLE(2048);
COSTABLE(4096);
COSTABLE(8192);
COSTABLE(16384);
COSTABLE(32768);
COSTABLE(65536);
COSTABLE(131072);
float * const ff_cos_tabs[] = {
    NULL, NULL, NULL, NULL,
    ff_cos_16, ff_cos_32, ff_cos_64, ff_cos_128, ff_cos_256, ff_cos_512, ff_cos_1024,
    ff_cos_2048, ff_cos_4096, ff_cos_8192, ff_cos_16384, ff_cos_32768, ff_cos_65536,
    ff_cos_131072
};

void ff_init_ff_cos_tabs(unsigned int index)
{
    unsigned int i, m = 1<<index;
    float freq = 2*(float)M_PI/m;
    float *tab = ff_cos_tabs[index];
    for(i=0; i<=m/4; i++)
        tab[i] = k_cosf(i*freq);
    for(i=1; i<m/4; i++)
        tab[m/2-i] = tab[i];
}

/*** BEGIN stuff taken from fft.c ***/

int split_radix_permutation(int i, int n, int inverse)
{
    int m;
    if(n <= 2) return i&1;
    m = n >> 1;
    if(!(i&m))            return split_radix_permutation(i, m, inverse)*2;
    m >>= 1;
    if(inverse == !(i&m)) return split_radix_permutation(i, m, inverse)*4 + 1;
    else                  return split_radix_permutation(i, m, inverse)*4 - 1;
}

void ff_fft_init(FFTContext *s, uint32_t nbits)
{
    int i, n;
    s->nbits = nbits;
    n = 1 << nbits;
    s->revtab_fwd = malloc(2 * n * sizeof(uint16_t));
    s->revtab_inv = s->revtab_fwd + n;
    for(i=0; i<n; i++) {
        unsigned int j = i;
        j = (j&~3) | ((j>>1)&1) | ((j<<1)&2);
        s->revtab_fwd[-split_radix_permutation(i, n, 0) & (n-1)] = j;
    }
    for(i=0; i<n; i++) {
        unsigned int j = i;
        j = (j&~3) | ((j>>1)&1) | ((j<<1)&2);
        s->revtab_inv[-split_radix_permutation(i, n, 1) & (n-1)] = j;
    }
}

#define sqrthalf (float)M_SQRT1_2
#define COS_1_8 0.3826834559f
#define COS_3_8 0.9238795042f

#define BF(x,y,a,b) {\
    x = a - b;\
    y = a + b;\
}

#define BUTTERFLIES(a0,a1,a2,a3) {\
    BF(t3, t5, t5, t1);\
    BF(a2.re, a0.re, a0.re, t5);\
    BF(a3.im, a1.im, a1.im, t3);\
    BF(t4, t6, t2, t6);\
    BF(a3.re, a1.re, a1.re, t4);\
    BF(a2.im, a0.im, a0.im, t6);\
}

#define TRANSFORM(a0,a1,a2,a3,wre,wim) {\
    t1 = a2.re * wre + a2.im * wim;\
    t2 = a2.im * wre - a2.re * wim;\
    t5 = a3.re * wre - a3.im * wim;\
    t6 = a3.im * wre + a3.re * wim;\
    BUTTERFLIES(a0,a1,a2,a3)\
}

#define TRANSFORM_HALF(a0,a1,a2,a3) {\
    t1 = a2.re + a2.im; t1 *= sqrthalf;\
    t2 = a2.im - a2.re; t2 *= sqrthalf;\
    t5 = a3.re - a3.im; t5 *= sqrthalf;\
    t6 = a3.im + a3.re; t6 *= sqrthalf;\
    BUTTERFLIES(a0,a1,a2,a3)\
}

#define TRANSFORM_ZERO(a0,a1,a2,a3) {\
    t1 = a2.re;\
    t2 = a2.im;\
    t5 = a3.re;\
    t6 = a3.im;\
    BUTTERFLIES(a0,a1,a2,a3)\
}

#if 0
DECLARE_ALIGNED(16, uint32_t, ps_p1p1m1p1)[4] = { 0, 0, 0x80000000, 0 };

void fft4(FFTComplex *z)
{
    __asm__ __volatile__(
    ".intel_syntax noprefix \n\t"
    "movaps xmm0, [%0]\n\t"
    "movaps xmm1, [%0+0x10]\n\t"
    "movaps xmm2, xmm0\n\t"
    "subps  xmm2, xmm1\n\t"
    "addps  xmm0, xmm1\n\t"
    "xorps  xmm2, [%1]\n\t"
    "movaps xmm1, xmm2\n\t"
    "movaps xmm2, xmm0\n\t"
    "shufps xmm2, xmm1, 190\n\t"
    "shufps xmm0, xmm1, 68\n\t"
    "movaps xmm1, xmm0\n\t"
    "subps  xmm1, xmm2\n\t"
    "addps  xmm0, xmm2\n\t"
    "movaps [%0], xmm0\n\t"
    "movaps [%0+0x10], xmm1\n\t"
    ".att_syntax \n\t"
    :"+r"(z) :"r"(ps_p1p1m1p1));
}
#endif

/* fft4/fft8 are not identical to the ones in fft.c, they've been slightly modified to work with the SWAP_LSBS permutation instead of the standard mixed-radix Cooley-Tukey one,
   so I can use the same permutation for C and SSE code. The rest are exactly the same. */
void fft4(FFTComplex *z)
{
    float t1, t2, t3, t4, t5, t6, t7, t8;

    BF(t3, t1, z[0].re, z[2].re);
    BF(t8, t6, z[3].re, z[1].re);
    BF(t4, t2, z[0].im, z[2].im);
    BF(t7, t5, z[1].im, z[3].im);
    BF(z[3].im, z[1].im, t4, t8);
    BF(z[3].re, z[1].re, t3, t7);
    BF(z[2].im, z[0].im, t2, t5);
    BF(z[2].re, z[0].re, t1, t6);
}

void fft8(FFTComplex *z)
{
    float t1, t2, t3, t4, t5, t6;

    fft4(z);

    BF(t1, z[6].re, z[4].re, -z[6].re);
    BF(t2, z[6].im, z[4].im, -z[6].im);
    BF(t5, z[7].re, z[5].re, -z[7].re);
    BF(t6, z[7].im, z[5].im, -z[7].im);
    BUTTERFLIES(z[0],z[2],z[4],z[5]);

    t1 = z[6].re * sqrthalf + z[6].im * sqrthalf;
    t2 = z[6].im * sqrthalf - z[6].re * sqrthalf;
    t5 = z[7].re * sqrthalf - z[7].im * sqrthalf;
    t6 = z[7].im * sqrthalf + z[7].re * sqrthalf;
    BUTTERFLIES(z[1],z[3],z[6],z[7]);
    t1 = z[6].re; z[6].re = z[5].re; z[5].re = t1;
    t1 = z[6].im; z[6].im = z[5].im; z[5].im = t1;
}

void fft16(FFTComplex *z)
{
    float t1, t2, t3, t4, t5, t6;

    fft8(z);
    fft4(z+8);
    fft4(z+12);

    t1 = z[8].re; t2 = z[8].im;
    t5 = z[12].re; t6 = z[12].im;
    BUTTERFLIES(z[0],z[4],z[8],z[12]);
    t1 = z[10].re * sqrthalf + z[10].im * sqrthalf;
    t2 = z[10].im * sqrthalf - z[10].re * sqrthalf;
    t5 = z[14].re * sqrthalf - z[14].im * sqrthalf;
    t6 = z[14].im * sqrthalf + z[14].re * sqrthalf;
    BUTTERFLIES(z[2],z[6],z[10],z[14]);
    t1 = z[ 9].re * COS_3_8 + z[ 9].im * COS_1_8;
    t2 = z[ 9].im * COS_3_8 - z[ 9].re * COS_1_8;
    t5 = z[13].re * COS_3_8 - z[13].im * COS_1_8;
    t6 = z[13].im * COS_3_8 + z[13].re * COS_1_8;
    BUTTERFLIES(z[1],z[5],z[9],z[13]);
    t1 = z[11].re * COS_1_8 + z[11].im * COS_3_8;
    t2 = z[11].im * COS_1_8 - z[11].re * COS_3_8;
    t5 = z[15].re * COS_1_8 - z[15].im * COS_3_8;
    t6 = z[15].im * COS_1_8 + z[15].re * COS_3_8;
    BUTTERFLIES(z[3],z[7],z[11],z[15]);
}

extern COSTABLE(16);
extern COSTABLE(32);
extern COSTABLE(64);
extern COSTABLE(128);
extern COSTABLE(256);
extern COSTABLE(512);
extern COSTABLE(1024);
extern COSTABLE(2048);
extern COSTABLE(4096);
extern COSTABLE(8192);
extern COSTABLE(16384);
extern COSTABLE(32768);
extern COSTABLE(65536);

/* z[0...8n-1], w[1...2n-1] */
static void pass(FFTComplex *z, const float *wre, unsigned int n)
{
    float t1, t2, t3, t4, t5, t6;
    int o1 = 2*n;
    int o2 = 4*n;
    int o3 = 6*n;
    const float *wim = wre+o1;
    n--;

    TRANSFORM_ZERO(z[0],z[o1],z[o2],z[o3]);
    TRANSFORM(z[1],z[o1+1],z[o2+1],z[o3+1],wre[1],wim[-1]);
    do {
        z += 2;
        wre += 2;
        wim -= 2;
        TRANSFORM(z[0],z[o1],z[o2],z[o3],wre[0],wim[0]);
        TRANSFORM(z[1],z[o1+1],z[o2+1],z[o3+1],wre[1],wim[-1]);
    } while(--n);
}

#define DECL_FFT(n,n2,n4)\
void fft##n(FFTComplex *z)\
{\
    fft##n2(z);\
    fft##n4(z+n4*2);\
    fft##n4(z+n4*3);\
    pass(z,ff_cos_##n,n4/2);\
}

DECL_FFT(32,16,8)
DECL_FFT(64,32,16)
DECL_FFT(128,64,32)
DECL_FFT(256,128,64)
DECL_FFT(512,256,128)
DECL_FFT(1024,512,256)
DECL_FFT(2048,1024,512)
DECL_FFT(4096,2048,1024)
DECL_FFT(8192,4096,2048)
DECL_FFT(16384,8192,4096)
DECL_FFT(32768,16384,8192)
DECL_FFT(65536,32768,16384)

static void (* const fft_dispatch[])(FFTComplex*) = {
    fft4, fft8, fft16, fft32, fft64, fft128, fft256, fft512, fft1024,
    fft2048, fft4096, fft8192, fft16384, fft32768, fft65536
};

void ff_fft_calc_c(FFTComplex *z, uint32_t nbits)
{
    fft_dispatch[nbits-2](z);
}

/*** END stuff taken from fft.c ***/

/*
 * This is identical to fft5.c from imdct15.c, but simplified by hand. Code output from this is approximately HALF(!) the size of the original when compiled with recent gcc.
 * Unsurprisingly, it's almost twice as fast.
 */
void fft5(FFTComplex *out, const FFTComplex *in, size_t stride)
{
    FFTComplex z0[4], t0, t1, t2, t3, t4, t5;

    t0.re = in[1 * stride].re + in[4 * stride].re;
    t0.im = in[1 * stride].im + in[4 * stride].im;
    t1.im = in[1 * stride].re - in[4 * stride].re;
    t1.re = in[1 * stride].im - in[4 * stride].im;
    t2.re = in[2 * stride].re + in[3 * stride].re;
    t2.im = in[2 * stride].im + in[3 * stride].im;
    t3.im = in[2 * stride].re - in[3 * stride].re;
    t3.re = in[2 * stride].im - in[3 * stride].im;

    out[0].re = in[0].re + t0.re + t2.re;
    out[0].im = in[0].im + t0.im + t2.im;

    t4.re = -(0.809017f * t0.re - 0.309017f * t2.re);
    t4.im = -(0.809017f * t0.im - 0.309017f * t2.im);
    t0.re =  (0.309017f * t0.re - 0.809017f * t2.re);
    t0.im =  (0.309017f * t0.im - 0.809017f * t2.im);
    t5.re =  0.951057f * t3.re - 0.587785f * t1.re;
    t5.im =  0.951057f * t3.im - 0.587785f * t1.im;
    t1.re =  0.951057f * t1.re + 0.587785f * t3.re;
    t1.im =  0.951057f * t1.im + 0.587785f * t3.im;
    
    z0[0].re = t0.re - t1.re;
    z0[0].im = t0.im - t1.im;
    z0[1].re = t4.re + t5.re;
    z0[1].im = t4.im + t5.im;

    z0[2].re = t4.re - t5.re;
    z0[2].im = t4.im - t5.im;
    z0[3].re = t0.re + t1.re;
    z0[3].im = t0.im + t1.im;

    out[1].re = in[0].re + z0[3].re;
    out[1].im = in[0].im + z0[0].im;
    out[2].re = in[0].re + z0[2].re;
    out[2].im = in[0].im + z0[1].im;
    out[3].re = in[0].re + z0[1].re;
    out[3].im = in[0].im + z0[2].im;
    out[4].re = in[0].re + z0[0].re;
    out[4].im = in[0].im + z0[3].im;
}

/*
 * fft15 exptab hardcoded. It's just 40 float values (the last line is a repeat of the first, not needed here but may be helpful for SIMD variants).
 */
static const FFTComplex fft15_exptab[20] = {
    {  1.0000000,  0.00000000 },
    {  0.9135454, -0.40673665 }, 
    {  0.6691306, -0.74314480 },
    {  0.3090170, -0.95105650 },
    { -0.1045285, -0.99452190 },
    { -0.5000000, -0.86602550 },
    { -0.8090170, -0.58778525 },
    { -0.9781476, -0.20791170 },
    { -0.9781476,  0.20791170 },
    { -0.8090170,  0.58778525 },
    { -0.5000000,  0.86602550 },
    { -0.1045285,  0.99452190 },
    {  0.3090170,  0.95105650 },
    {  0.6691306,  0.74314480 },
    {  0.9135454,  0.40673665 },
    {  1.0000000,  0.00000000 },
    {  0.9135454, -0.40673665 },
    {  0.6691306, -0.74314480 }, 
    {  0.3090170, -0.95105650 },
    { -0.1045285, -0.99452190 },
};

/*
 * IMPORTANT: This is NOT the same as fft15 in imdct15.c
 *            It assumes an INPUT stride of 1 and puts OUTPUTS stride values apart.
 *            This is so I can use it as a base in the PFA algorithm, it also makes it so fft5 does not need to handle strides other than 1 or 3, making SIMDing it dramatically easier.
 */
void fft15(FFTComplex *out, const FFTComplex *in, size_t stride)
{
    FFTComplex tmp[5], tmp1[5], tmp2[5];
    unsigned int k;

    fft5(tmp,  in    , 3);
    fft5(tmp1, in + 1, 3);
    fft5(tmp2, in + 2, 3);

    for (k = 0; k < 5; k++) {
        FFTComplex t1, t2;

        CMUL3(t1, tmp1[k], fft15_exptab[k]);
        CMUL3(t2, tmp2[k], fft15_exptab[2 * k]);
        out[stride*k].re = tmp[k].re + t1.re + t2.re;
        out[stride*k].im = tmp[k].im + t1.im + t2.im;

        CMUL3(t1, tmp1[k], fft15_exptab[k + 5]);
        CMUL3(t2, tmp2[k], fft15_exptab[2 * (k + 5)]);
        out[stride*(k + 5)].re = tmp[k].re + t1.re + t2.re;
        out[stride*(k + 5)].im = tmp[k].im + t1.im + t2.im;

        CMUL3(t1, tmp1[k], fft15_exptab[k + 10]);
        CMUL3(t2, tmp2[k], fft15_exptab[2 * k + 5]);
        out[stride*(k + 10)].re = tmp[k].re + t1.re + t2.re;
        out[stride*(k + 10)].im = tmp[k].im + t1.im + t2.im;
    }
}

/* 1, 2, 4, 8, ..., 2^17 */
static const uint16_t invmod2vals[18] = { 0x0, 0x0, 0x3, 0x7, 0xf, 0xf, 0x2f, 0x6f, 0xef, 0xef, 0x2ef, 0x6ef, 0xeef, 0xeef, 0x2eef, 0x6eef, 0xeeef, 0xeeef };

/*
 * The actual PFA FFT transform for a 2^k * 15 size transform.
 * s needs to be initialized to an FFT Context for a 2^k FFT.
 */
void ff_fft15_transform(FFTContext *s, FFTComplex *out, FFTComplex *in)
{
    FFTComplex fft15in[15];
    unsigned int nbits = s->nbits, n = (1U << nbits), Ntot = 15*n;
    unsigned int inv_1 = (n << ((4 - nbits)&3));
    unsigned int inv_2 = (0xeeeeeeef & ((1U << nbits)-1));
    //unsigned int inv_2 = invmod2vals[nbits];
    unsigned int i, j;

    for (i = 0; i < n; i++) {
        unsigned int l = s->revtab_fwd[i];
        for (j = 0; j < 15; j++) {
            unsigned int k = 15*i, q = (n*j)/15 + i;
            q >>= nbits;
            k += ((j - q*15)<<nbits);
            fft15in[j] = in[k];
        }
        fft15(out+l, fft15in, n);
    }
    for (i = 0; i < 15; i++) {
        ff_fft_calc_c(out+n*i, nbits);
    }
    memcpy(in, out, Ntot * sizeof(FFTComplex));
    for (j = 0; j < 15; j++) {
        for (i = 0; i < n; i++) {
            unsigned int k, q= (((j*inv_1)/15)+(i*inv_2))>>nbits; // this fairly odd way of doing things is to avoid divisions by non-constant non-power-of-2 values
            k = (i*inv_2*15 + j*inv_1) - q*Ntot;
            out[k] = in[n*j + i];
        }
    }
}

#ifdef TEST

static void disable_denormals(void)
{
    unsigned int a = 0x9f80;
    __asm__ __volatile__ ("ldmxcsr %0" : "+m" (a));
}

static void check_diff(float *tab1, float *tab2, unsigned int n)
{
    unsigned int i, n2 = (n>>1);
    float max= 0;
    float error= 0, err2= 0;

    for (i = 0; i < n2; i++) {
        float e1= fabsf(tab1[2*i] - tab2[2*i]);
        float e2= fabsf(tab1[2*i+1] - tab2[2*i+1]);
        if (e1 >= 3e-3f || e2 >= 3e-3f) {
            printf("ERROR %u: (%.5f,%.5f) (%.5f,%.5f)\n", i, tab1[2*i], tab1[2*i+1], tab2[2*i], tab2[2*i+1]);
        }
        error+= e1*e1;
        error+= e2*e2;
        if(e1>max) max= e1;
        if(e2>max) max= e2;
    }
    err2 = sqrtf(error)/n;
    printf("max:%f e:%g\n", max, err2);
}

static uint64_t rngstate = 0xdeadbeefdeadbeef; /* The state can be seeded with any value. */

uint64_t splitmix64(void) {
    uint64_t z;
    rngstate += 0x9E3779B97F4A7C15ULL;
    z = rngstate;
    z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
    z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
    return z ^ (z >> 31);
}

int main(int argc, char **argv)
{
    fftwf_plan fft480;
    FFTContext s;
    FFTComplex in[122880], fft15out[122880], fft15ref[122880];
    fftwf_complex *z;
    fftwf_complex fft_data;
    unsigned int i, n, nbits = (argc > 1) ? strtoul(argv[1], NULL, 10) : 4, Ntot; /* default to n=32 for the power-of-2 fft, corresponding to a size-480 transform for a 960-point MDCT */

    if (nbits > 12) {
        printf("Usage: pfa_test [nbits] (2 <= nbits <= 12)\n");
        return -1;
    }

    disable_denormals();
    for(i=5; i<=17; i++) {
        ff_init_ff_cos_tabs(i);
    }
    n = (1U << nbits);
    Ntot = 15*n;
    z = (fftwf_complex*) memalign(32, sizeof(fftwf_complex) * Ntot);
    ff_fft_init(&s, nbits);
    for (i = 0; i < Ntot; i++) {
        in[i].re = (float)((int16_t)splitmix64() * 3.0517578125e-05f);
        in[i].im = (float)((int16_t)splitmix64() * 3.0517578125e-05f);
        z[i][0] = in[i].re;
        z[i][1] = in[i].im;
    }
    fft480 = fftwf_plan_dft_1d(Ntot, z, &fft_data, FFTW_FORWARD, FFTW_ESTIMATE);
    ff_fft15_transform(&s, fft15out, in);
    fftwf_execute_dft(fft480, z, (fftwf_complex*)fft15ref);
    check_diff((float*)fft15out, (float*)fft15ref, Ntot);
    fftwf_destroy_plan(fft480);
    return 0;
}

#endif