import numpy as np
from dataclasses import dataclass
from typing import Tuple, Optional
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
import warnings
warnings.filterwarnings('ignore')

"""
非接触式心理压力监测系统
参考 Frontiers in Psychiatry 2026 Meta-Analysis

模块:
1. rPPG 心率提取（从面部视频）
2. HRV 特征工程
3. 压力分类（正常/压力/过载）
"""

@dataclass
class rPPGResult:
    """rPPG 提取结果"""
    heart_rate: float          # 心率(bpm)
    hr_mean_nn: float          # RR间期均值(ms)
    hr_sdnn: float             # RR间期标准差(ms)
    hr_rmssd: float            # 连续差值均方根(ms)
    hr_pnn50: float            # 相邻RR差>50ms比例(%)
    lf_hf_ratio: float         # LF/HF比（交感/副交感平衡）
    stress_score: float        # 综合压力分数 0-1
    confidence: float          # 置信度


class rPPGExtractor:
    """
    远程光电容积描记（rPPG）心率提取器
    
    使用 POS (Plane-Orthogonal-to-Skin) 算法
    从 RGB 视频中提取心率信号
    
    参考: Wang et al., IEEE TBIOM 2021
    """
    
    def __init__(self, fps: int = 30, window_sec: int = 10):
        self.fps = fps
        self.window_sec = window_sec
        self.window_frames = fps * window_sec
        
    def extract_from_face_roi(self, face_frames: np.ndarray) -> rPPGResult:
        """
        从面部 ROI 帧序列提取心率
        
        Args:
            face_frames: (N, H, W, 3) RGB 面部帧
        
        Returns:
            rPPGResult
        """
        N = len(face_frames)
        
        if N < self.window_frames:
            return rPPGResult(0, 0, 0, 0, 0, 0, 0, 0)
        
        # 1. 空间平均 -> 时序信号
        # 每帧面部 ROI 的 RGB 均值
        rgb_means = np.array([
            frame.reshape(-1, 3).mean(axis=0) 
            for frame in face_frames
        ])  # (N, 3)
        
        # 2. POS 算法：投影到正交于皮肤的平面
        # S = 3R - 2G, H = 1.5R - G
        S = 3 * rgb_means[:, 0] - 2 * rgb_means[:, 1]
        H = 1.5 * rgb_means[:, 0] - rgb_means[:, 1]
        
        # 归一化
        S = (S - np.mean(S)) / (np.std(S) + 1e-8)
        H = (H - np.mean(H)) / (np.std(H) + 1e-8)
        
        # 组合信号
        ppg_signal = S + H * 0.3  # 加权组合
        
        # 3. 带通滤波 (0.7-3.5 Hz = 42-210 bpm)
        from scipy.signal import butter, filtfilt
        nyq = self.fps / 2
        low, high = 0.7 / nyq, 3.5 / nyq
        b, a = butter(3, [low, high], btype='band')
        ppg_filtered = filtfilt(b, a, ppg_signal)
        
        # 4. 峰值检测 -> RR 间期
        from scipy.signal import find_peaks
        peaks, _ = find_peaks(ppg_filtered, distance=self.fps*0.4)
        
        if len(peaks) < 3:
            return rPPGResult(0, 0, 0, 0, 0, 0, 0, 0)
        
        rr_intervals = np.diff(peaks) / self.fps * 1000  # ms
        
        # 5. HRV 特征
        hr = 60000 / np.mean(rr_intervals)  # bpm
        mean_nn = np.mean(rr_intervals)
        sdnn = np.std(rr_intervals)
        diff_rr = np.diff(rr_intervals)
        rmssd = np.sqrt(np.mean(diff_rr**2)) if len(diff_rr) > 0 else 0
        pnn50 = np.mean(np.abs(diff_rr) > 50) * 100 if len(diff_rr) > 0 else 0
        
        # 频域：简化 LF/HF
        if len(rr_intervals) > 10:
            # 简化：用 SDNN/RMSSD 比作为 LF/HF 近似
            lf_hf = sdnn / (rmssd + 1e-8)
        else:
            lf_hf = 1.0
        
        # 6. 压力分数
        # HRV 降低 + 心率升高 = 压力增大
        stress_score = self._compute_stress_score(hr, sdnn, rmssd, lf_hf)
        
        # 7. 置信度（基于信号质量）
        confidence = min(1.0, len(peaks) / (self.window_frames / self.fps * 1.5))
        
        return rPPGResult(
            heart_rate=round(hr, 1),
            hr_mean_nn=round(mean_nn, 1),
            hr_sdnn=round(sdnn, 1),
            hr_rmssd=round(rmssd, 1),
            hr_pnn50=round(pnn50, 1),
            lf_hf_ratio=round(lf_hf, 2),
            stress_score=round(stress_score, 3),
            confidence=round(confidence, 3)
        )
    
    def _compute_stress_score(self, hr, sdnn, rmssd, lf_hf):
        """
        综合压力评分
        
        基于 meta-analysis 发现：
        - 心率升高 → 压力
        - HRV(SDNN)降低 → 压力
        - LF/HF升高 → 交感占优 → 压力
        """
        # 各分量归一化
        hr_component = max(0, min(1, (hr - 60) / 40))  # 60-100 bpm
        sdnn_component = max(0, min(1, 1 - sdnn / 50))  # SDNN<50ms = 压力
        lf_hf_component = max(0, min(1, (lf_hf - 1) / 3))  # LF/HF 1-4
        
        # 加权
        score = (hr_component * 0.35 + 
                sdnn_component * 0.35 + 
                lf_hf_component * 0.30)
        
        return score


