# mmWave雷达回波信号预处理后的输入
# shape: (batch_size, channels, time_steps)
# 典型配置：8通道 × 1024时间步
radar_input = torch.randn(32, 8, 1024)  # batch=32

import torch
import torch.nn as nn
from torchdiffeq import odeint

class ODEFunc(nn.Module):
    """
    ODE函数：定义ECG波形的动力学先验
    
    论文Section 3.2：使用ODE约束ECG波形的生物合理性
    核心思想：ECG波形遵循心脏电生理动力学，
    可用微分方程描述PQRST波的生成过程
    """
    def __init__(self, hidden_dim: int = 64):
        super().__init__()
        self.fc1 = nn.Linear(hidden_dim, hidden_dim * 2)
        self.fc2 = nn.Linear(hidden_dim * 2, hidden_dim)
        self.activation = nn.Softplus()  # 平滑激活，适合生理信号
        
    def forward(self, t: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        """
        ODE右端函数: dx/dt = f(t, x)
        
        Args:
            t: 时间点, scalar
            x: 隐状态, shape=(B, hidden_dim)
            
        Returns:
            dx/dt: 隐状态的时间导数
        """
        h = self.activation(self.fc1(x))
        dxdt = self.fc2(h)
        return dxdt


class MorphologicalDecoder(nn.Module):
    """
    ODE-based 形态解码器
    
    将共享特征映射到ECG波形的PQRST形态
    使用Neural ODE保证波形的动力学合理性
    """
    def __init__(self, hidden_dim: int = 512, ode_dim: int = 64, 
                 ecg_length: int = 256):
        super().__init__()
        self.proj = nn.Linear(hidden_dim, ode_dim)
        self.ode_func = ODEFunc(ode_dim)
        self.output_proj = nn.Linear(ode_dim, 1)
        self.ecg_length = ecg_length
        
    def forward(self, shared_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            shared_features: (B, 512) 来自共享骨干
            
        Returns:
            ecg_morph: (B, ecg_length) 重建的ECG形态
        """
        # 初始状态
        z0 = self.proj(shared_features)  # (B, 64)
        
        # ODE求解：从t=0积分到t=ecg_length
        t_span = torch.linspace(0, 1, self.ecg_length).to(z0.device)
        z_traj = odeint(self.ode_func, z0, t_span)  # (T, B, 64)
        
        # 投影到ECG采样点
        ecg_morph = self.output_proj(z_traj).squeeze(-1)  # (T, B)
        return ecg_morph.transpose(0, 1)  # (B, T)

class AnchorDecoder(nn.Module):
    """
    R峰定位解码器
    
    二分类任务：判断每个时间步是否为R峰
    使用1D-CNN提取局部时序模式
    """
    def __init__(self, hidden_dim: int = 512, kernel_size: int = 7):
        super().__init__()
        self.conv = nn.Conv1d(hidden_dim, 128, kernel_size, padding=kernel_size//2)
        self.bn = nn.BatchNorm1d(128)
        self.act = nn.ReLU()
        self.head = nn.Conv1d(128, 1, 1)
        
    def forward(self, shared_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            shared_features: (B, 512, T) 时序特征
            
        Returns:
            anchor_logits: (B, T) R峰位置的二分类logits
        """
        x = self.act(self.bn(self.conv(shared_features)))
        return self.head(x).squeeze(1)  # (B, T)

class CycleLengthDecoder(nn.Module):
    """
    周期长度解码器
    
    回归任务：预测每个心动周期的Peak-to-Peak间隔（PPI）
    """
    def __init__(self, hidden_dim: int = 512):
        super().__init__()
        self.lstm = nn.LSTM(hidden_dim, 128, batch_first=True, num_layers=2)
        self.head = nn.Linear(128, 1)
        
    def forward(self, shared_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            shared_features: (B, 512, T)
            
        Returns:
            cycle_lengths: (B, T) 每个时间步的周期长度预测（秒）
        """
        lstm_out, _ = self.lstm(shared_features.transpose(1, 2))
        return self.head(lstm_out).squeeze(-1)  # (B, T)

class EccentricGradientAlignment:
    """
    EGA: 偏心梯度对齐策略
    
    核心思想：
    1. 计算每个任务梯度的主方向
    2. 通过正交投影消除任务间的冲突方向
    3. 用"偏心向量"优先照顾收敛慢的困难任务
    
    论文Section 4.3
    """
    
    def __init__(self, num_tasks: int = 3, beta: float = 0.5):
        self.num_tasks = num_tasks
        self.beta = beta  # 偏心权重
        
    def compute_task_gradients(self, losses: list, 
                                shared_params: list) -> list:
        """计算每个任务对共享参数的梯度"""
        task_grads = []
        for loss in losses:
            grad = torch.autograd.grad(
                loss, shared_params, 
                retain_graph=True, 
                create_graph=False
            )
            # 展平为向量
            flat_grad = torch.cat([g.flatten() for g in grad])
            task_grads.append(flat_grad)
        return task_grads
    
    def eccentric_alignment(self, task_grads: list, 
                            task_losses: list) -> torch.Tensor:
        """
        偏心梯度对齐
        
        Args:
            task_grads: 各任务梯度列表, 每个 (D,)
            task_losses: 各任务当前loss值
            
        Returns:
            aligned_grad: 对齐后的合并梯度
        """
        # 1. 找到最难的任务（loss最大的任务）
        hardest_idx = torch.argmax(torch.tensor(task_losses))
        
        # 2. 以最难任务为"偏心"方向
        eccentric_vec = task_grads[hardest_idx]
        eccentric_norm = torch.norm(eccentric_vec) + 1e-8
        
        # 3. 对其他任务梯度做正交投影，消除与偏心方向冲突的分量
        aligned_grads = [eccentric_vec]  # 困难任务保持原样
        for i, grad in enumerate(task_grads):
            if i == hardest_idx:
                continue
            # 计算与偏心方向的余弦相似度
            cos_sim = torch.dot(grad, eccentric_vec) / (
                torch.norm(grad) * eccentric_norm + 1e-8
            )
            if cos_sim < 0:
                # 冲突：投影到偏心方向的正交补空间
                proj = (torch.dot(grad, eccentric_vec) / 
                        (eccentric_norm ** 2)) * eccentric_vec
                aligned = grad - proj
            else:
                # 不冲突：保持原样
                aligned = grad
            aligned_grads.append(aligned)
        
        # 4. 加权合并（偏心权重）
        weights = torch.zeros(self.num_tasks)
        weights[hardest_idx] = self.beta
        other_weight = (1 - self.beta) / (self.num_tasks - 1)
        for i in range(self.num_tasks):
            if i != hardest_idx:
                weights[i] = other_weight
        
        aligned_grad = sum(w * g for w, g in zip(weights, aligned_grads))
        return aligned_grad

class RadarODEMTL(nn.Module):
    """
    radarODE-MTL 完整模型
    
    多任务学习框架：
    - 共享ResNet骨干
    - 3个并行解码器（形态/R峰/周期）
    - EGA梯度对齐（训练阶段）
    """
    def __init__(self, radar_channels: int = 8, 
                 hidden_dim: int = 512,
                 ecg_length: int = 256):
        super().__init__()
        # 共享骨干
        self.backbone = ResNetBackbone(radar_channels, hidden_dim)
        # 三个解码器
        self.morph_decoder = MorphologicalDecoder(hidden_dim, ecg_length=ecg_length)
        self.anchor_decoder = AnchorDecoder(hidden_dim)
        self.cycle_decoder = CycleLengthDecoder(hidden_dim)
        
    def forward(self, radar_input: torch.Tensor):
        """
        Args:
            radar_input: (B, C, T) mmWave雷达信号
            
        Returns:
            ecg_morph: (B, ecg_length) ECG形态
            anchor_logits: (B, T) R峰logits
            cycle_lengths: (B, T) 周期长度
        """
        shared = self.backbone(radar_input)  # (B, 512, T)
        
        ecg_morph = self.morph_decoder(shared)
        anchor_logits = self.anchor_decoder(shared)
        cycle_lengths = self.cycle_decoder(shared)
        
        return ecg_morph, anchor_logits, cycle_lengths


# ===== 训练代码 =====
def train_radarode_mtl(model, dataloader, optimizer, 
                       ega: EccentricGradientAlignment,
                       num_epochs: int = 100):
    """
    radarODE-MTL 训练流程
    
    关键步骤：
    1. 前向传播得到三个任务输出
    2. 分别计算三个loss
    3. EGA对齐梯度
    4. 反向传播
    """
    model.train()
    morph_loss_fn = nn.MSELoss()
    anchor_loss_fn = nn.BCEWithLogitsLoss()
    cycle_loss_fn = nn.SmoothL1Loss()
    
    for epoch in range(num_epochs):
        for batch in dataloader:
            radar, ecg_gt, anchor_gt, cycle_gt = batch
            radar, ecg_gt = radar.cuda(), ecg_gt.cuda()
            anchor_gt, cycle_gt = anchor_gt.cuda(), cycle_gt.cuda()
            
            # 前向
            ecg_pred, anchor_pred, cycle_pred = model(radar)
            
            # 三个任务loss
            loss_morph = morph_loss_fn(ecg_pred, ecg_gt)
            loss_anchor = anchor_loss_fn(anchor_pred, anchor_gt)
            loss_cycle = cycle_loss_fn(cycle_pred, cycle_gt)
            
            # EGA梯度对齐
            shared_params = list(model.backbone.parameters())
            task_grads = ega.compute_task_gradients(
                [loss_morph, loss_anchor, loss_cycle],
                shared_params
            )
            task_losses = [loss_morph.item(), loss_anchor.item(), loss_cycle.item()]
            aligned_grad = ega.eccentric_alignment(task_grads, task_losses)
            
            # 手动梯度更新
            optimizer.zero_grad()
            # 用对齐后的梯度更新共享参数
            pointer = 0
            for param in shared_params:
                num = param.numel()
                param.grad = aligned_grad[pointer:pointer+num].view_as(param).clone()
                pointer += num
            optimizer.step()
            
            # 分别更新各解码器
            for loss, decoder_params in [
                (loss_morph, model.morph_decoder.parameters()),
                (loss_anchor, model.anchor_decoder.parameters()),
                (loss_cycle, model.cycle_decoder.parameters()),
            ]:
                decoder_grad = torch.autograd.grad(loss, decoder_params, 
                                                    retain_graph=True)
                for p, g in zip(decoder_params, decoder_grad):
                    if p.grad is None:
                        p.grad = g.clone()
                    else:
                        p.grad += g
            optimizer.step()
            
        print(f"Epoch {epoch+1}/{num_epochs} | "
              f"Morph: {loss_morph:.4f} | "
              f"Anchor: {loss_anchor:.4f} | "
              f"Cycle: {loss_cycle:.4f}")

# IMS集成示例：雷达ECG用于驾驶员心脏健康监测
class DriverCardiacMonitor:
    """
    基于mmWave雷达的驾驶员心脏监测模块
    
    应用场景：
    1. 实时心率/HRV监测 → 疲劳/压力评估
    2. 心律异常检测 → 紧急停车预警
    3. 心率变异性(HRV) → 自主神经平衡评估
    
    硬件复用：
    - 60GHz CPD雷达 → 同时做CPD和ECG重建
    - 无需接触式电极，驾驶员无感知
    """
    def __init__(self, radar_config):
        self.ecg_model = RadarODEMTL(radar_channels=radar_config.channels)
        self.ecg_model.load_state_dict(torch.load("radarode_mtl.pth"))
        self.ecg_model.eval()
        
    def process_radar_frame(self, radar_adc: np.ndarray) -> dict:
        """
        处理一帧雷达ADC数据
        
        Args:
            radar_adc: (C, T) 雷达原始ADC数据
            
        Returns:
            {
                'ecg': 重建ECG波形,
                'heart_rate': 心率(bpm),
                'hrv_rmssd': HRV指标,
                'r_peaks': R峰位置列表,
                'anomaly': 是否检测到心律异常
            }
        """
        with torch.no_grad():
            radar_tensor = torch.FloatTensor(radar_adc).unsqueeze(0).cuda()
            ecg, anchors, cycles = self.ecg_model(radar_tensor)
            
        ecg_np = ecg.squeeze().cpu().numpy()
        
        # 提取心率
        r_peak_indices = torch.where(anchors.squeeze() > 0.5)[0].cpu().numpy()
        if len(r_peak_indices) > 1:
            rr_intervals = np.diff(r_peak_indices) / 500.0  # 500Hz
            heart_rate = 60.0 / np.mean(rr_intervals)
            rmssd = np.sqrt(np.mean(np.diff(rr_intervals) ** 2)) * 1000
        else:
            heart_rate = 0
            rmssd = 0
            
        return {
            'ecg': ecg_np,
            'heart_rate': heart_rate,
            'hrv_rmssd': rmssd,
            'r_peaks': r_peak_indices,
            'anomaly': self._detect_anomaly(ecg_np, heart_rate)
        }

"""
radarODE-MTL 测试脚本

测试内容：
1. 模型前向传播
2. ECG重建质量评估
3. R峰检测准确率
4. 噪声鲁棒性测试
"""

import torch
import numpy as np
from scipy.signal import find_peaks
from scipy.stats import pearsonr

def test_forward_pass():
    """测试前向传播"""
    model = RadarODEMTL(radar_channels=8, ecg_length=256)
    model.eval()
    
    # 模拟雷达输入
    radar = torch.randn(4, 8, 1024)  # batch=4
    ecg, anchor, cycle = model(radar)
    
    assert ecg.shape == (4, 256), f"ECG shape error: {ecg.shape}"
    assert anchor.shape == (4, 1024), f"Anchor shape error: {anchor.shape}"
    assert cycle.shape == (4, 1024), f"Cycle shape error: {cycle.shape}"
    print("✓ Forward pass test passed")
    return model


def test_ecg_quality():
    """测试ECG重建质量"""
    model = test_forward_pass()
    
    # 生成模拟ECG（使用高斯函数模拟PQRST）
    t = np.linspace(0, 1, 256)
    ecg_gt = (
        0.1 * np.exp(-(t - 0.1)**2 / 0.001) +   # P wave
        -0.3 * np.exp(-(t - 0.2)**2 / 0.0005) + # Q wave
        1.0 * np.exp(-(t - 0.25)**2 / 0.0005) + # R wave
        -0.4 * np.exp(-(t - 0.3)**2 / 0.0005) + # S wave
        0.3 * np.exp(-(t - 0.4)**2 / 0.002) +   # T wave
    )
    
    # 模拟雷达输入（ECG的变形版本+噪声）
    radar_sim = np.random.randn(8, 1024) * 0.1
    for i in range(8):
        radar_sim[i] += np.interp(np.linspace(0, 1, 1024), t, ecg_gt) * (0.5 + 0.1*i)
    
    radar_tensor = torch.FloatTensor(radar_sim).unsqueeze(0)
    ecg_pred, _, _ = model(radar_tensor)
    ecg_pred_np = ecg_pred.squeeze().detach().numpy()
    
    # 计算PCC
    ecg_gt_resampled = np.interp(np.linspace(0, 1, 256), t, ecg_gt)
    pcc, _ = pearsonr(ecg_pred_np, ecg_gt_resampled)
    print(f"ECG重建 PCC: {pcc:.2f}")
    assert pcc > 0.5, "PCC too low"
    print("✓ ECG quality test passed")


def test_noise_robustness():
    """测试噪声鲁棒性"""
    model = test_forward_pass()
    
    t = np.linspace(0, 1, 1024)
    clean_signal = np.sin(2 * np.pi * 5 * t)
    
    for snr_db in [20, 10, 0, -5]:
        noise_power = np.var(clean_signal) / (10 ** (snr_db / 10))
        noisy = clean_signal + np.random.randn(1024) * np.sqrt(noise_power)
        
        radar = torch.FloatTensor(noisy).unsqueeze(0).unsqueeze(0).expand(1, 8, 1024)
        ecg, _, _ = model(radar)
        
        signal_power = np.var(ecg.squeeze().detach().numpy())
        print(f"SNR={snr_db:>3}dB | Output variance={signal_power:.4f}")
    
    print("✓ Noise robustness test passed")


def test_r_peak_detection():
    """测试R峰检测"""
    model = test_forward_pass()
    
    # 生成含R峰的ECG
    t = np.linspace(0, 5, 2560)
    ecg = np.zeros_like(t)
    for i in range(5):  # 5个心跳
        center = 0.5 + i * 1.0
        ecg += np.exp(-(t - center)**2 / 0.001)
    
    # 模拟雷达输入
    radar = np.tile(ecg, (8, 1))[:, :1024]
    radar += np.random.randn(8, 1024) * 0.1
    
    radar_tensor = torch.FloatTensor(radar).unsqueeze(0)
    _, anchor_logits, _ = model(radar_tensor)
    
    # 检测R峰
    anchor_probs = torch.sigmoid(anchor_logits).squeeze().detach().numpy()
    peaks, _ = find_peaks(anchor_probs, height=0.5, distance=100)
    
    print(f"检测到 {len(peaks)} 个R峰")
    assert len(peaks) >= 3, "R峰检测不足"
    print("✓ R-peak detection test passed")


if __name__ == "__main__":
    print("=" * 60)
    print("radarODE-MTL 测试套件")
    print("=" * 60)
    
    test_forward_pass()
    test_ecg_quality()
    test_noise_robustness()
    test_r_peak_detection()
    
    print("\n" + "=" * 60)
    print("所有测试通过 ✓")
    print("=" * 60)