#include <cmath>
#include <vector>
#include <complex>
#include "const.h"
#include "FreqResponse.h"
#include "FFTHelper.h"

void FreqResponse::iir_freqz(double b[3], double a[3], int N, double* w, double* m)
{
    double num, den;
    int i;
    double div;

    for (i = 0; i < N; i++)
    {
        num = b[0] * b[0] + b[1] * b[1] + b[2] * b[2] + 2 * b[0] * b[1] * w[i] + 2 * b[1] * b[2] * w[i] + 2 * b[0] * b[2] * (2 * w[i] * w[i] - 1);
        den = a[0] * a[0] + a[1] * a[1] + a[2] * a[2] + 2 * a[0] * a[1] * w[i] + 2 * a[1] * a[2] * w[i] + 2 * a[0] * a[2] * (2 * w[i] * w[i] - 1);

        m[i] = sqrt(num / den);
    }
}

void FreqResponse::iir_anglez(double b[3], double a[3], int N, int fs, double* f, double* p)
{
    int i;
    double* w = new double[N];
    double a0, b0, c0, d0;

    w[0] = f[0] = 0;
    for (i = 1; i < N; i++)
    {
        w[i] = w[i - 1] + pi / N;
        f[i] = w[i] * fs / 2 / pi;
    }

    for (i = 0; i < N; i++)
    {
        a0 = b[0] + b[1] * cos(w[i]) + b[2] * cos(2 * w[i]);
        c0 = a[0] + a[1] * cos(w[i]) + a[2] * cos(2 * w[i]);

        b0 = -(b[1] * sin(w[i]) + b[2] * sin(2 * w[i]));
        d0 = -(a[1] * sin(w[i]) + a[2] * sin(2 * w[i]));

        p[i] = atan2(b0 * c0 - a0 * d0, a0 * c0 + b0 * d0);
    }
    delete[] w;
}

int FreqResponse::freqz(double coeffs[], int size, int smpl, int fs, double f[], double gain[], double angle[])
{
    int i;
    int n = smpl;

    if (size > smpl)
    {
        return 0;
    }
    if ((n & (n - 1)) > 0)
    {
        return 0;
    }
    double* re_x = new double[n];
    double* im_x = new double[n];
    for (i = 0; i < size; i++)
    {
        re_x[i] = (float)coeffs[i];
        im_x[i] = 0;
    }
    for (; i < n; i++)
    {
        re_x[i] = 0;
        im_x[i] = 0;
    }

    FFTHelper fft;
    fft.FFT(re_x, im_x, n);

    double interval = fs*1.0 / n;
    for (i = 0; i <= n / 2; i++)
    {
        //f[i] = i * interval;  //fÍâ²ãÒѾ­Ìî³äƵÂÊ
        gain[i] = 20 * log10(sqrt(re_x[i] * re_x[i] + im_x[i] * im_x[i]));
        angle[i] = atan2(im_x[i], re_x[i]); //-pi~pi
    }
    delete[] re_x;
    delete[] im_x;

    return n / 2 +1;
}

int FreqResponse::freqz(double(*sos)[6], int section, int smpl, int fs, double f[], double gain[])
{
    int i, j = 1;
    double a[3],b[3];
    int N = smpl;
    double* w = new double[N];
    double* m = new double[N];

    for (i = 0; i < N; i++){
        w[i] = cos(pi * 2 * f[i] / fs);
    }

    b[0] = sos[0][0]; b[1] = sos[0][1]; b[2] = sos[0][2];
    a[0] = sos[0][3]; a[1] = sos[0][4]; a[2] = sos[0][5];

    iir_freqz(b, a, N, w, gain);

    for (i = 1; i < section; i++) {
        b[0] = sos[i][0]; b[1] = sos[i][1]; b[2] = sos[i][2];
        a[0] = sos[i][3]; a[1] = sos[i][4]; a[2] = sos[i][5];

        iir_freqz(b, a, N, w, m);

        for (j = 0; j < N; j++) {
            gain[j] *= m[j];
        }
    }

    for (i = 0; i < N; i++){
        if (isnan(gain[i]) || isinf(gain[i])) {
            gain[i] = 1e-10;
        }
        gain[i] = 20 * log10(gain[i]);
       
    }
    delete[]w;
    delete[]m;

    return 0;
}

