#include "calculation.hpp"

Calculation::Calculation() {}

Calculation::~Calculation() {}

/************************************************************************/
/* 一维数据的复数快速傅里叶变换                                             */
/************************************************************************/
void Calculation::FFT(int n, fftw_complex *in, fftw_complex *out) {
    if (in == NULL || out == NULL) return;
    fftw_plan p;
    p = fftw_plan_dft_1d(n, in, out, FFTW_FORWARD, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
    fftw_cleanup();
}

/************************************************************************/
/* 一维数据的实数快速傅里叶变换                                             */
/************************************************************************/
void Calculation::FFT_R(int n, std::vector<float> &vecData, fftw_complex *out) {
    double in[n];
    for (int i = 0; i < n; i++) {
        in[i] = vecData[i];
    }

    for (int i = 0; i < n; i++) {
        out[i][0] = (double)vecData[i];
        out[i][1] = 0;
    }
    // create a DFT plan and execute it
    fftw_plan plan = fftw_plan_dft_r2c_1d(n, in, out, FFTW_ESTIMATE);
    fftw_execute(plan);
    // destroy the plan to prevent a memory leak
    fftw_destroy_plan(plan);
    fftw_cleanup();
}

/************************************************************************/
/* 一维数据的快速傅里叶逆变换                                          */
/************************************************************************/
void Calculation::iFFT(int n, fftw_complex *in, fftw_complex *out) {
    if (in == NULL || out == NULL) return;
    fftw_plan p;
    p = fftw_plan_dft_1d(n, in, out, FFTW_BACKWARD, FFTW_ESTIMATE);
    fftw_execute(p);
    fftw_destroy_plan(p);
    fftw_cleanup();
}

void Calculation::_iFFT(std::vector<float> &vecrealData, std::vector<float> &vecimageData, std::vector<float> &veciFFTData) {
    fftw_complex *inFFt, *outFFt;
    int N = vecrealData.size();
    inFFt = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * N);
    outFFt = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * N);

    for (int j = 0; j < N; j++) {
        inFFt[j][0] = (double)vecrealData[j];
        inFFt[j][1] = (double)vecimageData[j];
    }
    iFFT(N, inFFt, outFFt);
    for (int i = 0; i < N; i++) {
        outFFt[i][0] *= 1. / N;
        outFFt[i][1] *= 1. / N;
        veciFFTData.push_back(outFFt[i][0]);
    }
    fftw_free(inFFt);
    fftw_free(outFFt);
}

void Calculation::_FFT(std::vector<float> &vecData, std::vector<float> &vecFFTrealData, std::vector<float> &vecFFTimageData) {
    fftw_complex *inHilFFt, *outHilFFt;
    inHilFFt = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * vecData.size());
    outHilFFt = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * vecData.size());

    for (size_t j = 0; j < vecData.size(); j++) {
        inHilFFt[j][0] = (double)vecData[j];
        inHilFFt[j][1] = 0;
    }

    FFT(vecData.size(), inHilFFt, outHilFFt);
    // fftShift(outHilFFt,  vecData.size());
    for (size_t i = 0; i < vecData.size(); i++) {
        vecFFTrealData.push_back(outHilFFt[i][0]);
        vecFFTimageData.push_back(outHilFFt[i][1]);
    }
    fftw_free(inHilFFt);
    fftw_free(outHilFFt);
}

float Calculation::mean(std::vector<float> &vecData) {
    float meanTemp = 0;
    for (size_t i = 0; i < vecData.size(); i++) {
        meanTemp += vecData[i];
    }
    return meanTemp / vecData.size();
}

float Calculation::getSample_variance(std::vector<float> a) {
    float ss = 0;
    float s = 0;
    float mx = mean(a);
    int len = a.size();
    for (int i = 0; i < len; i++) {
        s = a[i] - mx;
        ss += pow(s, 2);
    }
    return ss / (len - 1);
}

