#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <locale.h>
#include <float.h>
#include "emg_util.h"
#include <assert.h>

// SOS coefficients (20-450 Hz bandpass, fs=1000Hz)
// 将常量数组类型从 double 改为 float，并添加 'f' 后缀
const float sos[2][6] = { // double -> float
    {1.0000000000f,  2.0000000000f, 1.0000000000f, 1.0f,  -1.750100000f, 0.7797000000f}, // 所有数值添加 f 后缀
    {1.0000000000f, -2.0000000000f, 1.0000000000f, 1.0f, -1.9912000000f, 0.9912000000f}  // 所有数值添加 f 后缀
};

// Notch filter parameters (50Hz notch, fs=1000Hz)
// 将系数类型从 double 改为 float，并添加 'f' 后缀
const float NOTCH_B[3] = {0.9955f, -1.8936f, 0.9955f}; // double -> float, 添加 f 后缀
const float NOTCH_A[3] = {1.0000f, -1.8936f, 0.9911f}; // double -> float, 添加 f 后缀

typedef struct {
    float b[3];   // Numerator coefficients: b0, b1, b2 (double -> float)
    float a[3];   // Denominator coefficients: a0, a1, a2 (a0 is typically normalized to 1) (double -> float)
    float x1, x2; // Input delay units: x[n-1], x[n-2] (double -> float)
    float y1, y2; // Output delay units: y[n-1], y[n-2] (double -> float)
} NotchFilter;

// 这些变量如果是在全局范围内使用，建议加上 extern 声明并在某个.c文件中定义。
extern float f0;          // Power line frequency (50Hz) (double -> float)
extern float q;            // Quality factor (double -> float)
extern float low_cut;      // Low cutoff frequency (double -> float)
extern float high_cut;     // High cutoff frequency (double -> float)
extern int order;          // Filter order (保持不变)

// Butterworth filter structure (现在为单精度 float)
typedef struct {
    float x1, x2; // double -> float
    float y1, y2; // double -> float
} BiquadState;

// Symlet 4 (sym4) wavelet filter coefficients (normalized)
// 将系数类型从 double 改为 float，并添加 'f' 后缀
static float sym4_h[8] = { // double -> float
     0.0266700579010166f,
     0.0607514995432352f,
    -0.0656160944756398f,
    -0.1498879107562273f,
     0.4829629131445341f,
     0.8365163037378608f,
     0.2242848007341553f,
    -0.1077300958383618f
};

static float sym4_g[8] = { // double -> float
    -0.1077300958383618f,
    -0.2242848007341553f,
     0.8365163037378608f,
    -0.4829629131445341f,
    -0.1498879107562273f,
     0.0656160944756398f,
     0.0607514995432352f,
    -0.0266700579010166f
};

/**
 * Reset notch filter states
 * Function: Clear the input and output history values of the filter to ensure new filtering process is not affected by previous states
 * Parameter: filter - pointer to the notch filter structure
 */
void notchFilterReset(NotchFilter *filter) {
    // Reset input history (x1: previous input, x2: input from two steps ago)
    filter->x1 = filter->x2 = 0.0f; // 改为 0.0f
    // Reset output history (y1: previous output, y2: output from two steps ago)
    filter->y1 = filter->y2 = 0.0f; // 改为 0.0f
}

/**
 * Process a single sample with notch filter
 * Function: Filter a single input data sample based on IIR filter difference equation
 * Parameters:
 *   filter - pointer to notch filter structure (contains coefficients and historical states)
 *   input  - current input sample data
 * Return: Filtered output data
 */
float notchFilterProcess(NotchFilter *filter, float input) { // double -> float
    // Calculate current output according to IIR filter difference equation
    // Formula: y[n] = b0*x[n] + b1*x[n-1] + b2*x[n-2] - a1*y[n-1] - a2*y[n-2]
    float output = // double -> float
        filter->b[0] * input +         // Current input term
        filter->b[1] * filter->x1 +    // Previous input term
        filter->b[2] * filter->x2 -    // Input term from two steps ago
        filter->a[1] * filter->y1 -    // Previous output feedback term
        filter->a[2] * filter->y2;     // Output feedback term from two steps ago

    // Update input history (shift operation,保存最新的两个输入值)
    filter->x2 = filter->x1;  // Two steps ago = previous step
    filter->x1 = input;       // Previous step = current input

    // Update output history (shift operation,保存最新的两个输出值)
    filter->y2 = filter->y1;  // Two steps ago = previous step
    filter->y1 = output;      // Previous step = current output

    return output;
}

