"""
核心问题：
mmWave雷达接触式心率感知中，相干叠加的多散射体回波
可能抑制心跳成分，即使在相似观测几何下也会发生。

这意味着：不能假设"看起来好的测量就一定有好的心率"
必须学习评估每次测量的可观测性。

HEAR方案：
1. 训练双任务Transformer：同时预测可观测性分数+心率
2. 仅用仿真数据训练 → 零样本迁移到真实数据
3. 可观测性分数可用于选择性预测（只在高置信度时输出）
"""

# 示意代码,不可直接运行(为解释论文 Section 3 思想的简化示意,非官方实现)
import torch
import torch.nn as nn
import numpy as np
from typing import Tuple

class ControllableFMCWSimulator:
    """
    可控多散射体FMCW雷达仿真器
    
    论文 Section 3: 生成不同可观测条件训练数据
    
    核心思想：
    - 模拟多个散射体（胸壁、心脏、呼吸、衣物）
    - 控制各成分的幅度和相位
    - 自动生成可观测性标签
    """
    
    def __init__(
        self,
        freq_ghz: float = 60.0,
        num_samples: int = 1024,
        sample_rate: int = 1000
    ):
        self.freq = freq_ghz * 1e9  # Hz
        self.c = 3e8  # 光速
        self.num_samples = num_samples
        self.sample_rate = sample_rate
        
    def simulate(
        self,
        heart_rate_bpm: float = 72.0,
        breathing_rate_bpm: float = 18.0,
        heart_amplitude: float = 0.1e-3,  # 心脏微动幅度（米）
        breathing_amplitude: float = 1.0e-3,  # 呼吸幅度
        num_scatterers: int = 5,
        noise_level: float = 0.01,
        observability: str = 'high'  # 'high', 'medium', 'low'
    ) -> Tuple[np.ndarray, float]:
        """
        仿真FMCW雷达相位信号
        
        Args:
            heart_rate_bpm: 心率（次/分）
            breathing_rate_bpm: 呼吸率
            heart_amplitude: 心脏微动幅度
            breathing_amplitude: 呼吸幅度
            num_scatterers: 散射体数量
            noise_level: 噪声水平
            observability: 可观测性等级
        Returns:
            phase_signal: 相位信号
            true_observability: 真实可观测性分数
        """
        t = np.arange(self.num_samples) / self.sample_rate
        
        # 心跳信号
        heart_freq = heart_rate_bpm / 60  # Hz
        heart_signal = heart_amplitude * np.sin(2 * np.pi * heart_freq * t)
        
        # 呼吸信号
        breath_freq = breathing_rate_bpm / 60
        breath_signal = breathing_amplitude * np.sin(2 * np.pi * breath_freq * t)
        
        # 多散射体叠加
        scatterer_signals = []
        for i in range(num_scatterers):
            # 每个散射体有不同的距离和相位
            distance = 0.5 + 0.1 * i + np.random.randn() * 0.05
            phase_offset = np.random.uniform(0, 2 * np.pi)
            
            # 散射体反射信号（包含心脏+呼吸微动）
            s = (heart_signal + breath_signal) * np.cos(phase_offset) \
                + noise_level * np.random.randn(len(t))
            scatterer_signals.append(s)
        
        # 相干叠加
        combined = np.sum(scatterer_signals, axis=0)
        
        # 提取相位
        phase_signal = np.angle(combined + 1j * np.imag(combined))
        
        # 计算真实可观测性
        # 频谱中心跳峰与已知心率的吻合度
        fft = np.fft.rfft(phase_signal)
        freqs = np.fft.rfftfreq(len(phase_signal), d=1/self.sample_rate)
        
        heart_mask = np.abs(freqs - heart_freq) < 0.05
        heart_power = np.sum(np.abs(fft[heart_mask]))
        
        breath_mask = np.abs(freqs - breath_freq) < 0.05
        breath_power = np.sum(np.abs(fft[breath_mask]))
        
        total_power = np.sum(np.abs(fft))
        
        # 可观测性 = 心跳功率 / (总功率 - 呼吸功率)
        true_observability = heart_power / (total_power - breath_power + 1e-10)
        
        # 根据可观测性设置调整
        if observability == 'low':
            # 增加散射体干扰降低可观测性
            combined += 0.5 * heart_amplitude * np.random.randn(len(t))
            true_observability *= 0.3
        
        return phase_signal, float(true_observability)
    
    def generate_dataset(
        self,
        n_samples: int = 10000
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        生成训练数据集
        """
        signals = []
        heart_rates = []
        observabilities = []
        
        for _ in range(n_samples):
            # 随机参数
            hr = np.random.uniform(50, 120)  # 50-120 BPM
            br = np.random.uniform(12, 30)   # 12-30 BPM
            obs = np.random.choice(['high', 'medium', 'low'], p=[0.4, 0.4, 0.2])
            
            signal, true_obs = self.simulate(hr, br, observability=obs)
            
            signals.append(signal)
            heart_rates.append(hr)
            observabilities.append(true_obs)
        
        return (np.array(signals), np.array(heart_rates), 
                np.array(observabilities))


# 测试仿真器
if __name__ == "__main__":
    sim = ControllableFMCWSimulator(freq_ghz=60)
    
    # 生成一个样本
    signal, obs = sim.simulate(heart_rate_bpm=72, observability='high')
    print(f"信号长度: {len(signal)}")
    print(f"可观测性分数: {obs:.4f}")
    
    # 生成数据集
    signals, hrs, obs_list = sim.generate_dataset(100)
    print(f"\n数据集: {signals.shape}")
    print(f"心率范围: {hrs.min():.0f}-{hrs.max():.0f} BPM")
    print(f"可观测性范围: {obs_list.min():.4f}-{obs_list.max():.4f}")
    
    print("\n✅ FMCW仿真器测试通过")

# 示意代码,不可直接运行(为解释论文 Section 4 思想的简化示意,非官方实现)
class HEARTransformer(nn.Module):
    """
    HEAR: Heartbeat Estimation with Assessed Reliability
    
    论文 Section 4: 紧凑双任务Transformer
    - 任务1: 预测可观测性分数（0-1）
    - 任务2: 预测心率（BPM）
    
    输入: 频谱幅度 + 相对呼吸基频的频率
    """
    
    def __init__(
        self,
        input_dim: int = 256,  # 频谱bin数
        hidden_dim: int = 128,
        num_heads: int = 4,
        num_layers: int = 3,
        max_freq: float = 5.0  # 最大频率 Hz
    ):
        super().__init__()
        self.input_dim = input_dim
        self.max_freq = max_freq
        
        # 输入嵌入
        self.input_proj = nn.Linear(input_dim * 2, hidden_dim)  # 幅度+相对频率
        
        # 位置编码
        self.pos_encoding = nn.Parameter(torch.randn(1, input_dim, hidden_dim))
        
        # Transformer编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            activation='gelu',
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # 池化
        self.pool = nn.AdaptiveAvgPool1d(1)
        
        # 可观测性预测头
        self.obs_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1),
            nn.Sigmoid()
        )
        
        # 心率预测头
        self.hr_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1)
        )
        
    def forward(
        self,
        spectral_magnitudes: torch.Tensor,
        frequencies: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            spectral_magnitudes: 频谱幅度 (B, F)
            frequencies: 频率值 (B, F) in Hz
        Returns:
            obs_score: 可观测性分数 (B, 1), 0-1
            heart_rate: 心率预测 (B, 1), BPM
        """
        B, F = spectral_magnitudes.shape
        
        # 拼接幅度和频率
        x = torch.cat([
            spectral_magnitudes.unsqueeze(-1),
            frequencies.unsqueeze(-1)
        ], dim=-1)  # (B, F, 2)
        
        # 投影
        x = self.input_proj(x)  # (B, F, hidden)
        
        # 位置编码
        x = x + self.pos_encoding[:, :F, :]
        
        # Transformer
        x = self.transformer(x)  # (B, F, hidden)
        
        # 池化
        x = self.pool(x.transpose(1, 2)).squeeze(-1)  # (B, hidden)
        
        # 双任务输出
        obs_score = self.obs_head(x)    # (B, 1)
        heart_rate = self.hr_head(x)    # (B, 1)
        
        return obs_score, heart_rate
    
    def predict_with_confidence(
        self,
        spectral_magnitudes: torch.Tensor,
        frequencies: torch.Tensor,
        coverage: float = 0.5
    ) -> dict:
        """
        选择性预测：基于可观测性分数选择性输出
        
        Args:
            coverage: 输出覆盖率（0-1），越高输出越多
        Returns:
            results: 包含预测和置信度的字典
        """
        with torch.no_grad():
            obs_score, hr = self.forward(spectral_magnitudes, frequencies)
            
            # 基于可观测性分数排序
            scores = obs_score.squeeze().cpu().numpy()
            n_output = int(len(scores) * coverage)
            top_indices = np.argsort(scores)[-n_output:]
            
            results = {
                'all_heart_rates': hr.squeeze().cpu().numpy(),
                'all_observability': scores,
                'selected_indices': top_indices,
                'selected_hr': hr.squeeze().cpu().numpy()[top_indices],
                'coverage': coverage,
                'mean_confidence': scores[top_indices].mean() if n_output > 0 else 0
            }
        
        return results


# 测试
if __name__ == "__main__":
    model = HEARTransformer(input_dim=256, hidden_dim=128)
    
    # 模拟输入
    B, F = 4, 256
    magnitudes = torch.randn(B, F)
    freqs = torch.linspace(0, 5, F).unsqueeze(0).expand(B, -1)
    
    obs, hr = model(magnitudes, freqs)
    print(f"可观测性分数: {obs.shape} → {obs.squeeze().tolist()}")
    print(f"心率预测: {hr.shape} → {hr.squeeze().tolist()}")
    
    # 选择性预测
    result = model.predict_with_confidence(magnitudes, freqs, coverage=0.5)
    print(f"\n选择性预测 (50%覆盖):")
    print(f"  选中数: {len(result['selected_indices'])}")
    print(f"  平均置信度: {result['mean_confidence']:.4f}")
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n模型参数量: {total_params:,} ({total_params/1e6:.2f}M)")
    print("✅ HEAR模型测试通过")

# 示意代码,不可直接运行(仅为集成路线的梳理笔记)
class IMSRadarVitalSigns:
    """IMS 60GHz雷达生命体征监测方案"""
    
    def __init__(self):
        # 60GHz雷达已在CPD中引入（参考Murata方案）
        # HEAR框架可直接集成
        self.radar_config = {
            'frequency': '60 GHz',
            'application': 'CPD + 心率监测',
            'advantage': '已有雷达硬件，复用做生命体征',
        }
        
    def integration_plan(self):
        return {
            'Phase 1': 'CPD雷达检测活体（呼吸频率）',
            'Phase 2': '+ 心率监测（HEAR框架）',
            'Phase 3': '+ 可观测性评估（选择性输出）',
            'Phase 4': '健康座舱（持续生命体征监控）',
        }