#include "lmath_fft.h"
#include "lmath_internal.h"
#include <fftw3.h>
#include <stdlib.h>
#include <string.h>

struct st_lmath_fft_r2c_ctx {
    int size;
    int fft_size;
    double *in;
    fftw_complex *fft;
    fftw_plan plan;
};

t_lmath_fft_r2c_ctx lmath_fft_r2c_ctx_create(int size, int flags) {
    t_lmath_fft_r2c_ctx ctx = malloc(sizeof(*ctx));
    if (ctx) {
        ctx->size = size;
        ctx->fft_size = size/2 + 1;
        ctx->fft = (fftw_complex*)fftw_malloc(sizeof(fftw_complex)*ctx->fft_size);
        ctx->in  = (double*)fftw_malloc(sizeof(double)*size);
        ctx->plan = fftw_plan_dft_r2c_1d(size, ctx->in, ctx->fft, FFTW_ESTIMATE);
        if ((ctx->fft == NULL) || (ctx->in == NULL) || (ctx->plan == NULL)) {
            lmath_fft_r2c_ctx_destroy(ctx);
            ctx = NULL;
        }
    }
    return ctx;
}

void lmath_fft_r2c_ctx_destroy(t_lmath_fft_r2c_ctx ctx) {
    if (ctx) {
        fftw_destroy_plan(ctx->plan);
        fftw_free(ctx->fft);
        fftw_free(ctx->in);
        free(ctx);
    }
}

double *lmath_fft_r2c_ctx_in_buf(const t_lmath_fft_r2c_ctx ctx) {
    return ctx->in;
}

int lmath_fft_r2c_ctx_in_size(const t_lmath_fft_r2c_ctx ctx) {
    return ctx->size;
}

int lmath_fft_r2c_ctx_fft_size(const t_lmath_fft_r2c_ctx ctx) {
    return ctx->fft_size;
}

double lmath_fft_r2c_ctx_df(const t_lmath_fft_r2c_ctx ctx, double dt) {
    return 1./(dt*ctx->size);
}

t_lmath_errs lmath_fft_r2c_calc_amp_pha_spectrum_peak(t_lmath_fft_r2c_ctx ctx,
                                                      t_lmath_phase_mode pha_mode,
                                                      double *amp, double *pha) {
    int i;
    fftw_execute(ctx->plan);
    for (i=0; i < ctx->fft_size; i++) {
        if (amp != NULL)
            amp[i] = sqrt(ctx->fft[i][0]*ctx->fft[i][0] + ctx->fft[i][1]*ctx->fft[i][1])*
                     ((i == 0) || (i == (ctx->fft_size-1)) ? 1 : 2)/ctx->size;
        if (pha != NULL)
            pha[i] = lmath_phase_calc(ctx->fft[i][0], ctx->fft[i][1], pha_mode);
    }
    return LMATH_ERR_OK;
}

t_lmath_errs lmath_fft_r2c_calc_amp_pha_spectrum_rms(t_lmath_fft_r2c_ctx ctx,
                                                     t_lmath_phase_mode pha_mode,
                                                     double *amp, double *pha) {
    int ret = lmath_fft_r2c_calc_amp_pha_spectrum_peak(ctx, pha_mode, amp, pha);
    if ((ret == 0) && (amp != NULL)) {
        int i;
        for (i=1; i < ctx->fft_size; i++) {
            amp[i] /= sqrt(2);
        }
    }
    return LMATH_ERR_OK;
}




