#include <stdlib.h>
#include <stdint.h>
#include <math.h>
#include <string.h>
#include "stft.h"
#include "operations.h"

#define PI 3.14159265358979323846
#define FRAME_SIZE 512
#define HOP_SIZE 32
#define FIXED_POINT_SHIFT 16
#define FFT_SIZE_IN_BITS 9
#define N_WAVE      1024    /* full length of Sinewave[] */
#define LOG2_N_WAVE 10      /* log2(N_WAVE) */
#define FRACT_BITS FIXED_POINT_SHIFT-1
#define FRACT_SCALE (1 << FRACT_BITS)

int16_t** sliding_window(int16_t* x, int num_samples, int framesize, int hopsize) {
    int num_frames = (num_samples - framesize) / hopsize + 1;
    int16_t** frames = malloc(num_frames * sizeof(int16_t*)); 

    for (int i = 0; i < num_frames; i++) {
        frames[i] = malloc(framesize * sizeof(int16_t));
        for (int j = 0; j < framesize; j++) {
            frames[i][j] = x[i * hopsize + j];
        }
    }

    return frames;
}

void apply_hann_window_forward(int16_t *input) {
    // Create a Hann window
    int16_t window[FRAME_SIZE];
    for (int i = 0; i < FRAME_SIZE; i++) {
        double value = 0.5 * (1.0 - cos(2.0 * PI * i / (FRAME_SIZE - 1)));
        window[i] = (int16_t)(value * pow(2, FIXED_POINT_SHIFT));
    }

    // Apply the Hann window to the input data
    for (int i = 0; i < FRAME_SIZE; i++) {
        input[i] = (input[i] * window[i]) >> FIXED_POINT_SHIFT;
    }
}

void apply_hann_window_reverse(int16_t *input) {
    // Create a Hann window
    int16_t window[FRAME_SIZE];
    for (int i = 0; i < FRAME_SIZE; i++) {
        double value = 0.5 * (1.0 - cos(2.0 * PI * i / (FRAME_SIZE - 1)));
        window[i] = (int16_t)(value * pow(2, FIXED_POINT_SHIFT));
    }

    // Apply the Hann window to the input data
    for (int i = 0; i < FRAME_SIZE; i++) {
        input[i] = (input[i] * window[i]) >> FIXED_POINT_SHIFT;
    }
}

// Radix-2 RFFT function
void rfft_fixed(complex_fixed *x, int N, int inverse) {
    // Bit-reversal
    int k, j, m;
    for (k = 1, j = 0; k < N; ++k) {
        int i;
        for (i = N >> 1; i > (j ^= i); i >>= 1);
        if (k < j) {
            complex_fixed temp = x[j];
            x[j] = x[k];
            x[k] = temp;
        }
    }

    // Cooley-Tukey algorithm
    int n = 2;
    int shift = 1;
    while (n <= N) {
        int mmax = n >> 1;
        int angle = 0;
        int m;
        for (m = 0; m < mmax; ++m) {
            int k;
            int16_t wr, wi;
            if (inverse) {
                wr = cos(angle * PI / (n >> 1)) * 32767.0;
                wi = sin(angle * PI / (n >> 1)) * -32767.0;
            } else {
                wr = cos(angle * PI / (n >> 1)) * 32767.0;
                wi = sin(angle * PI / (n >> 1)) * 32767.0;
            }
            angle += shift * 2 * mmax;

            for (k = m; k < N; k += n) {
                int16_t tr = wr * x[k + mmax].real - wi * x[k + mmax].imag;
                int16_t ti = wr * x[k + mmax].imag + wi * x[k + mmax].real;
                x[k + mmax].real = x[k].real - tr;
                x[k + mmax].imag = x[k].imag - ti;
                x[k].real += tr;
                x[k].imag += ti;
            }
        }
        n <<= 1;
        shift = inverse ? -shift : -shift;
    }

    if (inverse) {
        // Scale the output
        int i;
        for (i = 0; i < N; ++i) {
            x[i].real /= N;
            x[i].imag /= N;
        }
    }
}