/**
 * Process array data with notch filter
 * Function: Continuously filter an entire segment of input data by processing samples one by one
 * Parameters:
 *   filter - pointer to notch filter structure
 *   input  - pointer to input data array
 *   output - pointer to output data array (memory should be pre-allocated)
 *   length - length of input/output data (number of samples)
 */
void notchFilterProcessArray(NotchFilter *filter, const float *input, float *output, int length) { // double -> float
    // Iterate through input array and filter each sample
    int i;
    for (i = 0; i < length; i++) {
        output[i] = notchFilterProcess(filter, input[i]);
    }
}

// Apply a biquad section
float biquad_filter(float input, const float *sos_row, BiquadState *state) { // double -> float
    float b0 = sos_row[0]; // double -> float
    float b1 = sos_row[1];
    float b2 = sos_row[2];
    float a1 = sos_row[4];  // Note index: the 5th element is a1
    float a2 = sos_row[5];  // The 6th element is a2
    float output; // double -> float
    output = b0 * input + b1 * state->x1 + b2 * state->x2
                             - a1 * state->y1 - a2 * state->y2;

    state->x2 = state->x1;
    state->x1 = input;
    state->y2 = state->y1;
    state->y1 = output;

    return output;
}

// Reverse array
void reverse_array(float *data, int n) { // double -> float
    int i;
    for (i = 0; i < n / 2; i++) {
        float temp = data[i]; // double -> float
        data[i] = data[n - 1 - i];
        data[n - 1 - i] = temp;
    }
}

/**
 * @brief Zero-phase bandpass filtering (similar to MATLAB filtfilt)
 * @param indata: Input signal array
 * @param outdata: Output signal array (filtered result)
 * @param len: Signal length
 * @param sos: 2x6 second-order sections coefficient matrix
 * @param num_stages: Number of sections (2 here)
 */
void filtfilt_sos(const float *indata, float *outdata, int len, const float sos[2][6], int num_stages) { // double -> float
    // === All variable declarations moved to top ===
    float *temp; // double -> float
    BiquadState *states;
    int i, j;
    float sample; // double -> float

    temp = (float *)malloc(len * sizeof(float)); // double -> float
    states = (BiquadState *)calloc(num_stages, sizeof(BiquadState));

    // === Forward filtering ===
    for (i = 0; i < len; i++) {
        sample = indata[i];
        for (j = 0; j < num_stages; j++) {
            sample = biquad_filter(sample, sos[j], &states[j]);
        }
        temp[i] = sample;
    }

    memset(states, 0, num_stages * sizeof(BiquadState));

    reverse_array(temp, len);

    for (i = 0; i < len; i++) {
        sample = temp[i];
        for (j = 0; j < num_stages; j++) {
            sample = biquad_filter(sample, sos[j], &states[j]);
        }
        temp[i] = sample;
    }

    reverse_array(temp, len);

    memcpy(outdata, temp, len * sizeof(float)); // double -> float

    free(temp);
    free(states);
}

// Quick sort (ascending order)
void quicksort(float *arr, int n) { // double -> float
    // All variables moved to top
    float pivot, temp; // double -> float
    int i, j;

    if (n <= 1) return;

    pivot = arr[n / 2];
    i = 0;
    j = n - 1;

    while (i <= j) {
        while (arr[i] < pivot) i++;
        while (arr[j] > pivot) j--;
        if (i <= j) {
            temp = arr[i]; arr[i] = arr[j]; arr[j] = temp;
            i++; j--;
        }
    }

    quicksort(arr, j + 1);
    if (i < n) quicksort(arr + i, n - i);
}