void Calculation::Hanning(std::vector<float> &vecData, std::vector<float> &vecHanningData) {
    int N = vecData.size();

    float *w = NULL;
    w = (float *)calloc(N, sizeof(float));
    int half, i, idx;
    if (N % 2 == 0) {
        half = N / 2;
        for (i = 0; i < half; i++) // CALC_HANNING   Calculates Hanning window samples.
            w[i] = 0.5 * (1 - cos(2 * M_PI * (i + 1) / (N + 1)));

        idx = half - 1;
        for (i = half; i < N; i++) {
            w[i] = w[idx];
            idx--;
        }
    } else {
        half = (N + 1) / 2;
        for (i = 0; i < half; i++) // CALC_HANNING   Calculates Hanning window samples.
            w[i] = 0.5 * (1 - cos(2 * M_PI * (i + 1) / (N + 1)));

        idx = half - 2;
        for (i = half; i < N; i++) {
            w[i] = w[idx];
            idx--;
        }
    }
    for (int j = 0; j < N; j++) {
        vecHanningData.push_back(w[j]);
    }
    free(w);
}
void Calculation::hilbert(std::vector<float> &vecData, std::vector<float> &vecHilbertData, int N) {
    double in[N];
    for (int i = 0; i < N; i++) {
        in[i] = vecData[i];
    }
    fftw_complex *out;
    out = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * N);

    for (int i = 0; i < N; ++i) {
        out[i][0] = (double)vecData[i];
        out[i][1] = 0;
    }
    // create a DFT plan and execute it
    fftw_plan plan = fftw_plan_dft_r2c_1d(N, in, out, FFTW_ESTIMATE);
    fftw_execute(plan);
    // destroy the plan to prevent a memory leak
    fftw_destroy_plan(plan);
    int hN = N >> 1; // half of the length (N /2)
    int numRem = hN; // the number of remaining elements
    // multiply the appropriate values by 2
    // (those that should be multiplied by 1 are left intact because they wouldn't change)
    for (int i = 1; i < hN; ++i) // 1,2,...,N/2 - 1 的项乘以2
    {
        out[i][0] *= 2;
        out[i][1] *= 2;
    }
    // if the length is even, the number of remaining elements decreases by 1
    if (N % 2 == 0) numRem--;
    // if it's odd and greater than 1, the middle value must be multiplied by 2
    else if (N > 1) // 奇数非空
    {
        out[hN][0] *= 2;
        out[hN][1] *= 2;
    }
    // set the remaining values to 0
    // (multiplying by 0 gives 0, so we don't care about the multiplicands)
    memset(&out[hN + 1][0], 0, numRem * sizeof(fftw_complex));
    // create an IDFT plan and execute it
    plan = fftw_plan_dft_1d(N, out, out, FFTW_BACKWARD, FFTW_ESTIMATE);
    fftw_execute(plan);
    // do some cleaning
    fftw_destroy_plan(plan);
    fftw_cleanup();
    // scale the IDFT output
    for (int i = 0; i < N; ++i) {
        out[i][0] /= N;
        out[i][1] /= N;
    }

    for (int n = 0; n < N; n++) //输出
    {
        //  xr[n]=cos(n*pi/6);//原始信号
        //  y_r[n] = s_i[n];
        complex complex_after;
        complex_after.real = out[n][1];
        complex_after.imag = out[n][0];
        float amp = sqrt(complex_after.real * complex_after.real + complex_after.imag * complex_after.imag);
        vecHilbertData.push_back(amp);
        //  printf("%d    %f\n",n,vecHilbertData[n]);
    }
    fftw_free(out);
}

void Calculation::fftShift(fftw_complex *in, int l) {
    double temp;
    double temp2;

    for (int j = 0; j < l / 2; j++) {
        temp = in[j + l / 2][0];
        temp2 = in[j + l / 2][1];
        in[j + l / 2][0] = in[j][0];
        in[j + l / 2][1] = in[j][1];
        in[j][0] = temp;
        in[j][1] = temp2;
    }
}

