import numpy as np
from dataclasses import dataclass

@dataclass
class HelmholtzResonator:
    """
    Helmholtz共振器物理模型
    
    共振频率由腔体几何形状决定：
    f = (c/2π) * sqrt(A/(V*L))
    
    其中：
    - c: 声速 (343 m/s at 20°C)
    - A: 颈部截面积 (m²)
    - V: 腔体体积 (m³)
    - L: 颈部长度 (m)
    """
    cavity_volume: float = 0.001  # V: 1L = 0.001 m³
    neck_area: float = 3.14e-4   # A: 半径1cm的圆面积
    neck_length: float = 0.02     # L: 2cm
    sound_speed: float = 343.0    # c: 声速
    
    @property
    def resonant_frequency(self) -> float:
        """计算共振频率"""
        f = (self.sound_speed / (2 * np.pi)) * \
            np.sqrt(self.neck_area / (self.cavity_volume * self.neck_length))
        return f
    
    @property
    def quality_factor(self) -> float:
        """
        Q因子：共振峰锐度
        
        Q = 2π × (储能/每周期耗能)
        高Q → 窄带高增益 → 心音频段选择性放大
        """
        # 简化估算
        f = self.resonant_frequency
        # 颈部空气质量
        mass = 1.225 * self.neck_area * self.neck_length  # ρ_air * A * L
        # 弹簧常数
        k = (1.225 * self.sound_speed**2 * self.neck_area**2) / self.cavity_volume
        # 阻尼（辐射+热损耗简化）
        R = 1.225 * self.sound_speed * 2 * np.pi * (self.neck_area**2) / (2 * np.pi * f)
        Q = np.sqrt(mass * k) / R
        return Q
    
    def frequency_response(self, freqs: np.ndarray) -> np.ndarray:
        """
        频率响应：在共振频率附近的增益
        
        Returns:
            gain_db: 各频率的增益(dB)
        """
        f0 = self.resonant_frequency
        Q = max(self.quality_factor, 1)
        # 二阶共振传递函数
        x = freqs / f0
        gain = 1.0 / np.sqrt((1 - x**2)**2 + (x/Q)**2)
        return 20 * np.log10(gain)


# 设计目标：50-120 Hz心音频段
resonator = HelmholtzResonator(
    cavity_volume=0.003,    # 3L腔体
    neck_area=7.85e-5,      # 半径0.5cm圆面积
    neck_length=0.015       # 1.5cm颈部
)

f0 = resonator.resonant_frequency
Q = resonator.quality_factor
print(f"共振频率: {f0:.1f} Hz")
print(f"品质因子 Q: {Q:.1f}")
print(f"3dB带宽: {f0/Q:.1f} Hz")
print(f"峰值增益: {20*np.log10(Q):.1f} dB")

# 频率响应曲线
freqs = np.linspace(10, 200, 1000)
gain_db = resonator.frequency_response(freqs)
print(f"\n在50Hz增益: {gain_db[np.argmin(np.abs(freqs-50))]:.1f} dB")
print(f"在80Hz增益: {gain_db[np.argmin(np.abs(freqs-80))]:.1f} dB")
print(f"在120Hz增益: {gain_db[np.argmin(np.abs(freqs-120))]:.1f} dB")
print(f"在200Hz增益: {gain_db[np.argmin(np.abs(freqs-200))]:.1f} dB (抑制)")

import torch
import torch.nn as nn

class AcousticToECGNet(nn.Module):
    """
    声学信号→ECG波形重建网络
    
    论文核心创新：学习心音→心电的跨模态映射
    
    架构：
    - Encoder: 1D-CNN提取声学特征
    - Transformer: 时序建模
    - Decoder: 重建ECG波形
    """
    
    def __init__(self, 
                 audio_sample_rate: int = 1000,
                 ecg_sample_rate: int = 250,
                 d_model: int = 256,
                 n_heads: int = 8,
                 n_layers: int = 6):
        super().__init__()
        
        # 声学特征提取 (1D CNN)
        self.audio_encoder = nn.Sequential(
            nn.Conv1d(1, 32, kernel_size=15, stride=2, padding=7),
            nn.BatchNorm1d(32),
            nn.GELU(),
            nn.Conv1d(32, 64, kernel_size=15, stride=2, padding=7),
            nn.BatchNorm1d(64),
            nn.GELU(),
            nn.Conv1d(64, d_model, kernel_size=15, stride=2, padding=7),
            nn.BatchNorm1d(d_model),
            nn.GELU(),
        )
        
        # 位置编码
        self.pos_encoding = self._create_pos_encoding(5000, d_model)  # max 5000 tokens
        
        # Transformer时序建模
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=0.1,
            batch_first=True,
            activation='gelu'
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, n_layers)
        
        # ECG波形解码
        self.decoder = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Linear(d_model // 2, 1),
            nn.Tanh()  # ECG信号归一化到[-1, 1]
        )
        
        # 下采样比例
        self.downsample_ratio = audio_sample_rate // ecg_sample_rate
    
    def _create_pos_encoding(self, max_len, d_model):
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).float().unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-np.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return pe.unsqueeze(0)
    
    def forward(self, audio: torch.Tensor) -> torch.Tensor:
        """
        Args:
            audio: (B, T_audio) 声学信号
        
        Returns:
            ecg_pred: (B, T_ecg) 重建ECG波形
        """
        # CNN特征提取
        x = self.audio_encoder(audio.unsqueeze(1))  # (B, d_model, T_feat)
        x = x.transpose(1, 2)  # (B, T_feat, d_model)
        
        # 位置编码
        T = x.size(1)
        x = x + self.pos_encoding[:, :T, :].to(x.device)
        
        # Transformer时序建模
        x = self.transformer(x)
        
        # 解码ECG
        ecg_pred = self.decoder(x)  # (B, T_feat, 1)
        ecg_pred = ecg_pred.squeeze(-1)  # (B, T_feat)
        
        return ecg_pred