// Calculate median
float median(float *arr, int n) { // double -> float
    // All variables moved to top
    float *copy; // double -> float
    float med; // double -> float
    int i;

    copy = (float*)malloc(n * sizeof(float)); // double -> float
    if (!copy) return 0.0f;  // Memory allocation failed, 0.0 -> 0.0f

    for (i = 0; i < n; i++) {
        copy[i] = arr[i];
    }

    quicksort(copy, n);

    if (n % 2 == 1) {
        med = copy[n/2];
    } else {
        med = (copy[n/2-1] + copy[n/2]) * 0.5f; // 0.5 -> 0.5f
    }

    free(copy);
    return med;
}
/**
 * Symmetric index mapping function, used to simulate 'sym' mode boundary extension in MATLAB.
 * 
 * @param i Currently calculated index position.
 * @param N Signal length.
 * @return Mapped index position.
 */
int symmetric_index(int i, int N) {
    while (i < 0) { // If index is less than 0, flip to positive range
        i = -i - 1;
    }
    while (i >= N) { // If index exceeds signal length, mirror back into signal
        i = 2 * N - i - 1;
    }
    return i;
}
// Convolution + downsampling (symmetric boundaries)
void dwt_conv_down(const float *input, int len_in, const float *filt, int len_filt, float *output) { // double -> float
    // All variables moved to top
    int out_len;
    int i, j, center, n, idx;
    float sum; // double -> float

    out_len = (len_in + 1) / 2;  // Corrected length

    for (i = 0; i < out_len; i++) {
        center = 2 * i;
        sum = 0.0f; // 0.0 -> 0.0f
        for (j = 0; j < len_filt; j++) {
            n = center - (len_filt - 1)/2 + j;
            idx = symmetric_index(n, len_in);
            sum += input[idx] * filt[len_filt - 1 - j];
        }
        output[i] = sum;
    }
}

// Upsampling + convolution (reconstruction)
void idwt_up_conv(const float *input, int len_in, const float *filt, int len_filt, float *output, int len_out) { // double -> float
    // All variables moved to top
    int i, j, pos_in, pos_out;
    int radius;

    memset(output, 0, len_out * sizeof(float)); // double -> float
    radius = (len_filt - 1) / 2;  // 3 for sym4

    for (i = 0; i < len_in; i++) {
        pos_in = 2 * i;
        for (j = 0; j < len_filt; j++) {
            pos_out = pos_in + j - radius;
            if (pos_out >= 0 && pos_out < len_out) {
                output[pos_out] += input[i] * filt[len_filt - 1 - j];
            }
        }
    }
}

// Soft threshold function
void soft_threshold(float *vec, int n, float thr) { // double -> float
    int i;
    for (i = 0; i < n; i++) {
        if (vec[i] > thr)
            vec[i] -= thr;
        else if (vec[i] < -thr)
            vec[i] += thr;
        else
            vec[i] = 0.0f; // 0.0 -> 0.0f
    }
}