class StressClassifier:
    """
    心理压力三级分类器
    
    0 = 正常 (stress_score < 0.3)
    1 = 轻度压力 (0.3-0.6)
    2 = 过载 (> 0.6)
    """
    
    def __init__(self):
        self.classifier = RandomForestClassifier(
            n_estimators=100, max_depth=8, random_state=42,
            class_weight='balanced'
        )
        self.feature_names = ['HR', 'MeanNN', 'SDNN', 'RMSSD', 'pNN50', 'LF_HF']
        self.is_trained = False
    
    def fit(self, X: np.ndarray, y: np.ndarray):
        """训练"""
        self.classifier.fit(X, y)
        self.is_trained = True
        
    def predict(self, rppg: rPPGResult) -> int:
        """预测压力等级"""
        if not self.is_trained:
            # 未训练时用规则
            if rppg.stress_score < 0.3:
                return 0
            elif rppg.stress_score < 0.6:
                return 1
            else:
                return 2
        
        features = np.array([[
            rppg.heart_rate, rppg.hr_mean_nn, rppg.hr_sdnn,
            rppg.hr_rmssd, rppg.hr_pnn50, rppg.lf_hf_ratio
        ]])
        return self.classifier.predict(features)[0]
    
    def evaluate(self, X: np.ndarray, y: np.ndarray) -> float:
        """10折交叉验证"""
        scores = cross_val_score(self.classifier, X, y, cv=10, scoring='accuracy')
        return np.mean(scores)


# ==================== 测试 ====================
if __name__ == "__main__":
    np.random.seed(42)
    
    print("=" * 70)
    print("非接触式心理压力监测系统测试")
    print("参考: Frontiers in Psychiatry 2026 Meta-Analysis")
    print("=" * 70)
    
    # 模拟 rPPG 信号
    extractor = rPPGExtractor(fps=30, window_sec=10)
    
    # 生成模拟面部帧
    num_frames = 300  # 10s @ 30fps
    
    # 正常状态
    normal_frames = np.random.randint(100, 120, (num_frames, 100, 100, 3), dtype=np.uint8)
    # 添加微弱脉搏波（70bpm = 1.17Hz）
    t = np.arange(num_frames) / 30
    pulse_normal = np.sin(2 * np.pi * 1.17 * t) * 2
    for i in range(num_frames):
        normal_frames[i] = np.clip(normal_frames[i].astype(float) + pulse_normal[i], 0, 255).astype(np.uint8)
    
    # 压力状态（90bpm = 1.5Hz, HRV降低）
    stress_frames = np.random.randint(100, 120, (num_frames, 100, 100, 3), dtype=np.uint8)
    pulse_stress = np.sin(2 * np.pi * 1.5 * t) * 2
    for i in range(num_frames):
        stress_frames[i] = np.clip(stress_frames[i].astype(float) + pulse_stress[i], 0, 255).astype(np.uint8)
    
    # 提取 rPPG
    print("\n--- 正常状态 ---")
    result_normal = extractor.extract_from_face_roi(normal_frames)
    print(f"  心率: {result_normal.heart_rate} bpm")
    print(f"  SDNN: {result_normal.hr_sdnn} ms")
    print(f"  RMSSD: {result_normal.hr_rmssd} ms")
    print(f"  LF/HF: {result_normal.lf_hf_ratio}")
    print(f"  压力分数: {result_normal.stress_score}")
    print(f"  置信度: {result_normal.confidence}")
    
    print("\n--- 压力状态 ---")
    result_stress = extractor.extract_from_face_roi(stress_frames)
    print(f"  心率: {result_stress.heart_rate} bpm")
    print(f"  SDNN: {result_stress.hr_sdnn} ms")
    print(f"  RMSSD: {result_stress.hr_rmssd} ms")
    print(f"  LF/HF: {result_stress.lf_hf_ratio}")
    print(f"  压力分数: {result_stress.stress_score}")
    print(f"  置信度: {result_stress.confidence}")
    
    # 分类器测试
    print("\n--- 分类器训练 ---")
    classifier = StressClassifier()
    
    # 生成训练数据
    num_samples = 200
    X = []
    y = []
    
    for i in range(num_samples):
        label = i % 3  # 0=正常, 1=轻度, 2=过载
        if label == 0:
            X.append([70 + np.random.normal(0, 5), 850, 50, 35, 8, 1.2])
        elif label == 1:
            X.append([85 + np.random.normal(0, 5), 700, 30, 20, 4, 2.5])
        else:
            X.append([100 + np.random.normal(0, 5), 600, 15, 10, 1, 4.0])
        y.append(label)
    
    X = np.array(X)
    y = np.array(y)
    
    classifier.fit(X, y)
    acc = classifier.evaluate(X, y)
    print(f"  10折交叉验证准确率: {acc:.4f}")
    
    # 对比 Meta-Analysis 结果
    print(f"\n{'='*70}")
    print("对比 Meta-Analysis 结果:")
    print(f"  Meta-Analysis 准确率: 81.3% (95% CI 68.4-91.5%)")
    print(f"  模拟准确率: {acc*100:.1f}%")
    print(f"  F1: 0.791 | 敏感度: 0.809 | 特异度: 0.697")
    
    # 三模态对比
    print(f"\n{'='*70}")
    print("非接触模态对比 (Meta-Analysis):")
    print(f"  {'模态':<15} {'准确率':>10} {'IMS可行性':>15}")
    print(f"  {'rPPG':<15} {'~82%':>10} {'✅ 高':>15}")
    print(f"  {'热成像':<15} {'~80%':>10} {'⚠️ 中':>15}")
    print(f"  {'雷达':<15} {'~79%':>10} {'✅ 高(复用CPD)':>15}")
    print(f"  {'融合':<15} {'~85%':>10} {'✅ 最佳':>15}")
    
    print(f"\n{'='*70}")