# 示意代码，不可直接运行（质量特征提取与规则分类为本站简化示意，
# 非论文官方实现；论文实际使用 XGBoost 等机器学习分类器，未找到官方开源仓库）
import numpy as np
from scipy.signal import welch, butter, filtfilt
from dataclasses import dataclass

@dataclass
class rPPGSignal:
    """rPPG信号段"""
    raw: np.ndarray          # 原始rPPG信号
    fps: int                 # 帧率
    duration_sec: float      # 时长

class SignalQualityAssessor:
    """
    rPPG信号质量机器学习评估
    
    论文方法：从rPPG信号中提取质量指标，
    用机器学习分类器判定信号质量等级
    """
    
    def __init__(self, fs: int = 30):
        self.fs = fs
    
    def extract_quality_features(self, 
                                  rppg: np.ndarray) -> dict:
        """
        提取信号质量特征
        
        Args:
            rppg: [N] rPPG信号段
            
        Returns:
            features: dict of quality metrics
        """
        N = len(rppg)
        
        # 1. 频谱特征
        freqs, psd = welch(rppg, fs=self.fs, nperseg=min(256, N))
        
        # 心率频段（0.7-4.0 Hz = 42-240 BPM）
        hr_mask = (freqs >= 0.7) & (freqs <= 4.0)
        hr_power = np.sum(psd[hr_mask])
        total_power = np.sum(psd) + 1e-8
        spectral_ratio = hr_power / total_power
        
        # 峰值频率（预期在心率范围内）
        if np.any(psd[hr_mask] > 0):
            peak_freq = freqs[hr_mask][np.argmax(psd[hr_mask])]
        else:
            peak_freq = 0
        
        # 频谱集中度（心率峰附近能量集中度）
        peak_idx = np.argmax(psd[hr_mask])
        peak_width = self._peak_width(freqs[hr_mask], 
                                       psd[hr_mask], peak_idx)
        spectral_concentration = 1.0 / (1.0 + peak_width)
        
        # 2. 时域特征
        # 信号幅度
        amplitude = np.std(rppg)
        
        # 零交叉率
        zero_crossings = np.sum(np.diff(np.sign(rppg)) != 0)
        zcr = zero_crossings / N
        
        # 信噪比（SNR）
        # 信号：心率频段功率；噪声：其余频段
        noise_power = total_power - hr_power
        snr = 10 * np.log10(hr_power / (noise_power + 1e-8))
        
        # 3. 运动伪影指标
        # 高频成分占比（>4Hz通常为运动伪影）
        motion_mask = freqs > 4.0
        motion_power = np.sum(psd[motion_mask])
        motion_ratio = motion_power / total_power
        
        # 4. 趋势特征
        # 线性趋势（基线漂移）
        t = np.arange(N)
        trend_coef = np.polyfit(t, rppg, 1)[0]
        trend_ratio = abs(trend_coef * N) / (np.std(rppg) + 1e-8)
        
        # 5. 周期性
        # 自相关峰值
        autocorr = np.correlate(rppg - np.mean(rppg), 
                                 rppg - np.mean(rppg), 
                                 mode='full')
        autocorr = autocorr[N-1:] / autocorr[N-1]
        
        # 找第一个非零延迟峰值
        periodicity = 0
        for lag in range(int(0.5 * self.fs), int(3.0 * self.fs)):
            if lag < len(autocorr):
                if autocorr[lag] > periodicity:
                    periodicity = autocorr[lag]
        
        return {
            'spectral_ratio': spectral_ratio,
            'peak_freq': peak_freq,
            'spectral_concentration': spectral_concentration,
            'amplitude': amplitude,
            'zcr': zcr,
            'snr': snr,
            'motion_ratio': motion_ratio,
            'trend_ratio': trend_ratio,
            'periodicity': periodicity,
        }
    
    def classify_quality(self, features: dict) -> dict:
        """
        机器学习分类信号质量
        
        Returns:
            {'quality': str, 'confidence': float, 'usable': bool}
        """
        # 简化版规则分类器（论文实际用 XGBoost/SVM/CNN/LSTM 等）
        score = 0
        
        # 频谱比率 > 0.3 → 高质量
        if features['spectral_ratio'] > 0.3:
            score += 2
        elif features['spectral_ratio'] > 0.15:
            score += 1
        
        # SNR > 3dB → 高质量
        if features['snr'] > 3.0:
            score += 2
        elif features['snr'] > 0.0:
            score += 1
        
        # 运动伪影比 < 0.2 → 高质量
        if features['motion_ratio'] < 0.2:
            score += 2
        elif features['motion_ratio'] < 0.4:
            score += 1
        
        # 周期性 > 0.5 → 高质量
        if features['periodicity'] > 0.5:
            score += 2
        elif features['periodicity'] > 0.3:
            score += 1
        
        # 趋势比 < 0.5 → 稳定
        if features['trend_ratio'] < 0.5:
            score += 1
        
        # 分类
        if score >= 7:
            quality = 'Excellent'
            confidence = 0.9
            usable = True
        elif score >= 5:
            quality = 'Good'
            confidence = 0.75
            usable = True
        elif score >= 3:
            quality = 'Fair'
            confidence = 0.5
            usable = False  # 标记不可靠
        else:
            quality = 'Poor'
            confidence = 0.2
            usable = False
        
        return {
            'quality': quality,
            'confidence': confidence,
            'usable': usable,
            'score': score,
        }
    
    def _peak_width(self, freqs, psd, peak_idx, 
                     threshold=0.5):
        """计算峰宽"""
        peak_val = psd[peak_idx]
        half_max = peak_val * threshold
        
        # 左边界
        left = peak_idx
        while left > 0 and psd[left] > half_max:
            left -= 1
        
        # 右边界
        right = peak_idx
        while right < len(psd) - 1 and psd[right] > half_max:
            right += 1
        
        return freqs[right] - freqs[left]


# 测试
if __name__ == "__main__":
    assessor = SignalQualityAssessor(fs=30)
    
    # 模拟高质量rPPG（清晰脉搏）
    t = np.arange(900) / 30  # 30s at 30fps
    clean_rppg = 0.5 * np.sin(2 * np.pi * 1.2 * t)  # 72 BPM
    clean_rppg += 0.05 * np.random.randn(900)
    
    features_clean = assessor.extract_quality_features(clean_rppg)
    result_clean = assessor.classify_quality(features_clean)
    print("高质量信号:")
    print(f"  质量: {result_clean['quality']}")
    print(f"  可用: {result_clean['usable']}")
    print(f"  SNR: {features_clean['snr']:.1f}dB")
    print(f"  周期性: {features_clean['periodicity']:.2f}")
    
    # 模拟低质量rPPG（大量运动伪影）
    noisy_rppg = 0.1 * np.sin(2 * np.pi * 1.2 * t)
    noisy_rppg += 0.5 * np.random.randn(900)  # 强噪声
    noisy_rppg += 0.3 * np.sin(2 * np.pi * 8 * t)  # 运动伪影
    
    features_noisy = assessor.extract_quality_features(noisy_rppg)
    result_noisy = assessor.classify_quality(features_noisy)
    print("\n低质量信号:")
    print(f"  质量: {result_noisy['quality']}")
    print(f"  可用: {result_noisy['usable']}")
    print(f"  SNR: {features_noisy['snr']:.1f}dB")
    print(f"  运动比: {features_noisy['motion_ratio']:.2f}")