void Calculation::FFTSpec(std::vector<float> &vecData, std::vector<float> &vecFFTSpecData) {
    fftw_complex *inFFt, *outFFt;
    inFFt = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * vecData.size());
    outFFt = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * vecData.size());

    for (size_t j = 0; j < vecData.size(); j++) {
        inFFt[j][0] = (double)vecData[j];
        inFFt[j][1] = 0;
    }
    FFT(vecData.size(), inFFt, outFFt);
    for (size_t j = 0; j < vecData.size() / 2; j++) {
        vecFFTSpecData.push_back(sqrt(outFFt[j][0] * outFFt[j][0] + outFFt[j][1] * outFFt[j][1]) * 2 / vecData.size());
    }
}

void Calculation::envSpec(std::vector<float> &vecData, std::vector<float> &vecEnvSpecData, int StartFrequency, int EndFrequency) {
    std::vector<float> vecFFTrealData, vecFFTimageData;
    std::vector<float> vecRealData, vecImageData;
    std::vector<float> veciFFtData;
    std::vector<float> veciFFtData2;
    std::vector<float> vecHilbertData;
    _FFT(vecData, vecFFTrealData, vecFFTimageData);
    for (size_t i = 0; i < vecFFTrealData.size(); i++) {
        if (i > (size_t)StartFrequency && i < (size_t)EndFrequency) {
            vecRealData.push_back(vecFFTrealData.at(i));
            vecImageData.push_back(vecFFTimageData.at(i));
        } else {
            vecRealData.push_back(0);
            vecImageData.push_back(0);
        }
    }
    _iFFT(vecRealData, vecImageData, veciFFtData);
    for (size_t j = 0; j < veciFFtData.size(); j++) {
        veciFFtData2.push_back(veciFFtData[j] * 2);
    }
    hilbert(veciFFtData2, vecHilbertData, veciFFtData2.size());
    FFTSpec(vecHilbertData, vecEnvSpecData);
}

void Calculation::GenerateSin(std::vector<float> &vecData) {
    double frequency = 800.0;        // Frequency of the sine wave in Hz
    double sampling_rate = 12800.0;  // Sampling rate in Hz (8 kHz)
    size_t num_samples = 12800;      // Total number of samples (1 second of data)
    double dt = 1.0 / sampling_rate; // Time step in seconds

    // Generate the sine wave
    for (size_t i = 0; i < num_samples; ++i) {
        vecData.push_back(std::sin(2 * M_PI * frequency * i * dt));
    }
}

