"""
基于rPPG的非接触压力检测框架

复现meta分析中的主流rPPG方法
"""
import numpy as np
import torch
import torch.nn as nn

class rPPGStressDetector(nn.Module):
    """
    rPPG压力检测模型
    
    流程：
    1. 面部检测+ROI提取（额头/双颊）
    2. 像素平均→时间序列（原始rPPG信号）
    3. 滤波→心率/HRV特征
    4. 时序模型→压力分类
    """
    
    def __init__(self, fps=30, window_sec=60):
        super().__init__()
        self.fps = fps
        self.window = int(window_sec * fps)
        
        # 信号处理参数
        self.hr_band = (0.7, 3.5)  # 42-210 bpm
        
        # 时序模型
        self.lstm = nn.LSTM(input_size=32, hidden_size=64, 
                            num_layers=2, batch_first=True)
        self.classifier = nn.Sequential(
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(32, 2),  # binary: stressed/not
        )
    
    def extract_rppg_signal(self, video_frames, facial_rois):
        """
        从视频帧提取rPPG信号
        
        方法：POS (Plane Orthogonal to Skin)
        
        Args:
            video_frames: shape=(T, H, W, 3), RGB视频帧
            facial_rois: list of (x,y,w,h), 面部ROI区域
            
        Returns:
            rppg_signal: shape=(T,), 原始rPPG信号
        """
        T = len(video_frames)
        rppg = np.zeros(T)
        
        for t in range(T):
            frame = video_frames[t]
            roi = facial_rois[t]
            x, y, w, h = roi
            
            # 提取ROI区域
            patch = frame[y:y+h, x:x+w, :]
            
            # 空间平均
            r_mean = np.mean(patch[:,:,0])
            g_mean = np.mean(patch[:,:,1])
            b_mean = np.mean(patch[:,:,2])
            
            # POS投影
            # C = [3x1] 颜色均值
            # 投影到与皮肤色调正交的平面
            s = np.array([3*r_mean - 2*g_mean,
                          -1.5*r_mean + g_mean + 1.5*b_mean])
            rppg[t] = s[0]  # 取第一分量
        
        # 去趋势+带通滤波
        from scipy.signal import detrend, butter, filtfilt
        rppg = detrend(rppg)
        nyq = self.fps / 2
        b, a = butter(3, [self.hr_band[0]/nyq, self.hr_band[1]/nyq], 
                      btype='band')
        rppg = filtfilt(b, a, rppg)
        
        return rppg
    
    def extract_hrv_features(self, rppg_signal):
        """提取HRV特征"""
        # 峰值检测（心跳）
        from scipy.signal import find_peaks
        peaks, _ = find_peaks(rppg_signal, distance=self.fps*0.4)
        
        if len(peaks) < 10:
            return np.zeros(32)  # 不足
        
        # RR间隔
        rr_intervals = np.diff(peaks) / self.fps * 1000  # ms
        
        features = np.zeros(32)
        # 时域特征
        features[0] = np.mean(rr_intervals)  # mean RR
        features[1] = np.std(rr_intervals)   # SDNN
        features[2] = np.sqrt(np.mean(np.diff(rr_intervals)**2))  # RMSSD
        features[3] = len(peaks) / (len(rppg_signal)/self.fps) * 60  # HR
        
        # 频域特征
        from scipy.signal import periodogram
        f, pxx = periodogram(rr_intervals, fs=1.0/np.mean(rr_intervals)*1000)
        vlf = np.sum(pxx[(f<0.04)])  # VLF
        lf = np.sum(pxx[(f>=0.04)&(f<0.15)])  # LF
        hf = np.sum(pxx[(f>=0.15)&(f<0.4)])  # HF
        features[4] = lf
        features[5] = hf
        features[6] = lf/hf if hf > 0 else 0  # LF/HF ratio
        
        return features
    
    def forward(self, rppg_signal):
        """压力分类"""
        # 提取特征序列（滑动窗口）
        features = self.extract_hrv_features(rppg_signal)
        
        # LSTM分类
        x = torch.tensor(features, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
        out, _ = self.lstm(x)
        out = self.classifier(out[:, -1, :])
        return torch.softmax(out, dim=1)

"""
IMS非接触压力检测融合方案

rPPG + 热成像 + 驾驶行为 三模态融合
预期：准确率>88%, 特异度>85%
"""
class MultiModalStressDetector:
    def __init__(self):
        self.rppg_model = rPPGStressDetector()
        self.thermal_model = ThermalStressModel()
        self.driving_model = DrivingBehaviorStressModel()
        
    def predict(self, video, thermal, can_signals):
        """三模态融合"""
        rppg_prob = self.rppg_model(video)
        thermal_prob = self.thermal_model(thermal)
        driving_prob = self.driving_model(can_signals)
        
        # 加权融合
        weights = [0.5, 0.3, 0.2]  # rPPG最可靠
        fused = (weights[0]*rppg_prob + weights[1]*thermal_prob 
                 + weights[2]*driving_prob)
        return fused