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

// 模拟PPG信号生成函数，包含运动伪迹
std::vector<float> generate_test_ppg_with_artifacts(int sample_rate, float duration_seconds) {
    int num_samples = static_cast<int>(sample_rate * duration_seconds);
    std::vector<float> signal(num_samples);
    
    std::random_device rd;
    std::mt19937 gen(rd());
    std::normal_distribution<float> noise_dist(0.0f, 0.05f);
    std::uniform_real_distribution<float> artifact_dist(0.0f, 1.0f);
    
    for (int i = 0; i < num_samples; ++i) {
        float t = static_cast<float>(i) / sample_rate;
        
        // 基础PPG信号（心率约60bpm）
        float ppg_component = 0.5 * sin(2 * M_PI * 1.0 * t);
        
        // 呼吸调制
        float respiration = 0.1 * sin(2 * M_PI * 0.2 * t);
        
        // 基线漂移
        float baseline_drift = 0.05 * sin(2 * M_PI * 0.01 * t);
        
        // 高频噪声
        float high_freq_noise = 0.02 * sin(2 * M_PI * 10.0 * t) + 
                               0.02 * sin(2 * M_PI * 15.0 * t);
        
        // 运动伪迹（随机出现）
        float motion_artifact = 0.0f;
        if (artifact_dist(gen) < 0.01f) { // 1%概率出现运动伪迹
            float artifact_type = artifact_dist(gen);
            if (artifact_type < 0.3f) {
                // 尖峰伪迹
                motion_artifact = 0.8f * exp(-std::pow((t - static_cast<int>(t)) * 10, 2));
            } else if (artifact_type < 0.6f) {
                // 基线跳跃
                motion_artifact = 0.3f * (artifact_dist(gen) - 0.5f);
            } else {
                // 高频振荡
                motion_artifact = 0.2f * sin(2 * M_PI * 25.0 * t) * exp(-std::pow((t - static_cast<int>(t)) * 5, 2));
            }
        }
        
        // 组合所有成分
        signal[i] = ppg_component + respiration + baseline_drift + high_freq_noise + motion_artifact + noise_dist(gen);
    }
    
    return signal;
}

// 简化的增强版运动伪迹检测算法（用于测试）
std::vector<float> enhanced_motion_artifact_removal(const std::vector<float>& signal, double sample_rate) {
    if (signal.empty()) return signal;
    
    std::vector<float> result = signal;
    const size_t min_window_size = static_cast<size_t>(0.1 * sample_rate);
    
    if (signal.size() < min_window_size) return result;
    
    // 多尺度检测窗口
    std::vector<size_t> window_sizes = {
        static_cast<size_t>(0.1 * sample_rate),
        static_cast<size_t>(0.5 * sample_rate),
        static_cast<size_t>(1.0 * sample_rate)
    };
    
    std::vector<bool> artifact_mask(signal.size(), false);
    
    // 统计特征检测
    for (size_t ws : window_sizes) {
        if (signal.size() < ws) continue;
        
        for (size_t i = ws; i < signal.size(); ++i) {
            std::vector<float> window_data;
            for (size_t j = i - ws; j < i; ++j) {
                window_data.push_back(signal[j]);
            }
            
            // 计算统计特征
            float sum = 0.0f, sum_sq = 0.0f;
            for (float val : window_data) {
                sum += val;
                sum_sq += val * val;
            }
            
            float mean = sum / ws;
            float variance = (sum_sq / ws) - (mean * mean);
            
            if (variance > 1e-6f) {
                float std_dev = std::sqrt(variance);
                float current_val = signal[i];
                float z_score = std::abs(current_val - mean) / (std_dev + 1e-6f);
                
                // 检测异常值
                if (z_score > 4.0f) {
                    artifact_mask[i] = true;
                }
                
                // 梯度异常检测
                if (i > 0) {
                    float gradient = std::abs(current_val - signal[i-1]);
                    float avg_gradient = 0.0f;
                    for (size_t j = 1; j < std::min(ws, i); ++j) {
                        avg_gradient += std::abs(signal[i-j] - signal[i-j-1]);
                    }
                    avg_gradient /= std::min(ws, i);
                    
                    if (gradient > 3.0f * avg_gradient) {
                        artifact_mask[i] = true;
                    }
                }
            }
        }
    }
    
    // 形态学检测
    for (size_t i = 2; i < signal.size() - 2; ++i) {
        float current = signal[i];
        float prev = signal[i-1];
        float next = signal[i+1];
        
        // 尖峰检测
        if ((current > prev && current > next && 
             current - prev > 2.0f * std::abs(next - prev)) ||
            (current < prev && current < next && 
             prev - current > 2.0f * std::abs(next - prev))) {
            artifact_mask[i] = true;
        }
    }
    
    // 智能修复
    size_t artifact_count = 0;
    for (size_t i = 0; i < signal.size(); ++i) {
        if (artifact_mask[i]) {
            artifact_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()) {
                if (neighbors.size() >= 3) {
                    std::sort(neighbors.begin(), neighbors.end());
                    result[i] = neighbors[neighbors.size() / 2];
                } else {
                    float sum = 0.0f;
                    for (float val : neighbors) sum += val;
                    result[i] = sum / neighbors.size();
                }
            }
        }
    }
    
    std::cout << "检测到 " << artifact_count << " 个运动伪迹点" << std::endl;
    return result;
}