int FreqResponse::extra(double f[], double gain[], double phase[], int num
    , double new_f[], double new_gain[], double new_phase[], int new_num)
{
    for (int i = 0; i < new_num; i++) {
        int j = 0;
        double target_f = new_f[i];
        if (target_f < f[0]) {
            // µÍÓÚ×îСƵÂÊ - ʹÓõÚÒ»¸öÇø¼äµÄÍâÍÆ
            //double log_f0 = log10(f[0]);
            //double log_f1 = log10(f[1]);
            //double log_target = log10(new_f[i]);

            //double t = (log_target - log_f0) / (log_f1 - log_f0);
            //new_gain[i] = gain[0] + t * (gain[1] - gain[0]);
            //new_phase[i] = phase[0] + t * (phase[1] - phase[0]);
            if(new_gain)
                new_gain[i] = gain[0] ;
            if (new_phase)
                new_phase[i] = phase[0] ;
        }
        else if (target_f >= f[num - 1]) {
            // ¸ßÓÚ×î´óƵÂÊ - ʹÓÃ×îºóÒ»¸öÇø¼äµÄÍâÍÆ
            //double log_fn2 = log10(f[num - 2]);
            //double log_fn1 = log10(f[num - 1]);
            //double log_target = log10(target_f);

            //double t = (log_target - log_fn2) / (log_fn1 - log_fn2);
            //new_gain[i] = gain[num - 2] + t * (gain[num - 1] - gain[num - 2]);
            //new_phase[i] = phase[num - 2] + t * (phase[num - 1] - phase[num - 2]);
            if (new_gain)
                new_gain[i] = gain[0];
            if (new_phase)
                new_phase[i] = phase[0];
        }
        else {
            // ÔÚ·¶Î§ÄڵIJåÖµ
            int idx = 0;
            // ÕÒµ½Ä¿±êƵÂÊËùÔÚµÄÇø¼ä
            for (int j = 0; j < num - 1; j++) {
                if (target_f >= f[j] && target_f <= f[j + 1]) {
                    idx = j;
                    break;
                }
            }

            // ÔÚ¶ÔÊý³ß¶ÈÉϽøÐÐÏßÐÔ²åÖµ
            double log_f_prev = log10(f[idx]);
            double log_f_next = log10(f[idx + 1]);
            double log_target = log10(target_f);

            double t = (log_target - log_f_prev) / (log_f_next - log_f_prev);
            if (new_gain)
                new_gain[i] = gain[idx] + t * (gain[idx + 1] - gain[idx]);
            if (new_phase)
                new_phase[i] = phase[idx] + t * (phase[idx + 1] - phase[idx]);
        }
    }
    
    return ret_success;
}

int FreqResponse::logspace(double f1, double f2, double f[], int num)
{
    double log_min = log10(f1);
    double log_max = log10(f2);
    double log_range = log_max - log_min;

    for (int i = 0; i < num; i++) {
        double log_freq = log_min + (log_range * i) / (num - 1);
        // ת»»»ØÏßÐÔ¿Õ¼ä
        f[i] = pow(10, log_freq);
    }

    return ret_success;
}

int FreqResponse::linearspace(double f1, double f2, double f[], int num)
{
    // ¼ÆËãƵÂʼä¸ô
    double step = (f2 - f1) / (num - 1);

    // Éú³ÉÏßÐÔƵÂÊÊý×é
    for (int i = 0; i < num; i++) {
        f[i] = f1 + i * step;
    }

    return ret_success;
}

//void FreqResponse::hn_energy_center(double* ifft_real, int real_size, double* hn, int tap)
//{
//    double energy_sum = 0.0;
//    double energy_center = 0.0;
//
//    // ¼ÆËãÄÜÁ¿ÖÐÐÄ
//    for (int n = 0; n < real_size; n++) {
//        double e = ifft_real[n] * ifft_real[n];
//        energy_sum += e;
//        energy_center += n * e;
//    }
//    energy_center /= energy_sum;
//
//    // ½ØÈ¡Æðµã
//    int start = (int)round(energy_center - tap / 2.0);
//    if (start < 0) start = 0;
//    if (start + tap > real_size) start = real_size - tap;
//    start = 0;
//    // ¿½±´FIR
//    for (int i = 0; i < tap; i++) {
//        hn[i] = ifft_real[start + i];
//    }
//}