void denoise_emg_wavelab(float *data, int N, float fs, float *denoised) { // double -> float
    // === All variable declarations moved to top ===
    const int levels = 6;
    float *work, *tmp_low, *tmp_high; // double -> float
    float *details[6]; // double -> float
    int det_len[6];
    int current_len;
    int lev, out_len;
    float *abs_d6; // double -> float
    float med_abs, sigma, thr; // double -> float
    int i;

    // Allocate memory
    work = (float*)malloc(N * sizeof(float)); // double -> float
    tmp_low = (float*)malloc(N * sizeof(float));
    tmp_high = (float*)malloc(N * sizeof(float));
    if (!work || !tmp_low || !tmp_high) {
        // Error handling
        if (work) free(work);
        if (tmp_low) free(tmp_low);
        if (tmp_high) free(tmp_high);
        return;
    }

    memcpy(work, data, N * sizeof(float)); // double -> float
    current_len = N;

    // === 1. Multi-level wavelet decomposition ===
    for (lev = 0; lev < levels; lev++) {
        out_len = (current_len + 1) / 2;
        dwt_conv_down(work, current_len, sym4_h, 8, tmp_low);
        dwt_conv_down(work, current_len, sym4_g, 8, tmp_high);

        details[lev] = (float*)malloc(out_len * sizeof(float)); // double -> float
        det_len[lev] = out_len;
        memcpy(details[lev], tmp_high, out_len * sizeof(float)); // double -> float
        memcpy(work, tmp_low, out_len * sizeof(float)); // double -> float
        current_len = out_len;
    }

    // === 2. Noise estimation ===
    abs_d6 = (float*)malloc(det_len[5] * sizeof(float)); // double -> float

    for (i = 0; i < det_len[5]; i++) {
        abs_d6[i] = fabsf(details[5][i]); // fabs -> fabsf for float
    }
    med_abs = median(abs_d6, det_len[5]);
    sigma = med_abs / 0.6745f/4.2; // 0.6745 -> 0.6745f
    // thr = sigma * sqrt(2.0 * log((double)N));
    thr = sigma * sqrtf(2.0f * logf((float)N)); // sqrt -> sqrtf, log -> logf, (double)N -> (float)N, 2.0 -> 2.0f
    printf("Threshold: %.15f  sigma:%.15f\n", thr,sigma);
    free(abs_d6);

    // === 3. Soft thresholding ===
    for (lev = 0; lev < levels; lev++) {
        soft_threshold(details[lev], det_len[lev], thr);
    }

    // === 4. Reconstruction ===
    memcpy(tmp_low, work, current_len * sizeof(float)); // double -> float
    for (lev = levels - 1; lev >= 0; lev--) {
        out_len = (lev == 5) ? N : det_len[lev] * 2;
        idwt_up_conv(tmp_low, current_len, sym4_h, 8, work, out_len);
        idwt_up_conv(details[lev], det_len[lev], sym4_g, 8, tmp_high, out_len);
        for (i = 0; i < out_len; i++) {
            work[i] += tmp_high[i];
        }
        memcpy(tmp_low, work, out_len * sizeof(float)); // double -> float
        current_len = out_len;
    }

    memcpy(denoised, work, N * sizeof(float)); // double -> float

    free(work);
    free(tmp_low);
    free(tmp_high);
    for (i = 0; i < levels; i++) {
        if (details[i]) free(details[i]);
    }
}

void remove_dc_offset(const float* src, float* dest, int n) { // double -> float
    float sum, mean; // double -> float
    int i;

    if (n <= 0) return;

    sum = 0.0f; // 0.0 -> 0.0f
    for (i = 0; i < n; i++) {
        sum += src[i];
    }
    mean = sum / n;

    for (i = 0; i < n; i++) {
        dest[i] = src[i] - mean;
    }
}

// ================= Main processing function =================
void emg_denoised(float *emg_raw, int length, int fs, float *emg_denoised) { // double -> float
    // === All variable declarations moved to top ===
    // float *temp1; // double -> float
    // float *temp2; // double -> float
    NotchFilter filter;
    int i;

    // Validate input validity
    if (!emg_raw || length <= 0 || !emg_denoised) {
        fprintf(stderr, "Error: Invalid input parameters\n");
        return;
    }

    // Initialize notch filter (using predefined coefficients)
    for (i = 0; i < 3; i++) {
        filter.b[i] = NOTCH_B[i]; // 确保 NOTCH_B 和 NOTCH_A 在 emg_util.h 中也是 float 类型
        filter.a[i] = NOTCH_A[i];
    }

    // Reset filter states
    notchFilterReset(&filter);

    // Remove DC offset
    remove_dc_offset(emg_raw, emg_denoised, length);

    // Notch filtering (50Hz or 60Hz)
    //notchFilterProcessArray(&filter, emg_denoised, emg_denoised, length);

    // Bidirectional second-order section filtering (zero-phase bandpass filtering)
    filtfilt_sos(emg_denoised, emg_denoised, length, sos, 2); // 确保 sos 系数数组在 emg_util.h 中也是 float 类型
   for(i=0;i<length;i++)
   {
       emg_denoised[i]=emg_denoised[i]/1000;
   }
    // Wavelet denoising (using sym4)
    //denoise_emg_wavelab(emg_denoised, length, fs, emg_denoised);
}