#include "lmath.h"

#include <math.h>
#ifdef LMATH_USE_FFTW
#include <fftw3.h>
#endif
#include <stdlib.h>
#include <string.h>

#ifndef M_PI
    #define M_PI 3.1415926535
#endif

#ifndef MAX
    #define MAX(x,y) ((x) > (y) ? (x) : (y))
#endif


/** Параметры используемого окна. Заполняются при вызове функций
 *  lmath_window()/lmath_scaled_window() и используются другими
 *  функциями при дальнейшей обработке */
typedef struct {
    double inherent_gain; /**< Сумма коэф. окна в квадрате / size */
    double coherent_gain; /**< Квадрат суммы коэф. окна, деленной на size */
    double coherent_gain_sqrt; /**< Сумма коэф. окна, деленной на size */
    double enbw; /**< Equivalent Noise Bandwidth = inherent_gain/coherent_gain */
    int freq_span; /**< Ширина пика одной частоты в спектре */
} t_lmath_win_params;


struct st_lmath_win_context {
    double *coefs;
    int size;
    t_lmath_win_params params;
};


typedef double (*t_win_pont_func)(int i, int n, void* param);

static double f_rectangle_win_point(int i, int n, void *param);
static double f_hanning_win_point(int i, int n, void *param);
static double f_hamming_win_point(int i, int n, void* param);
static double f_bh4_win_point(int i, int n, void *param);
static double f_bh7_win_pint(int i, int n, void *param);

static t_win_pont_func f_win_func_tbl[] = {
    f_rectangle_win_point,
    f_hanning_win_point,
    f_hamming_win_point,
    0,
    0,
    0,
    0,
    f_bh4_win_point,
    f_bh7_win_pint
};


static double f_rectangle_win_point(int i, int n, void *param) {
    return 1.;
}


static double f_hanning_win_point(int i, int n, void* param) {
    double a= (2.*M_PI*i)/n;
    return 0.5*(1-cos(a));
}

static double f_hamming_win_point(int i, int n, void* param) {
    double a= (2.*M_PI*i)/n;
    return 0.54-0.46*cos(a);
}

static double f_bh4_win_point(int i, int n, void *param) {
    static const double a0 = 0.35875;
    static const double a1 = 0.48829;
    static const double a2 = 0.14128;
    static const double a3 = 0.01168;
    double a= (2.*M_PI*i)/n;
    return a0 - a1*cos(a) + a2*cos(2*a) - a3*cos(3*a);
}


static double f_bh7_win_pint(int i, int n, void *param) {
    static const double a0 = 0.27105140069342415;
    static const double a1 = 0.433297939234486060;
    static const double a2 = 0.218122999543110620;
    static const double a3 = 0.065925446388030898;
    static const double a4 = 0.010811742098372268;
    static const double a5 = 7.7658482522509342E-4;
    static const double a6 = 1.3887217350903198E-5;
    double a= (2.*M_PI*i)/n;
    return a0 - a1*cos(a) + a2*cos(2*a) - a3*cos(3*a) + a4*cos(4*a) - a5*cos(5*a) + a6*cos(6*a);
}




t_lmath_window_context lmath_window_context_create(int wintype, void *param, int size) {
    t_lmath_window_context ctx = NULL;
    if ((wintype < sizeof(f_win_func_tbl)/sizeof(f_win_func_tbl[0])) ||
            (f_win_func_tbl[wintype] != NULL)) {
        ctx = (t_lmath_window_context)malloc(sizeof(*ctx));
        if (ctx != NULL) {
            ctx->size = size;
            ctx->coefs = (double*)malloc(size*sizeof(ctx->coefs[0]));
            if (ctx->coefs == NULL) {
                free(ctx);
                ctx = NULL;
            } else {
                double sum =0.0;
                double sum2 = 0;
                int i;
                t_win_pont_func func = f_win_func_tbl[wintype];
                for (i=0; i < size; i++) {
                    double val = func(i, size, param);
                    sum+=val;
                    sum2 += (val*val);
                    ctx->coefs[i] = val;
                }

                ctx->params.coherent_gain_sqrt =  sum/size;
                ctx->params.coherent_gain = ctx->params.coherent_gain_sqrt * ctx->params.coherent_gain_sqrt;
                ctx->params.inherent_gain = sum2/size;
                ctx->params.enbw = ctx->params.inherent_gain/ctx->params.coherent_gain;
                ctx->params.freq_span = 9;
            }
        }
    }
    return ctx;
}

