#include <stdlib.h>
#include <math.h>
#include <string.h>
#include "transform.h"

struct Transform {
    int                 fftsize;
    TransformCallback   callback;
    void *              callback_data;
    fftwf_plan          fft;
    fftwf_plan          ifft;
    float   *           real;
    fftwf_complex *     cplx;
    float   *           window;
    int16_t *           input_save;
    float  *            output_save;
};

static void generate_hanning(float *window, int size) {
    int i;
    float factor = 2.f * M_PI / (size - 1.f);
    for (i = 0; i < size; i++) {
        window[i] = (1.f - cos(i * factor)) / 2.f;
    }
}

Transform *transform_new(int fftsize, TransformCallback callback, void *data) {
    Transform *self     = calloc(1, sizeof(Transform));

    self->fftsize       = fftsize;
    self->callback      = callback;
    self->callback_data = data;
    self->cplx          = fftwf_malloc((fftsize / 2 + 1) *  sizeof(*self->cplx));
    self->real          = fftwf_malloc(fftsize * sizeof(*self->real));
    self->input_save    = calloc(fftsize / 2, sizeof(*self->input_save));
    self->output_save   = calloc(fftsize / 2, sizeof(*self->output_save));
    self->fft           = fftwf_plan_dft_r2c_1d(fftsize, self->real, self->cplx, FFTW_ESTIMATE);
    self->ifft          = fftwf_plan_dft_c2r_1d(fftsize, self->cplx, self->real, FFTW_ESTIMATE);
    self->window        = calloc(fftsize, sizeof(*self->window));

    generate_hanning(self->window, fftsize);

    return self;
}

void transform_destroy(Transform *self) {
    fftwf_free(self->real);
    fftwf_free(self->cplx);
    free(self->input_save);
    free(self->output_save);
    free(self->window);
    fftwf_destroy_plan(self->ifft);
    fftwf_destroy_plan(self->fft);
    free(self);
}

void transform_process(Transform *self, int16_t *input, int16_t *output) {

    int i;
    int buffer_size     = self->fftsize / 2;
    float *real1      = self->real;
    float *real2      = self->real + buffer_size;
    float *window1    = self->window;
    float *window2    = self->window + buffer_size;

    for (i = 0; i < buffer_size; i++) {
        real1[i] = self->input_save[i] / (float) INT16_MAX * window1[i];
        real2[i] = input[i]  / (float) INT16_MAX * window2[i];
    }

    memcpy(self->input_save, input, buffer_size * sizeof(*input));

    fftwf_execute(self->fft);

    self->callback(self->cplx, self->callback_data);

    fftwf_execute(self->ifft);

    float scale = INT16_MAX / (float) self->fftsize;
    for (i = 0; i < buffer_size; i++) {
        output[i] = (real1[i] + self->output_save[i]) * scale;
    }

    memcpy(self->output_save, real2, buffer_size * sizeof(*real2));
}