import torch
import torch.nn as nn
import numpy as np
from typing import Tuple

class PhysiologicalPreprocessing(nn.Module):
    """
    SpiRadar 预处理模块
    
    论文 Section III-A
    解决强制呼气动作中的非自主运动伪影
    
    关键问题：
    - 强制呼气（FEV1 maneuver）产生大幅躯干运动
    - 非自主运动叠加在呼吸信号上
    - 需要区分胸腔呼吸运动与整体躯干位移
    """
    def __init__(self, fs=200, cutoff_freq=5.0):
        super().__init__()
        self.fs = fs  # 雷达采样率
        
        # 带通滤波：呼吸频段 0.1-5 Hz
        from scipy.signal import butter, filtfilt
        self.b_low, self.a_low = butter(4, cutoff_freq / (fs / 2), btype='low')
        
    def forward(self, raw_phase: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            raw_phase: 雷达回波相位 (B, T) - 胸壁位移原始信号
        Returns:
            respiratory_signal: 呼吸信号 (B, T)
            motion_artifact: 运动伪影 (B, T)
        """
        # 转为 numpy 处理（实际部署用 torch 实现）
        phase_np = raw_phase.cpu().numpy()
        
        results = []
        artifacts = []
        for i in range(phase_np.shape[0]):
            signal = phase_np[i]
            
            # 1. 对数插值校正
            # 强制呼气时可能有相位跳变
            signal = self._log_interpolation_correct(signal)
            
            # 2. 低通滤波提取呼吸分量
            resp = filtfilt(self.b_low, self.a_low, signal)
            
            # 3. 运动伪影 = 原始 - 呼吸
            artifact = signal - resp
            
            results.append(resp)
            artifacts.append(artifact)
        
        return (torch.from_numpy(np.array(results)).to(raw_phase.device),
                torch.from_numpy(np.array(artifacts)).to(raw_phase.device))
    
    def _log_interpolation_correct(self, signal):
        """对数插值校正相位跳变"""
        # 检测大幅跳变
        diff = np.diff(signal)
        jumps = np.where(np.abs(diff) > np.pi)[0]
        
        for j in jumps:
            # 线性插值修复
            if j > 0 and j < len(signal) - 1:
                signal[j] = (signal[j-1] + signal[j+1]) / 2
        
        return signal


class SparsePolynomialTransform(nn.Module):
    """
    稀疏多项式变换模型
    
    论文 Section III-B
    将雷达测量的胸壁位移转换为肺活量体积曲线
    
    核心创新：
    1. 特征依赖的多项式（不同呼吸阶段使用不同系数）
    2. 稀疏优化（L1 正则化）实现跨被试泛化
    3. 无需个体校准
    """
    def __init__(self, poly_order=5, num_features=8):
        super().__init__()
        self.poly_order = poly_order
        self.num_features = num_features
        
        # 特征提取器
        self.feature_extractor = nn.Sequential(
            nn.Conv1d(1, 32, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(2),
            nn.Conv1d(32, 64, kernel_size=5, padding=2),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(1)
        )
        
        # 特征 -> 多项式系数
        # 每个特征组合产生一组多项式系数
        self.coeff_predictor = nn.Sequential(
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, (poly_order + 1) * num_features)
        )
        
        # 稀疏门控（L1 正则化效果）
        self.sparse_gate = nn.Sequential(
            nn.Linear(64, num_features),
            nn.Sigmoid()
        )
        
    def forward(self, displacement: torch.Tensor) -> torch.Tensor:
        """
        Args:
            displacement: 胸壁位移序列 (B, T)
        Returns:
            volume: 肺活量体积曲线 (B, T)
        """
        B, T = displacement.shape
        
        # 提取时序特征
        x = displacement.unsqueeze(1)  # (B, 1, T)
        feat = self.feature_extractor(x)  # (B, 64, 1)
        feat = feat.flatten(1)  # (B, 64)
        
        # 预测多项式系数
        coeffs = self.coeff_predictor(feat)  # (B, (poly_order+1) * num_features)
        coeffs = coeffs.view(B, self.num_features, self.poly_order + 1)
        
        # 稀疏门控
        gate = self.sparse_gate(feat)  # (B, num_features)
        
        # 加权多项式系数
        weighted_coeffs = coeffs * gate.unsqueeze(-1)  # (B, num_features, poly_order+1)
        final_coeffs = weighted_coeffs.sum(dim=1)  # (B, poly_order+1)
        
        # 构建多项式
        volume = torch.zeros_like(displacement)
        for k in range(self.poly_order + 1):
            volume = volume + final_coeffs[:, k:k+1] * displacement ** k
        
        return volume


class SpiRadarModel(nn.Module):
    """
    SpiRadar: 完整的非接触雷达肺活量测定模型
    
    论文：arXiv:2607.19798
    应用：座舱呼吸监测、哮喘管理、COPD 监测
    """
    def __init__(self):
        super().__init__()
        self.preprocessing = PhysiologicalPreprocessing(fs=200)
        self.transform = SparsePolynomialTransform(poly_order=5, num_features=8)
        
    def forward(self, radar_phase: torch.Tensor) -> dict:
        """
        Args:
            radar_phase: 雷达回波相位 (B, T)
        Returns:
            results: dict with volume, breathing_rate, clinical_params
        """
        # 预处理
        resp_signal, motion = self.preprocessing(radar_phase)
        
        # 转换为体积曲线
        volume = self.transform(resp_signal)
        
        # 计算临床参数
        pev1 = volume[:, :int(0.1 * volume.shape[1])].max(dim=1)[0]  # 简化
        fvc = volume.max(dim=1)[0]
        pef = volume[:, :int(0.05 * volume.shape[1])].max(dim=1)[0]
        
        return {
            'volume_curve': volume,
            'breathing_signal': resp_signal,
            'motion_artifact': motion,
            'FEV1': pev1,
            'FVC': fvc,
            'PEF': pef
        }


# ==================== 测试 ====================

if __name__ == "__main__":
    print("=" * 60)
    print("SpiRadar 非接触雷达肺活量测定测试")
    print("=" * 60)
    
    model = SpiRadarModel()
    
    # 模拟雷达回波相位（呼吸信号 + 运动伪影）
    T = 2000  # 10秒 @ 200Hz
    t = np.linspace(0, 10, T)
    breathing = 0.5 * np.sin(2 * np.pi * 0.3 * t)  # 呼吸 0.3 Hz
    motion = 0.2 * np.random.randn(T)  # 随机运动
    radar_phase = torch.tensor(breathing + motion, dtype=torch.float32).unsqueeze(0)
    
    results = model(radar_phase)
    
    print(f"输入: {radar_phase.shape} (B, T)")
    print(f"体积曲线: {results['volume_curve'].shape}")
    print(f"FEV1: {results['FEV1'].item():.3f} L")
    print(f"FVC: {results['FVC'].item():.3f} L")
    print(f"PEF: {results['PEF'].item():.3f} L/s")
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params / 1e3:.1f}K")

class CabinBreathingMonitoring:
    """座舱呼吸监测集成方案"""
    
    def __init__(self):
        self.integration_points = {
            'steering_wheel_radar': {
                'sensor': '60GHz FMCW 嵌入方向盘',
                'target': '驾驶员胸壁',
                'distance': '40-60cm',
                'application': '呼吸率 + 疲劳检测',
                'spiRadar_applicability': 'Level 1-2'
            },
            'roof_radar': {
                'sensor': '60GHz FMCW 车顶',
                'target': '全车乘员',
                'distance': '1-2m',
                'application': 'CPD + 乘员呼吸',
                'spiRadar_applicability': 'Level 1'
            },
            'seat_radar': {
                'sensor': '60GHz 嵌入座椅靠背',
                'target': '后排乘员',
                'distance': '5-20cm',
                'application': 'CPD 婴儿呼吸',
                'spiRadar_applicability': 'Level 1-3 (近距离最佳)'
            }
        }
    
    def get_deployment_priority(self):
        return {
            'P0': '方向盘雷达呼吸率（疲劳检测增强）',
            'P1': '座椅雷达 CPD 婴儿呼吸（ADAC 痛点）',
            'P2': 'SpiRadar 肺活量模型（哮喘驾驶员管理）',
            'P3': '后排乘员呼吸异常检测（突发疾病）'
        }