// 计算信号质量指标
struct SignalQualityMetrics {
    float snr_db;
    float variance;
    float peak_to_peak;
    float artifact_ratio;
};

SignalQualityMetrics calculate_signal_quality(const std::vector<float>& signal, double sample_rate) {
    SignalQualityMetrics metrics;
    
    if (signal.empty()) {
        metrics.snr_db = 0.0f;
        metrics.variance = 0.0f;
        metrics.peak_to_peak = 0.0f;
        metrics.artifact_ratio = 0.0f;
        return metrics;
    }
    
    // 计算基本统计量
    float sum = 0.0f, sum_sq = 0.0f;
    float min_val = signal[0], max_val = signal[0];
    
    for (float val : signal) {
        sum += val;
        sum_sq += val * val;
        min_val = std::min(min_val, val);
        max_val = std::max(max_val, val);
    }
    
    float mean = sum / signal.size();
    metrics.variance = (sum_sq / signal.size()) - (mean * mean);
    metrics.peak_to_peak = max_val - min_val;
    
    // 计算SNR
    float signal_power = 0.0f;
    float noise_power = 0.0f;
    
    for (float val : signal) {
        signal_power += std::pow(val - mean, 2);
    }
    signal_power /= signal.size();
    
    for (size_t i = 1; i < signal.size(); ++i) {
        float diff = signal[i] - signal[i-1];
        noise_power += diff * diff;
    }
    noise_power /= (signal.size() - 1);
    
    if (noise_power > 1e-6f) {
        metrics.snr_db = 10.0f * std::log10(signal_power / noise_power);
    } else {
        metrics.snr_db = 0.0f;
    }
    
    // 估算伪迹比例（通过异常梯度检测）
    size_t artifact_count = 0;
    for (size_t i = 1; i < signal.size(); ++i) {
        float gradient = std::abs(signal[i] - signal[i-1]);
        if (gradient > 3.0f * std::sqrt(metrics.variance)) {
            artifact_count++;
        }
    }
    
    metrics.artifact_ratio = static_cast<float>(artifact_count) / signal.size();
    
    return metrics;
}