int FreqResponse::impulse_response(double* f, double* g, double* p, int num
    , double* hn, int tap , double* wn, double delay_ms,bool phase_comp)
{
    int fft_size = 2*(num-1);
    double tau0 = phase_comp? delay_ms*(KERNEL_SAMPLE_RATE/1000) :(tap-1)/2;  // Ä¿±êȺÑÓʱÑù±¾Êý

    double tau_c = estimate_avg_group_delay(g, p, num); //ƽ¾ùȺÑÓʱÑù±¾Êý
    double delta_tau = tau0 - tau_c; //²Ð²î.

    // Éú³É¸´ÊýƵÂÊÏìÓ¦
    double* H_real = (double*)malloc(fft_size * sizeof(double));
    double* H_imag = (double*)malloc(fft_size * sizeof(double));

    for (int i = 0; i < num; i++) {
        double magnitude = pow(10.0, g[i] / 20.0);
        double phi;
        
        if (phase_comp) {
            phi = p[i];
        }
        else {
            //¹¹ÔìÏßÐÔÏàλFIR.
            double omega = i * pi / (fft_size / 2);
            phi = -tau0 * omega;
        }

        // ¼ÆËãʵ²¿ºÍÐ鲿
        H_real[i] = magnitude * cos(phi);
        H_imag[i] = magnitude * sin(phi);

        //Ïàλ²¹³¥
        if (phase_comp) {
            double phase = -2.0 * pi * i * delta_tau / fft_size;
            double a = H_real[i], b = H_imag[i];
            double c = cos(phase), s = sin(phase);
            H_real[i] = a * c - b * s;
            H_imag[i] = a * s + b * c;
        }
    }

    // È·±£¹²éî¶Ô³ÆÐÔ
    for (int i = 1; i < fft_size / 2; i++) {
        H_real[fft_size - i] = H_real[i];      // ʵ²¿¶Ô³Æ
        H_imag[fft_size - i] = -H_imag[i];     // Ð鲿·´¶Ô³Æ
    }

    FFTHelper().IFFT(H_real, H_imag, fft_size);
   
    //hn_energy_center(H_real, fft_size, hn, tap);
    int s = (int)round(tau0 - (tap - 1) / 2.0);
    if (s < 0) s = 0;
    if (s + tap > fft_size) s = fft_size - tap;

    for (int i = 0; i < tap; i++) {
        hn[i] = H_real[s + i] * wn[i];
    }

    //¹éÒ»»¯?
    
    // ÊÍ·ÅÄÚ´æ
    free(H_real);
    free(H_imag);

    return 0;
}

void FreqResponse::unwrap_phase(double* phi, int N) {
    double prev = phi[0];
    for (int k = 1; k < N; ++k) {
        double cur = phi[k];
        double d = cur - prev;
        if (d > pi) {
            /* cur wrapped positively -> subtract 2pi */
            while (d > pi) {
                cur -= 2 * pi;
                d = cur - prev;
            }
        }
        else if (d < -pi) {
            /* cur wrapped negatively -> add 2pi */
            while (d < -pi) {
                cur += 2 * pi;
                d = cur - prev;
            }
        }
        phi[k] = cur;
        prev = cur;
    }
}

void FreqResponse::unwrap_phase(double phase[], int num, double unwrapped_phase[])
{
    unwrapped_phase[0] = phase[0];
    double cumulative_shift = 0.0;

    for (int i = 1; i < num; i++) {
        double delta = phase[i] - phase[i - 1];

        // ¼ì²â²¢ÀÛ»ý2¦ÐÌø±ä
        if (delta > pi) {
            cumulative_shift -= 2.0 * pi;
        }
        else if (delta < -pi) {
            cumulative_shift += 2.0 * pi;
        }

        unwrapped_phase[i] = phase[i] + cumulative_shift;
    }
}