void lmath_window_context_destroy(t_lmath_window_context ctx) {
    if (ctx) {
        free(ctx->coefs);
        free(ctx);
    }
}

int lmath_window_context_win_size(const t_lmath_window_context winctx) {
    return winctx->size;
}

double lmath_window_context_inherent_gain(const t_lmath_window_context winctx) {
    return winctx->params.inherent_gain;
}

double lmath_window_context_coherent_gain(const t_lmath_window_context winctx) {
    return winctx->params.coherent_gain;
}

double lmath_window_context_enbw(const t_lmath_window_context winctx) {
    return winctx->params.enbw;
}


int lmath_window_apply(const double *src, const t_lmath_window_context winctx, double *res) {
    int i;
    for (i = 0; i < winctx->size; i++) {
        res[i] = src[i]*winctx->coefs[i];
    }
    return  0;
}

int lmath_window_apply_scaled(const double *src, t_lmath_window_context winctx, double *res) {
    double del = winctx->params.coherent_gain_sqrt;
    int i;
    for (i = 0; i < winctx->size; i++) {
        res[i] = src[i]*winctx->coefs[i]/del;
    }
    return  0;
}


#ifdef LMATH_USE_FFTW

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

t_lmath_fft_r2c_context lmath_fft_r2c_context_create(int size, int flags) {
    t_lmath_fft_r2c_context 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_context_destroy(ctx);
            ctx = NULL;
        }
    }
    return ctx;
}

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

double *lmath_fft_r2c_context_in_buf(const t_lmath_fft_r2c_context ctx) {
    return ctx->in;
}

int lmath_fft_r2c_context_in_size(const t_lmath_fft_r2c_context ctx) {
    return ctx->size;
}

int lmath_fft_r2c_context_fft_size(const t_lmath_fft_r2c_context ctx) {
    return ctx->fft_size;
}

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

int lmath_fft_r2c_calc_amp_pha_spectrum_peak(t_lmath_fft_r2c_context ctx, 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] = atan2(ctx->fft[i][1]*ctx->fft[i][1], ctx->fft[i][0]*ctx->fft[i][0]);
    }
    return 0;
}

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


int lmath_amp_pha_spectrum_peak(const double *src, int size, double dt, int flags,
                           double *amp, double *pha, double *df) {
    int err = 0;
    t_lmath_fft_r2c_context ctx = lmath_fft_r2c_context_create(size, flags);
    if (ctx == NULL) {
        err = LMATH_ERR_MEMORY_ALLOC;
    } else {
        memcpy(ctx->in, src, sizeof(double)*size);
        err = lmath_fft_r2c_calc_amp_pha_spectrum_peak(ctx, amp, pha);
        if (df != NULL)
            *df = lmath_fft_r2c_context_df(ctx, dt);
        lmath_fft_r2c_context_destroy(ctx);
    }

    return err;
}

int lmath_amp_pha_spectrum_rms(const double *src, int size, double dt, int flags,
                               double *amp, double *pha, double *df) {
    int err = 0;
    t_lmath_fft_r2c_context ctx = lmath_fft_r2c_context_create(size, flags);
    if (ctx == NULL) {
        err = LMATH_ERR_MEMORY_ALLOC;
    } else {
        memcpy(ctx->in, src, sizeof(double)*size);
        err = lmath_fft_r2c_calc_amp_pha_spectrum_rms(ctx, amp, pha);
        if (df != NULL)
            *df = lmath_fft_r2c_context_df(ctx, dt);
        lmath_fft_r2c_context_destroy(ctx);
    }

    return err;
}

#endif