// 保存信号到CSV文件
void save_signal_to_csv(const std::vector<float>& signal, const std::string& filename, double sample_rate) {
    std::ofstream file(filename);
    if (!file.is_open()) {
        std::cerr << "无法创建文件: " << filename << std::endl;
        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();
    std::cout << "信号已保存到: " << filename << std::endl;
}

int main() {
    std::cout << "=== 增强版运动伪迹检测和去除算法测试 ===" << std::endl;
    
    // 测试参数
    const int sample_rate = 100; // 100Hz采样率
    const float duration = 10.0f; // 10秒数据
    
    std::cout << "采样率: " << sample_rate << "Hz" << std::endl;
    std::cout << "信号长度: " << duration << "秒" << std::endl;
    
    // 1. 生成包含运动伪迹的测试信号
    std::cout << "\n1. 生成测试信号..." << std::endl;
    std::vector<float> original_signal = generate_test_ppg_with_artifacts(sample_rate, duration);
    std::cout << "测试信号生成完成，长度: " << original_signal.size() << " 样本" << std::endl;
    
    // 2. 计算原始信号质量
    std::cout << "\n2. 分析原始信号质量..." << std::endl;
    SignalQualityMetrics original_metrics = calculate_signal_quality(original_signal, sample_rate);
    std::cout << "原始信号质量指标:" << std::endl;
    std::cout << "  SNR: " << original_metrics.snr_db << " dB" << std::endl;
    std::cout << "  方差: " << original_metrics.variance << std::endl;
    std::cout << "  峰峰值: " << original_metrics.peak_to_peak << std::endl;
    std::cout << "  伪迹比例: " << (original_metrics.artifact_ratio * 100) << "%" << std::endl;
    
    // 3. 应用运动伪迹检测和去除
    std::cout << "\n3. 应用增强版运动伪迹检测和去除..." << std::endl;
    std::vector<float> cleaned_signal = enhanced_motion_artifact_removal(original_signal, sample_rate);
    
    // 4. 计算处理后信号质量
    std::cout << "\n4. 分析处理后信号质量..." << std::endl;
    SignalQualityMetrics cleaned_metrics = calculate_signal_quality(cleaned_signal, sample_rate);
    std::cout << "处理后信号质量指标:" << std::endl;
    std::cout << "  SNR: " << cleaned_metrics.snr_db << " dB" << std::endl;
    std::cout << "  方差: " << cleaned_metrics.variance << std::endl;
    std::cout << "  峰峰值: " << cleaned_metrics.peak_to_peak << std::endl;
    std::cout << "  伪迹比例: " << (cleaned_metrics.artifact_ratio * 100) << "%" << std::endl;
    
    // 5. 计算改善效果
    std::cout << "\n5. 改善效果分析..." << std::endl;
    float snr_improvement = cleaned_metrics.snr_db - original_metrics.snr_db;
    float variance_reduction = (original_metrics.variance - cleaned_metrics.variance) / original_metrics.variance * 100;
    float artifact_reduction = (original_metrics.artifact_ratio - cleaned_metrics.artifact_ratio) / original_metrics.artifact_ratio * 100;
    
    std::cout << "改善效果:" << std::endl;
    std::cout << "  SNR改善: " << snr_improvement << " dB" << std::endl;
    std::cout << "  方差减少: " << variance_reduction << "%" << std::endl;
    std::cout << "  伪迹减少: " << artifact_reduction << "%" << std::endl;
    
    // 6. 保存结果
    std::cout << "\n6. 保存结果..." << std::endl;
    save_signal_to_csv(original_signal, "original_signal_with_artifacts.csv", sample_rate);
    save_signal_to_csv(cleaned_signal, "cleaned_signal.csv", sample_rate);
    
    // 7. 生成对比报告
    std::ofstream report("motion_artifact_removal_report.txt");
    if (report.is_open()) {
        report << "=== 运动伪迹检测和去除算法测试报告 ===" << std::endl;
        report << "测试时间: " << std::time(nullptr) << std::endl;
        report << "采样率: " << sample_rate << "Hz" << std::endl;
        report << "信号长度: " << duration << "秒" << std::endl;
        report << "\n原始信号质量:" << std::endl;
        report << "  SNR: " << original_metrics.snr_db << " dB" << std::endl;
        report << "  方差: " << original_metrics.variance << std::endl;
        report << "  峰峰值: " << original_metrics.peak_to_peak << std::endl;
        report << "  伪迹比例: " << (original_metrics.artifact_ratio * 100) << "%" << std::endl;
        report << "\n处理后信号质量:" << std::endl;
        report << "  SNR: " << cleaned_metrics.snr_db << " dB" << std::endl;
        report << "  方差: " << cleaned_metrics.variance << std::endl;
        report << "  峰峰值: " << cleaned_metrics.peak_to_peak << std::endl;
        report << "  伪迹比例: " << (cleaned_metrics.artifact_ratio * 100) << "%" << std::endl;
        report << "\n改善效果:" << std::endl;
        report << "  SNR改善: " << snr_improvement << " dB" << std::endl;
        report << "  方差减少: " << variance_reduction << "%" << std::endl;
        report << "  伪迹减少: " << artifact_reduction << "%" << std::endl;
        report.close();
        std::cout << "测试报告已保存到: motion_artifact_removal_report.txt" << std::endl;
    }
    
    std::cout << "\n=== 测试完成 ===" << std::endl;
    std::cout << "请检查生成的CSV文件和测试报告" << std::endl;
    
    return 0;
}