void FreqResponse::sw_smooth(double oct, double f1,double f2, double* freq, double* data, double* out_data ,int num)
{
    int start_idx = 0;
    double octave_ratio = pow(2.0, oct);
    double freq_unit = freq[1] - freq[0];

    for (int i = 0; i < num; i++) {
        double center_freq = freq[i];
 
        if (center_freq < f1 || center_freq > f2) {
            continue;
        }
        double lower_freq = center_freq / octave_ratio;
        double upper_freq = center_freq * octave_ratio;

        int lower_freq_index = (int)floor(lower_freq / freq_unit);
        int upper_freq_index = (int)ceil(upper_freq / freq_unit);

        // ¼ÆËã´°¿ÚÄڵĵã
        double sum_data = 0.0;
        int count = upper_freq_index - lower_freq_index;

        for (int j = lower_freq_index; j < upper_freq_index; j++) {
            sum_data += data[j];
        }

        if (count > 0) {
            out_data[i] = sum_data / count;
        }
        else {
            out_data[i] = data[i];
        }
    }
}

double FreqResponse::estimate_avg_group_delay(const double* gain, const double* phi, int N)
{
    /* compute magnitude^2 */
    double* mag2 = new double[N];
    double max_mag2 = 0.0;
    for (int k = 0; k < N; ++k) {
        mag2[k] = pow(10, gain[k] / 20);
        if (mag2[k] > max_mag2) max_mag2 = mag2[k];
    }

    double eps = max_mag2 * 1e-6;

    double domega = 2.0 * pi / N;

    /* ÿ¸öƵÂʵãȺÑÓʱ */
    double* taug = new double[N];
    for (int k = 1; k < N - 1; ++k) {
        double dph = phi[k + 1] - phi[k - 1];
        taug[k] = -dph / (2.0 * domega);
    }
    taug[0] = -(phi[1] - phi[0]) / domega;
    taug[N - 1] = -(phi[N - 1] - phi[N - 2]) / domega;

    /* ¼ÆËãƽ¾ùȺÑÓʱ */
    double num = 0.0, den = 0.0;
    for (int k = 0; k < N; ++k) {
        if (mag2[k] > eps) {
            num += taug[k] * mag2[k];
            den += mag2[k];
        }
    }

    double tau_c = 0.0;
    if (den > 0.0) tau_c = num / den;
    else tau_c = (N - 1) / 2.0; /* fallback: geometric center */

    delete []mag2; delete []taug;
    return tau_c;
}

void FreqResponse::hn_group_delay_center(double* ifft_real, int real_size, double* hn, int tap)
{
    
}

void FreqResponse::mulimpulse(pulse_t* pulse, int pulse_num, int pulse_start)
{
    double* real_part = new double[KERNEL_HALF_FFT];
    double* imag_part = new double[KERNEL_HALF_FFT];
    pulse_t& output = pulse[pulse_num - 1];

    // ³õʼ»¯
    for (int i = 0; i < KERNEL_HALF_FFT; i++) {
        real_part[i] = 1.0;
        imag_part[i] = 0.0;
    }

    for (int p = 0; p < pulse_num -1; p++) {
        if (!pulse[p].enable || p < pulse_start) continue;

        for (int i = 0; i < KERNEL_HALF_FFT; i++) {
            double mag = std::pow(10.0, pulse[p].gain[i] / 20.0);
            double current_real = mag * std::cos(pulse[p].phase[i]);
            double current_imag = mag * std::sin(pulse[p].phase[i]);

            // ¸´Êý³Ë·¨
            double new_real = real_part[i] * current_real - imag_part[i] * current_imag;
            double new_imag = real_part[i] * current_imag + imag_part[i] * current_real;

            real_part[i] = new_real;
            imag_part[i] = new_imag;
        }
    }
    for (int i = 0; i < KERNEL_HALF_FFT; i++) {
        // ¼ÆËã·ù¶ÈºÍÏàλ
        double magnitude = std::sqrt(real_part[i] * real_part[i] + imag_part[i] * imag_part[i]);
        output.gain[i] = 20.0 * std::log10(magnitude);
        output.phase[i] = std::atan2(imag_part[i], real_part[i]);
    }
    delete[]real_part;
    delete[]imag_part;
}