# R-peak检测（从重建ECG提取心率）
class RPeakDetector(nn.Module):
    """
    从重建ECG中检测R-peak（心搏事件）
    
    方法：Pan-Tompkins改进 + 深度学习精修
    """
    
    def __init__(self, fs: int = 250):
        super().__init__()
        self.fs = fs
        
        # 带通滤波（5-15Hz，QRS频段）
        nyq = fs / 2
        self.bandpass = nn.Sequential(
            nn.Conv1d(1, 1, kernel_size=15, padding=7, bias=False),
        )
        
        # 自适应阈值
        self.threshold = 0.5
        
    def forward(self, ecg: torch.Tensor) -> list:
        """
        Args:
            ecg: (B, T) ECG信号
        
        Returns:
            peaks: list of [(batch_idx, peak_location)]
        """
        x = ecg.unsqueeze(1)
        filtered = self.bandpass(x).squeeze(1)
        
        # 梯度
        diff = torch.diff(filtered, dim=1)
        squared = diff ** 2
        
        # 移动平均
        kernel = torch.ones(1, 1, self.fs // 4) / (self.fs // 4)
        moving_avg = nn.functional.conv1d(
            squared.unsqueeze(1), 
            kernel.unsqueeze(0),
            padding=self.fs // 8
        ).squeeze(1)[:, :squared.size(1)]
        
        # 峰值检测
        peaks = []
        for b in range(ecg.size(0)):
            signal = moving_avg[b].cpu().numpy()
            above = signal > self.threshold * signal.max()
            
            # 找上升沿
            rising = np.where(np.diff(above.astype(int)) == 1)[0]
            for idx in rising:
                # 在窗口内找精确峰值
                window = 20
                start = max(0, idx)
                end = min(len(filtered[b]), idx + window)
                peak_idx = start + torch.argmax(filtered[b, start:end]).item()
                peaks.append((b, peak_idx))
        
        return peaks


# 测试
if __name__ == "__main__":
    model = AcousticToECGNet()
    detector = RPeakDetector()
    
    # 模拟10秒声学输入
    audio = torch.randn(2, 10000)  # batch=2, 10s@1kHz
    ecg = model(audio)
    print(f"输入声学: {audio.shape}")
    print(f"输出ECG: {ecg.shape}")
    
    # R-peak检测
    peaks = detector(ecg)
    for b, p in peaks:
        print(f"Batch {b}: R-peak at sample {p} ({p/250:.2f}s)")
    
    # 计算心率
    if len(peaks) > 1:
        rr_intervals = np.diff([p for _, p in peaks]) / 250
        hr = 60 / np.mean(rr_intervals)
        print(f"估算心率: {hr:.1f} BPM")

# 座舱集成方案概念
class CabinAcousticECG:
    """
    座舱声学心电监测集成方案
    
    安装位置：座椅头枕或靠背
    优势：
    1. 无需身体接触（驾驶员无需佩戴任何设备）
    2. 被动放大不耗电（共振腔物理特性）
    3. 不受光照影响（暗光/强光均可工作）
    4. 隐私保护（不采集图像）
    
    挑战：
    1. 车内噪声干扰（发动机/风噪/路噪）
    2. 驾驶员移动导致信号丢失
    3. 共振腔体积限制（需小型化）
    4. 需要个性化校准
    """
    
    def __init__(self):
        # 噪声抑制策略
        self.noise_sources = {
            'engine': (20, 200),      # 发动机噪声频段
            'wind': (100, 1000),      # 风噪频段
            'road': (5, 50),          # 路面振动频段
            'cardiac': (50, 120),     # 目标心音频段
        }
        
    def snr_estimation(self, signal, noise):
        """
        信噪比估计
        
        座舱内目标SNR需要 > 10dB
        共振器增益约 20-30dB
        车内背景噪声约 40-60dB SPL
        心音强度约 30-40dB SPL（1m距离）
        → 增益后心音 50-70dB vs 噪声 40-60dB
        → SNR约 0-15dB，需进一步降噪
        """
        signal_power = np.mean(signal**2)
        noise_power = np.mean(noise**2)
        snr_db = 10 * np.log10(signal_power / noise_power)
        return snr_db
    
    def adaptive_noise_cancellation(self, primary, reference):
        """
        自适应噪声消除（LMS算法）
        
        primary: 共振器输出（心音+噪声）
        reference: 座舱麦克风参考噪声
        """
        N = len(reference)
        M = 32  # 滤波器阶数
        w = np.zeros(M)
        mu = 0.01  # 步长
        
        output = np.zeros(N)
        for n in range(M, N):
            x = reference[n-M:n][::-1]
            y = np.dot(w, x)
            e = primary[n] - y
            w = w + mu * e * x
            output[n] = e
        
        return output