int lmath_calc_spectrum_params (const double *ampSpectrum, int size,  double df,
                                double peak_freq, int harm_cnt,
                                const t_lmath_window_context winctx,
                                double *SNR, double *THD, double *SINAD, double *SFDR, double *ENOB) {
    int res=0;
    int i;
    int ipeak = (int)(peak_freq/df + 0.5);
    int freq_span = winctx->params.freq_span;
    double peak_power=0;
    double noise_power=0;
    double harm_power=0;
    double harm_noise_power =0;
    double max_noise_amp = -100000000;
    double max_signal_amp = -100000000;

    if (SNR!=NULL)
        *SNR=0;
    if (THD!=NULL)
        *THD=0;
    if (SINAD!=NULL)
        *SINAD=0;
    if (SFDR!=NULL)
        *SFDR =0;
    if (ENOB!=NULL)
        *ENOB=0;

    if (ipeak < (freq_span / 2))
        return LMATH_ERR_INVALID_PARAMS;

    for(i=freq_span/2; i<size;i++) {
        int nearest_peak_idx = (i+(freq_span / 2))/ipeak;
        int nearest_peak_idx_last = (nearest_peak_idx*ipeak)+(freq_span / 2);
        int nearest_peak_idx_first = (nearest_peak_idx*ipeak)-(freq_span / 2);

        double amp = ampSpectrum[i];

        // используется только первые harm_cnt гармоник
        if ((i>=nearest_peak_idx_first) && (i<=nearest_peak_idx_last) &&
                (nearest_peak_idx>0) && (nearest_peak_idx<harm_cnt)) {
            if(nearest_peak_idx > 1) {
                max_signal_amp = MAX(max_signal_amp, amp);
                harm_power += (amp*amp);
            } else {
                max_noise_amp = MAX(max_noise_amp, amp);
                peak_power += (amp*amp);
            }
        } else {
            max_signal_amp = MAX(max_signal_amp, amp);
            noise_power+=(amp*amp);
        }
    }

    peak_power = peak_power/winctx->params.enbw;
    harm_power = harm_power/winctx->params.enbw;
    noise_power = noise_power/winctx->params.enbw;
    harm_noise_power = noise_power+harm_power;

    if(peak_power != 0) {
        double DynamicRange = max_noise_amp/max_signal_amp;
        if ((DynamicRange>0) && (SFDR!=NULL)) {
            *SFDR = 20 * log10(DynamicRange);
        }

        if (noise_power) {
            double SNRPower =  peak_power / noise_power;
            if((SNRPower>0) && (SNR!=NULL)) {
                *SNR = 10 * log10(SNRPower);
                if (*SNR < 0) {
                    *SNR = 0;
                    res=LMATH_ERR_INTERNAL_CALC;
                }
            }
        }

        if(harm_power) {
            double THDPower = harm_power/peak_power;
            if ((THDPower>0) && (THD!=NULL)) {
                *THD = 10 * log10(THDPower);
                if (*THD >0) {
                    *THD = 0;
                    res=LMATH_ERR_INTERNAL_CALC;
                }
            }
        }

        if (harm_noise_power) {
            double SINADPower = peak_power/harm_noise_power;
            if ((SINADPower>0) && (SINAD!=NULL)) {
                double sinad = 10*log10(SINADPower);
                if (sinad<0) {
                    res=LMATH_ERR_INTERNAL_CALC;
                } else {
                    if (SINAD!=NULL) {
                        *SINAD=sinad;
                    }

                    if (ENOB!=NULL) {
                        *ENOB = (sinad - 1.76) / 6.02;
                    }
                }
            }
        }
    }
    return res;
}


int lmath_acdc_estimation_with_window(const double *src,
                                      const t_lmath_window_context winctx,
                                      double *tmp_buf,
                                      double *dc, double *ac) {

    int i;
    double avg = 0;
    int err = (src == NULL) || (winctx == NULL) || (tmp_buf == NULL)
             || ((ac != NULL) && (tmp_buf == src)) ?
                LMATH_ERR_INVALID_PARAMS : 0;


    if (!err) {
        /* сперва накладываем окно на исходный сигнал для рассчета DC */
        err = lmath_window_apply(src, winctx, tmp_buf);
    }

    if (!err) {        
        for (i = 0; i < winctx->size; i++) {
            avg+=tmp_buf[i];
        }
        avg/=(winctx->size*winctx->params.coherent_gain_sqrt);

        if (dc != NULL)
            *dc = avg;
    }

    if (!err && (ac != NULL)) {
        double rms = 0;
        /* далее для рассчета AC заменяем сигнал на исходный без DC */
        for (i=0; i < winctx->size; i++) {
            tmp_buf[i]=(src[i] - avg);
        }

        /* накладываем окно на сигнал без DC */
        err = lmath_window_apply(tmp_buf, winctx, tmp_buf);
        for (i = 0; i < winctx->size; i++) {
            rms+= tmp_buf[i]*tmp_buf[i];
        }

        rms/=(winctx->size*winctx->params.inherent_gain);
        rms = sqrt(rms);
        if (!err)
            *ac = rms;
    }

    return err;
}

