#pragma once

#include "inttypes.h"
#include "math.h"
#include "stdlib.h"
#include "table.h"

#define Re(x) 2 * x
#define Im(x) 2 * x + 1
#define FFT_BIN(num, fs, size) (uint16_t)((float)(num) * (float)fs / (float)size)
#define FFT_INDEX(freq, fs, size) ((int)((float)freq / ((float)fs / (float)size)))

void bitreversal(q31_t *data, uint16_t len, uint16_t *bitrev, uint16_t step) {
    q31_t tmp;
    uint16_t n1 = len >> 1;
    uint16_t n2 = n1 + 1;

    for (uint16_t i = 0, j = 0; i <= (n1 - 2); i += 2) {
        if (i < j) {
            tmp = data[i];
            data[i] = data[j];
            data[j] = tmp;

            tmp = data[i + n2];
            data[i + n2] = data[j + n2];
            data[j + n2] = tmp;
        }
        tmp = data[i + 1];
        data[i + 1] = data[j + n1];
        data[j + n1] = tmp;

        j = *bitrev;
        bitrev += step;
    }
}

void butterfly(q15_t *data, uint16_t len, uint16_t step) {
    q15_t xt, yt, cosv, sinv;

    for (uint16_t m = len; m > 1; m = m >> 1) {
        uint16_t n1 = m;
        uint16_t n2 = m >> 1;

        for (uint16_t i = 0, ia = 0; i < n2; i++) {
            cosv = rotCoef[ia * 2];
            sinv = rotCoef[ia * 2 + 1];
            ia = ia + step;

            for (uint16_t j = i, k; j < len; j += n1) {
                k = j + n2;
                if (m == len) {
                    xt = (data[Re(j)] >> 2u) - (data[Re(k)] >> 2u);
                    yt = (data[Im(j)] >> 2u) - (data[Im(k)] >> 2u);

                    data[Re(j)] = (q15_t)(((data[Re(j)] >> 2u) + (data[Re(k)] >> 2u)) >> 1);
                    data[Im(j)] = (q15_t)(((data[Im(k)] >> 2u) + (data[Im(j)] >> 2u)) >> 1);

                    data[Re(k)] = ((((q31_t)xt * cosv) >> 16)) + ((((q31_t)yt * sinv) >> 16));
                    data[Im(k)] = ((((q31_t)yt * cosv) >> 16)) - ((((q31_t)xt * sinv) >> 16));
                } else if (m == 2) {
                    xt = data[Re(j)] - data[Re(k)];
                    yt = data[Im(j)] - data[Im(k)];

                    data[Re(j)] = data[Re(j)] + data[Re(k)];
                    data[Im(j)] = data[Im(k)] + data[Im(j)];

                    data[Re(k)] = xt;
                    data[Im(k)] = yt;
                } else {
                    xt = data[Re(j)] - data[Re(k)];
                    yt = data[Im(j)] - data[Im(k)];

                    data[Re(j)] = (q15_t)((data[Re(j)] + data[Re(k)]) >> 1);
                    data[Im(j)] = (q15_t)((data[Im(k)] + data[Im(j)]) >> 1);

                    data[Re(k)] = ((((q31_t)xt * cosv) >> 16)) + ((((q31_t)yt * sinv) >> 16));
                    data[Im(k)] = ((((q31_t)yt * cosv) >> 16)) - ((((q31_t)xt * sinv) >> 16));
                }
            }
        }
        step = step << 1;
    }
}

static inline void applyWindow(q15_t *src, const q15_t *window, uint16_t len) {
    while (len--) {
        *src = (q15_t)(((q31_t)(*src) * (*window)) >> 15);
        src++;
        window++;
    }
}

int ZeroFFT(q15_t *data, uint16_t len) {
    uint16_t step;
    uint16_t *bitrev;
    q15_t out[len * 2];

    switch (len) {
        case 1024u:
            step = 1u;
            bitrev = (uint16_t *)&armBitRevTable[0];
            applyWindow(data, window_hanning_1024, 1024);
            break;

        case 512u:
            step = 2u;
            bitrev = (uint16_t *)&armBitRevTable[1];
            applyWindow(data, window_hanning_512, 512);
            break;

        case 256u:
            step = 4u;
            bitrev = (uint16_t *)&armBitRevTable[3];
            applyWindow(data, window_hanning_256, 256);
            break;

        case 128u:
            step = 8u;
            bitrev = (uint16_t *)&armBitRevTable[7];
            applyWindow(data, window_hanning_128, 128);
            break;

        default:
            return -1;
            break;
    }

    for (uint16_t i = 0; i < len; i++) {
        out[Re(i)] = data[i];  // real
        out[Im(i)] = 0;        // imaginary
    }

    butterfly(out, len, step);
    bitreversal((q31_t *)out, len, bitrev, step);
    for (uint16_t i = 0; i < len; i++) data[i] = (q15_t)abs(out[2 * i]);

    return 0;
}