#include <iostream>
#include <vector>
#include <cmath>
#include <fstream>
#include <random>

// 生成包含运动伪迹的测试信号
std::vector<float> generate_test_signal(int sample_rate, float duration) {
    int num_samples = static_cast<int>(sample_rate * duration);
    std::vector<float> signal(num_samples);
    
    std::random_device rd;
    std::mt19937 gen(rd());
    std::normal_distribution<float> noise(0.0f, 0.05f);
    
    for (int i = 0; i < num_samples; ++i) {
        float t = static_cast<float>(i) / sample_rate;
        
        // 基础信号
        float base = 0.5 * sin(2 * M_PI * 1.0 * t);
        
        // 运动伪迹
        float artifact = 0.0f;
        if (i % 500 == 0) { // 每500个样本添加伪迹
            artifact = 0.8f * exp(-std::pow((t - static_cast<int>(t)) * 10, 2));
        }
        
        signal[i] = base + artifact + noise(gen);
    }
    
    return signal;
}

// 简化的增强版运动伪迹检测
std::vector<float> remove_motion_artifacts_simple(const std::vector<float>& signal, double sample_rate) {
    if (signal.empty()) return signal;
    
    std::vector<float> result = signal;
    const size_t window_size = static_cast<size_t>(0.5 * sample_rate);
    
    if (signal.size() < window_size) return result;
    
    std::vector<bool> artifact_mask(signal.size(), false);
    
    // 多尺度检测
    for (size_t i = window_size; i < signal.size(); ++i) {
        std::vector<float> window;
        for (size_t j = i - window_size; j < i; ++j) {
            window.push_back(signal[j]);
        }
        
        float sum = 0.0f, sum_sq = 0.0f;
        for (float val : window) {
            sum += val;
            sum_sq += val * val;
        }
        
        float mean = sum / window_size;
        float variance = (sum_sq / window_size) - (mean * mean);
        
        if (variance > 1e-6f) {
            float std_dev = std::sqrt(variance);
            float current = signal[i];
            float z_score = std::abs(current - mean) / (std_dev + 1e-6f);
            
            if (z_score > 4.0f) {
                artifact_mask[i] = true;
            }
        }
    }
    
    // 修复伪迹
    size_t count = 0;
    for (size_t i = 0; i < signal.size(); ++i) {
        if (artifact_mask[i]) {
            count++;
            
            // 使用邻域中值
            std::vector<float> neighbors;
            for (int offset = -2; offset <= 2; ++offset) {
                int idx = static_cast<int>(i) + offset;
                if (idx >= 0 && idx < static_cast<int>(signal.size()) && !artifact_mask[idx]) {
                    neighbors.push_back(signal[idx]);
                }
            }
            
            if (!neighbors.empty()) {
                std::sort(neighbors.begin(), neighbors.end());
                result[i] = neighbors[neighbors.size() / 2];
            }
        }
    }
    
    std::cout << "检测到 " << count << " 个运动伪迹点" << std::endl;
    return result;
}

// 保存信号到CSV
void save_to_csv(const std::vector<float>& signal, const std::string& filename, double sample_rate) {
    std::ofstream file(filename);
    if (!file.is_open()) return;
    
    file << "Time(s),Amplitude\n";
    for (size_t i = 0; i < signal.size(); ++i) {
        float time = static_cast<float>(i) / sample_rate;
        file << time << "," << signal[i] << "\n";
    }
    file.close();
}

int main() {
    std::cout << "=== 增强版运动伪迹检测测试 ===" << std::endl;
    
    const int sample_rate = 100;
    const float duration = 5.0f;
    
    // 生成测试信号
    std::vector<float> original = generate_test_signal(sample_rate, duration);
    std::cout << "生成测试信号，长度: " << original.size() << " 样本" << std::endl;
    
    // 应用运动伪迹检测和去除
    std::vector<float> cleaned = remove_motion_artifacts_simple(original, sample_rate);
    
    // 保存结果
    save_to_csv(original, "original_with_artifacts.csv", sample_rate);
    save_to_csv(cleaned, "cleaned_signal.csv", sample_rate);
    
    std::cout << "测试完成，结果已保存到CSV文件" << std::endl;
    return 0;
}