int lmath_acdc_estimation(const double *src, int size, double *dc, double *ac) {
    int err = 0;
    double *tmp_arr = (double *)malloc(sizeof(tmp_arr[0])*size);
    t_lmath_window_context winctx = lmath_window_context_create(LMATH_WINTYPE_BH_4TERM, 0, size);


    if ((tmp_arr==NULL) || (winctx == NULL)) {
        err = LMATH_ERR_MEMORY_ALLOC;
    } else {
        /* сперва накладываем окно на исходный сигнал для рассчета DC */
        err = lmath_acdc_estimation_with_window(src, winctx, tmp_arr, dc, ac);
    }

    free(tmp_arr);
    lmath_window_context_destroy(winctx);

    return err;
}

#ifdef LMATH_USE_FFTW
int lmath_find_peak_freq(const double *src, int size, double dt, double search_freq,
                         double *peak_freq, double *peak_pwr) {
    int err = 0;
    double *proc_vals = malloc(sizeof(proc_vals[0])*size);
    t_lmath_window_context winctx = lmath_window_context_create(LMATH_WINTYPE_BH_4TERM, 0, size);
    double df;


    if ((proc_vals == NULL) || (winctx ==NULL)) {
        err = LMATH_ERR_MEMORY_ALLOC;
    } else {
        /* амплитуду берем только на заданной частоте по спектру */
        err = lmath_window_apply_scaled(src, winctx, proc_vals);
    }

    if (!err) {
        err = lmath_amp_pha_spectrum_peak(proc_vals, size, dt, 0, proc_vals, 0, &df);
    }

    if (!err) {
        err = lmath_find_peak_freq_spectrum(proc_vals, size/2+1, df, search_freq,
                                            winctx, peak_freq, peak_pwr);
    }

    free(proc_vals);
    lmath_window_context_destroy(winctx);

    return err;
}
#endif





int lmath_find_peak_freq_spectrum(const double *amp_spectrum, int size, double df, double search_freq,
                                  const t_lmath_window_context winctx,
                                  double *peak_freq, double *peak_pwr) {
    int i, fnd_i;
    double fnd_pwr = 0;
    int err = 0;
    int freq_span;

    if ((amp_spectrum==NULL) || (winctx == NULL) ||  ((peak_freq==NULL) && (peak_pwr==0))) {
        err = LMATH_ERR_INVALID_PARAMS;
    } else {
        freq_span = winctx->params.freq_span;
        if ((df <= 0) || (freq_span < 1)) {
            err = LMATH_ERR_INVALID_PARAMS;
        }
    }

    if (!err) {
        /* если не задана частота, в окрестности которой нужно найти требуемую,
         * то ищем во всем спектре по максимальной мощьности (исключая пост. составляющую) */
        if (search_freq <= 0) {
            for (fnd_i=-1,i= freq_span/2 + freq_span/2; i < (size - freq_span/2); i++) {
                double pwr=0;
                int j;
                for (j=i-freq_span/2; j <= (i+freq_span/2); j++) {
                    pwr+=(amp_spectrum[j]*amp_spectrum[j]);
                }

                if (pwr > fnd_pwr) {
                    fnd_pwr = pwr;
                    fnd_i = i;
                }
            }

            if (fnd_i < 0)
                err = LMATH_ERR_INTERNAL_CALC;
        } else {
            /* иначе рассчитываем мощность только в окрестности указанной частоты */
            int j;
            fnd_i = (int)(search_freq/df + 0.5);
            if (fnd_i >= (size - freq_span/2))
                fnd_i = size - freq_span/2 - 1;
            if (fnd_i < freq_span)
                fnd_i = freq_span;

            for (fnd_pwr=0, j=fnd_i-freq_span/2; j <= (fnd_i+freq_span/2); j++) {
                fnd_pwr+=(amp_spectrum[j]*amp_spectrum[j]);
            }
        }

    }

    if (!err) {
        if (peak_freq!=NULL) {
            double fnd_freq=0;
            int j;
            /* поиск частоты сигнала */
            for (j= fnd_i-freq_span/2; j <= (fnd_i+freq_span/2); j++) {
                fnd_freq+=(amp_spectrum[j]*amp_spectrum[j])*j*df;
            }
            fnd_freq/=fnd_pwr;

            *peak_freq = fnd_freq;
        }

        if (peak_pwr!=NULL) {
            *peak_pwr = fnd_pwr/winctx->params.enbw;
        }
    }

    return err;
}


double lmath_proc_fir(const double *coef, double *prev_x, int ntap, double new_x) {
    double y=0;            //output sample
    int n;

    //shift the old samples
    for(n=ntap-1; n>0; n--)
       prev_x[n] = prev_x[n-1];

    //Calculate the new output
    prev_x[0] = new_x;
    for(n=0; n<ntap; n++)
        y += coef[n] * prev_x[n];

    return y;
}








