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

struct Transform {
    int                 blocksize;
    TransformCallback   callback;
    void *              callback_data;
    fftwf_plan          fft;
    fftwf_plan          ifft;
    float *             overlap; 
    int16_t *           outframes;
    int                 windex;
    int                 rindex;
    int                 rlen;
    float   *           hanning;
    float   *           real;
    fftwf_complex *     cplx;
    int                 begun;
    float               scale;
};

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 blocksize, TransformCallback callback, void *data) {
    Transform *self     = calloc(1, sizeof(Transform));
    self->blocksize     = blocksize;
    self->callback      = callback;
    self->callback_data = data;
    self->cplx          = fftwf_malloc((blocksize / 2 + 1) *  sizeof(*self->cplx));
    self->real          = fftwf_malloc(blocksize * sizeof(*self->real));
    self->fft           = fftwf_plan_dft_r2c_1d(blocksize, self->real, self->cplx, FFTW_ESTIMATE);
    self->ifft          = fftwf_plan_dft_c2r_1d(blocksize, self->cplx, self->real, FFTW_ESTIMATE);
    self->overlap       = calloc(blocksize + blocksize / 2, sizeof(*self->overlap));
    self->outframes     = calloc(blocksize / 2, sizeof(*self->outframes));
    self->hanning       = calloc(blocksize, sizeof(*self->hanning));
    generate_hanning(self->hanning, blocksize);
    self->scale         = INT16_MAX / (float) blocksize;

    return self;
}

void transform_destroy(Transform *self) {
    fftwf_free(self->real);
    fftwf_free(self->cplx);
    free(self->hanning);
    free(self->overlap);
    fftwf_destroy_plan(self->ifft);
    fftwf_destroy_plan(self->fft);
    free(self);
}

void transform_push(Transform *self, int16_t *frames) {

    int i, ii;

    for (i = 0; i < self->blocksize; i++) {
        self->real[i] = frames[i] / (float) INT16_MAX * self->hanning[i];
    }

    fftwf_execute(self->fft);

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

    fftwf_execute(self->ifft);

    ii = self->blocksize / 2;
    float *overlap = self->overlap + self->windex * self->blocksize / 2;
    for (i = 0; i < ii; i++) {
        overlap[i] += self->real[i] * self->scale;
    }
    if (self->windex == 2) {
        overlap      = self->overlap;
        self->windex = 0;
    } else {
        overlap += ii;
        self->windex++;
    }

    float *real2 = self->real + ii;
    for (i = 0; i < ii; i++) {
        overlap[i] = real2[i] * self->scale;
    }

    if (!self->begun) {
        self->begun  = 1;
    } else {
        self->rindex = (self->windex + 1) % 3;
        self->rlen   = 2;
    }
}


int16_t * transform_pop(Transform *self) {
    int16_t *outframes = NULL;
    if (self->rlen) {
        float *overlap = self->overlap + self->rindex++ * self->blocksize / 2;
        self->rindex %= 3;
        self->rlen--;
        outframes = self->outframes;
        int i;
        for (i = 0; i < self->blocksize / 2; i++)
            outframes[i] = overlap[i];
    }
    return outframes;
}