#include "lmath_spectrum_signoise.h"
#include "lmath_internal.h"
#include <stdlib.h>

typedef struct {
    int num;
    int start_idx;
    int end_idx;
} t_lmath_spectrum_harm_param;

struct st_lmath_spectrum_signoise_ctx {
    int fft_size;
    int freq_span;
    int start_idx;
    double enbw;
    int total_harm_cnt;
    int result_harm_cnt;
    unsigned flags;
    t_lmath_spectrum_harm_param *harm_params;
};


t_lmath_spectrum_signoise_ctx lmath_spectrum_signoise_ctx_create(int fft_size,  double df,
                                                                 const t_lmath_window_params *winpar,
                                                                 double peak_freq, int harm_cnt,
                                                                 unsigned flags) {
    t_lmath_spectrum_signoise_ctx ctx = NULL;
    int ipeak = (int)(peak_freq/df + 0.5);

    if (winpar && (ipeak > (winpar->freq_span / 2))) {
        ctx = (t_lmath_spectrum_signoise_ctx)malloc(sizeof(*ctx));
        if (ctx) {
            ctx->harm_params = (t_lmath_spectrum_harm_param*)malloc(sizeof(t_lmath_spectrum_harm_param) * harm_cnt);
            if (ctx->harm_params) {
                int harm_idx = 0;

                ctx->fft_size = fft_size;
                ctx->freq_span = winpar->freq_span;
                ctx->enbw = winpar->enbw;
                ctx->start_idx  = winpar->freq_span/2 + 1;
                ctx->flags = flags;
                ctx->total_harm_cnt = harm_cnt;

                int cur_harm_cnt = 0;

                for (harm_idx = 0; harm_idx < harm_cnt; harm_idx++) {
                    int harm_num = harm_idx + 1;
                    int harm_bin_idx, harm_bin_start, harm_bin_end;
                    harm_bin_idx = (int)(peak_freq/df * harm_num + 0.5);

                    if (harm_bin_idx > (fft_size-1)) {
                        int ov_cnt = harm_bin_idx / (fft_size-1);
                        harm_bin_idx -= ov_cnt * (fft_size-1);
                        if (ov_cnt & 1) {
                            harm_bin_idx = fft_size - harm_bin_idx - 1;
                        }
                    }
                    harm_bin_start = harm_bin_idx - winpar->freq_span/2;
                    if (harm_bin_start < ctx->start_idx)
                        harm_bin_start = ctx->start_idx;

                    harm_bin_end = harm_bin_idx + winpar->freq_span/2;
                    if (harm_bin_end > (fft_size-1))
                        harm_bin_end = fft_size-1;


                    if ((cur_harm_cnt == 0) || (harm_bin_start > ctx->harm_params[cur_harm_cnt-1].end_idx)) {
                        ctx->harm_params[cur_harm_cnt].num = harm_num-1;
                        ctx->harm_params[cur_harm_cnt].start_idx = harm_bin_start;
                        ctx->harm_params[cur_harm_cnt].end_idx = harm_bin_end;
                        ++cur_harm_cnt;
                    } else {
                        int harm_pos_idx;
                        for (harm_pos_idx = 0; (harm_pos_idx < cur_harm_cnt) &&
                             ctx->harm_params[harm_pos_idx].start_idx < harm_bin_start; ++harm_pos_idx) {
                            continue;
                        }

                        if ((harm_pos_idx > 0) && (ctx->harm_params[harm_pos_idx-1].end_idx >=harm_bin_start)) {
                            harm_bin_start = ctx->harm_params[harm_pos_idx-1].end_idx + 1;
                        }
                        if ((harm_pos_idx < (cur_harm_cnt-1)) && (ctx->harm_params[harm_pos_idx].start_idx <= harm_bin_end)) {
                            harm_bin_end = ctx->harm_params[harm_pos_idx].start_idx - 1;
                        }

                        if (harm_bin_start <= harm_bin_end) {
                            int fnd_pos_idx = harm_pos_idx;
                            for (harm_pos_idx = cur_harm_cnt; harm_pos_idx > fnd_pos_idx; --harm_pos_idx) {
                                ctx->harm_params[harm_pos_idx] = ctx->harm_params[harm_pos_idx-1];
                            }

                            ctx->harm_params[fnd_pos_idx].num = harm_num-1;
                            ctx->harm_params[fnd_pos_idx].start_idx = harm_bin_start;
                            ctx->harm_params[fnd_pos_idx].end_idx = harm_bin_end;
                            ++cur_harm_cnt;
                        }
                    }
                }

                ctx->result_harm_cnt = cur_harm_cnt;
            } else {
                free(ctx);
                ctx = NULL;
            }
        }
    }
    return ctx;
}