void Calculation::Integration(std::vector<float> &vecData, std::vector<float> &retData, double &resolution) {
    std::vector<float> realshiftfft;
    std::vector<float> imageshiftfft;
    std::vector<float> realvalue, imagevalue, ifftData;
    _FFT(vecData, realshiftfft, imageshiftfft);

    //在频域上进行5-1000 Hz的带通滤波
    for (size_t i = 0; i < vecData.size() / 2 + 1; ++i) {
        double frequency = i * resolution; // 计算当前频率分量
        if (frequency < 5 || frequency > 1000) {
            // 将5 Hz 到 1000 Hz 之外的频率成分设置为0
            realshiftfft[i] = 0.0;
            realshiftfft[i] = 0.0;
            imageshiftfft[i] = 0.0;
            imageshiftfft[i] = 0.0;
        }
    }

    for (size_t k = 1; k < realshiftfft.size() + 1; k++) {
        double frequency = k * resolution;                                                   // 计算当前频率分量
        realvalue.push_back((realshiftfft.at(k - 1) / (frequency * 2 * M_PI)) * 1000 * 2);   //单位转换mm/s,*1000  *2  精度损失
        imagevalue.push_back((imageshiftfft.at(k - 1) / (frequency * 2 * M_PI)) * 1000 * 2); //单位转换mm/s,*1000
    }

    _iFFT(realvalue, imagevalue, retData);
}
std::vector<float> Calculation::fftInterpolate(const std::vector<float>& input, size_t outputSize) {
    size_t inputSize = input.size();
    double in[inputSize];
    for (size_t i = 0; i < inputSize; i++) {
        in[i] = input[i];
    }
    // 1. FFTW 初始化
    fftw_complex *freqDomain = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * inputSize);
    fftw_complex *paddedFreqDomain = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * outputSize);
    fftw_plan forwardPlan = fftw_plan_dft_r2c_1d(inputSize, in, freqDomain, FFTW_ESTIMATE);
    
    // 2. 执行 FFT
    fftw_execute(forwardPlan);

    // 3. 频域插值：扩展频谱
    size_t halfSize = inputSize / 2 + 1; // 实数FFT的对称部分
    for (size_t i = 0; i < halfSize; ++i) {
        paddedFreqDomain[i][0] = freqDomain[i][0]; // 实部
        paddedFreqDomain[i][1] = freqDomain[i][1]; // 虚部
    }
    for (size_t i = halfSize; i < outputSize - halfSize; ++i) {
        paddedFreqDomain[i][0] = 0.0; // 实部填零
        paddedFreqDomain[i][1] = 0.0; // 虚部填零
    }
    for (size_t i = outputSize - halfSize; i < outputSize; ++i) {
        paddedFreqDomain[i][0] = freqDomain[inputSize - (outputSize - i)][0];
        paddedFreqDomain[i][1] = freqDomain[inputSize - (outputSize - i)][1];
    }

    // 4. IFFT 变换回时域
    std::vector<double> output(outputSize);
    fftw_plan inversePlan = fftw_plan_dft_c2r_1d(outputSize, paddedFreqDomain, output.data(), FFTW_ESTIMATE);
    fftw_execute(inversePlan);
    
    // 5. 缩放输出结果
    for (double& val : output) {
        val /= outputSize;
    }
    std::vector<float> output2(outputSize);
    for (size_t i = 0; i < outputSize; i++) {
        output2[i] = output[i];
    }
    // 清理 FFTW
    fftw_destroy_plan(forwardPlan);
    fftw_destroy_plan(inversePlan);
    fftw_free(freqDomain);
    fftw_free(paddedFreqDomain);

    return output2;
}

bool Calculation::freqDomainDecimateFFTW(const std::vector<float>& input,
    std::vector<float>& output,
    int factor)
{
    int N = input.size();

    if(N % factor != 0)
    {
        std::cout<<"Length not divisible by factor\n";
        return false;
    }

    int N_out = N / factor;

    // ---------- 分配 ----------
    double* in_time  = (double*)fftw_malloc(sizeof(double)*N);
    fftw_complex* spectrum =
    (fftw_complex*)fftw_malloc(sizeof(fftw_complex)*(N/2+1));

    for(int i=0;i<N;i++)
    in_time[i] = input[i];

    // ---------- FFT ----------
    fftw_plan plan_fwd =
    fftw_plan_dft_r2c_1d(N,in_time,spectrum,FFTW_ESTIMATE);

    fftw_execute(plan_fwd);

    // ---------- 构建新频谱 ----------
    fftw_complex* new_spec =
    (fftw_complex*)fftw_malloc(sizeof(fftw_complex)*(N_out/2+1));

    memset(new_spec,0,sizeof(fftw_complex)*(N_out/2+1));

    int half = N_out/2;

    for(int i=0;i<=half;i++)
    {
    new_spec[i][0] = spectrum[i][0];
    new_spec[i][1] = spectrum[i][1];
    }

    // ---------- IFFT ----------
    double* out_time =
    (double*)fftw_malloc(sizeof(double)*N_out);

    fftw_plan plan_inv =
    fftw_plan_dft_c2r_1d(N_out,new_spec,out_time,FFTW_ESTIMATE);

    fftw_execute(plan_inv);

    output.resize(N_out);

    for(int i=0;i<N_out;i++)
    output[i] = out_time[i] / N_out;

    // ---------- 清理 ----------
    fftw_destroy_plan(plan_fwd);
    fftw_destroy_plan(plan_inv);

    fftw_free(in_time);
    fftw_free(out_time);
    fftw_free(spectrum);
    fftw_free(new_spec);

    return true;
}