static void inline f_add_spectrum_bin(t_lmath_sepctral_component *comp, double amp) {
    comp->maxBinValue = MAX(comp->maxBinValue, amp);
    comp->power += (amp*amp);
}

t_lmath_errs lmath_spectrum_signoise_calc(const t_lmath_spectrum_signoise_ctx ctx,
                                          const double *amp_spectrum,
                                          t_lmath_spectrum_signoise_result *result) {
    t_lmath_errs err = LMATH_ERR_OK;

    if ((amp_spectrum==NULL) || (ctx == NULL) ||  (ctx->harm_params==NULL)) {
        err = LMATH_ERR_INVALID_PARAMS;
    } else {
        int i;
        int harm_pos = 0;
        result->signal.power = result->noise.power = result->harms.power = 0;
        result->signal.maxBinValue = result->noise.maxBinValue = result->harms.maxBinValue = -(double)INFINITY;
        result->params.SNR = result->params.THD = result->params.SINAD =
                result->params.SFDR = result->params.ENOB = 0;

        for(i = ctx->start_idx; i < ctx->fft_size; i++) {
            const t_lmath_spectrum_harm_param *harm_params = &ctx->harm_params[harm_pos];
            double amp = amp_spectrum[i];

            // используется только первые harm_cnt гармоник
            if ((i >= harm_params->start_idx) && (i <= harm_params->end_idx)) {
                if(harm_params->num > 0) {
                    f_add_spectrum_bin(&result->harms, amp);
                } else {
                    f_add_spectrum_bin(&result->signal, amp);
                }

                if ((i == harm_params->end_idx) && (harm_pos < (ctx->result_harm_cnt-1)))
                    ++harm_pos;
            } else {
                f_add_spectrum_bin(&result->noise, amp);
            }
        }

        result->signal.power /= ctx->enbw;
        result->noise.power /= ctx->enbw;
        result->harms.power /= ctx->enbw;

        if(result->signal.power > 0) {
            if (ctx->flags & LMATH_SIGNOISE_FLAGS_CALC_SFDR) {
                double DynamicRange = result->signal.maxBinValue/(MAX(result->harms.maxBinValue,
                                                                 result->noise.maxBinValue));
                if (DynamicRange > 0) {
                    result->params.SFDR = 20 * log10(DynamicRange);
                } else {
                    result->params.SFDR = -(double)INFINITY;
                }
            }

            if (ctx->flags & LMATH_SIGNOISE_FLAGS_CALC_SNR) {
                if (result->noise.power > 0) {
                    double SNRPower =  result->signal.power / result->noise.power;
                    if((SNRPower > 0)) {
                        result->params.SNR = 10 * log10(SNRPower);
                    } else {
                        result->params.SNR = -(double)INFINITY;
                    }
                }
            }

            if (ctx->flags & LMATH_SIGNOISE_FLAGS_CALC_THD) {
                if(result->signal.power > 0) {
                    double THDPower = result->harms.power/result->signal.power;
                    if (THDPower>0) {
                        result->params.THD = 10 * log10(THDPower);
                    } else {
                        result->params.THD = -(double)INFINITY;
                    }
                }
            }

            if ((ctx->flags & LMATH_SIGNOISE_FLAGS_CALC_SINAD) ||
                (ctx->flags & LMATH_SIGNOISE_FLAGS_CALC_ENOB)) {
                if (result->noise.power > 0) {
                    double SINADPower = result->signal.power / (result->noise.power + result->harms.power);
                    if (SINADPower>0) {
                        result->params.SINAD = 10*log10(SINADPower);
                        result->params.ENOB = (result->params.SINAD - 1.76) / 6.02;
                    } else {
                        result->params.SINAD = -(double)INFINITY;
                    }
                }
            }
        }
    }
    return err;
}

void lmath_spectrum_signoise_ctx_destroy(t_lmath_spectrum_signoise_ctx ctx) {
    if (ctx) {
        free(ctx->harm_params);
        free(ctx);
    }
}

t_lmath_errs lmath_spectrum_signoise(const double *amp_spectrum, int fft_size, double df,
                                          const t_lmath_window_params *winpar,
                                          double peak_freq,
                                          int harm_cnt, unsigned flags,
                                          t_lmath_spectrum_signoise_result *result) {
    t_lmath_errs err = LMATH_ERR_OK;
    t_lmath_spectrum_signoise_ctx ctx = lmath_spectrum_signoise_ctx_create(
                fft_size, df, winpar, peak_freq, harm_cnt, flags);
    if (ctx) {
        err = lmath_spectrum_signoise_calc(ctx, amp_spectrum, result);
        lmath_spectrum_signoise_ctx_destroy(ctx);
    } else {
        err = LMATH_ERR_INVALID_PARAMS;
    }
    return err;
}
