import torch
import torch.nn as nn
import numpy as np
from scipy import signal as sig

class Radar4DCabinPerception(nn.Module):
    """4D雷达座舱感知多任务模型"""
    
    def __init__(self, n_range=128, n_doppler=64, n_azimuth=32, 
                 n_elevation=16):
        super().__init__()
        
        # 输入: 4D张量 (B, 1, Elevation, Azimuth, Range, Doppler)
        # 简化为两个3D输入
        
        # Range-Azimuth-Elevation (空间信息)
        self.spatial_encoder = nn.Sequential(
            nn.Conv3d(1, 32, 3, padding=1),
            nn.BatchNorm3d(32), nn.ReLU(),
            nn.MaxPool3d(2),
            
            nn.Conv3d(32, 64, 3, padding=1),
            nn.BatchNorm3d(64), nn.ReLU(),
            nn.MaxPool3d(2),
            
            nn.Conv3d(64, 128, 3, padding=1),
            nn.BatchNorm3d(128), nn.ReLU(),
            nn.AdaptiveAvgPool3d(1),
        )
        
        # Range-Doppler (速度信息)
        self.velocity_encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32), nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64), nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        
        # 融合
        self.fusion = nn.Sequential(
            nn.Linear(128 * 2, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
        )
        
        # 多任务头
        self.occupant_head = nn.Linear(256, 5)      # 0-4人
        self.cpd_head = nn.Linear(256, 2)            # 有人/无人
        self.breathing_head = nn.Linear(256, 1)      # 呼吸频率
        self.posture_head = nn.Linear(256, 5)        # 姿态分类
        self.gesture_head = nn.Linear(256, 8)        # 手势分类
        
    def forward(self, spatial_data, velocity_data):
        """
        Args:
            spatial_data: (B, 1, E, A, R) 空间RDA
            velocity_data: (B, 1, R, D) Range-Doppler
        
        Returns:
            dict of task outputs
        """
        spatial_feat = self.spatial_encoder(spatial_data).flatten(1)
        velocity_feat = self.velocity_encoder(velocity_data).flatten(1)
        
        fused = self.fusion(torch.cat([spatial_feat, velocity_feat], dim=1))
        
        return {
            'occupant_count': self.occupant_head(fused),
            'cpd_presence': self.cpd_head(fused),
            'breathing_rate': self.breathing_head(fused),
            'posture': self.posture_head(fused),
            'gesture': self.gesture_head(fused),
        }


# 微多普勒生命体征提取
class VitalSignsExtractor:
    """4D雷达微多普勒生命体征提取"""
    
    def __init__(self, fs=20, range_bin=None):
        self.fs = fs  # 帧率
        self.range_bin = range_bin  # 目标距离bin
        
    def extract_breathing(self, range_doppler_sequence):
        """
        从Range-Doppler序列提取呼吸
        
        Args:
            range_doppler_sequence: (T, R, D) 时序R-D图
        
        Returns:
            breathing_rate: 呼吸频率 (Hz)
            breathing_signal: 呼吸波形
        """
        # 1. 找到目标距离bin（最大能量）
        if self.range_bin is None:
            energy = np.sum(np.abs(range_doppler_sequence)**2, axis=2)
            self.range_bin = np.argmax(np.sum(energy, axis=0))
        
        # 2. 提取目标bin的相位序列
        target = range_doppler_sequence[:, self.range_bin, :]
        
        # 3. 低速Doppler区域（呼吸0.1-0.5Hz对应低速）
        low_doppler = target[:, :8]  # 前8个Doppler bin
        phase = np.angle(np.sum(low_doppler, axis=1))
        phase = np.unwrap(phase)
        
        # 4. 带通滤波 0.1-0.5 Hz
        b, a = sig.butter(4, [0.1, 0.5], btype='band', fs=self.fs)
        breathing_signal = sig.filtfilt(b, a, phase)
        
        # 5. FFT找呼吸频率
        freqs = np.fft.rfftfreq(len(breathing_signal), 1/self.fs)
        spectrum = np.abs(np.fft.rfft(breathing_signal))
        breathing_rate = freqs[np.argmax(spectrum)]
        
        return breathing_rate, breathing_signal
    
    def extract_heartbeat(self, range_doppler_sequence):
        """提取心跳（0.8-2.0 Hz）"""
        if self.range_bin is None:
            energy = np.sum(np.abs(range_doppler_sequence)**2, axis=2)
            self.range_bin = np.argmax(np.sum(energy, axis=0))
        
        target = range_doppler_sequence[:, self.range_bin, :]
        low_doppler = target[:, :8]
        phase = np.unwrap(np.angle(np.sum(low_doppler, axis=1)))
        
        # 去除呼吸成分
        b_breath, a_breath = sig.butter(4, [0.1, 0.5], btype='band', fs=self.fs)
        breathing = sig.filtfilt(b_breath, a_breath, phase)
        residual = phase - sig.filtfilt(
            sig.butter(4, [0.1, 0.5], btype='band', fs=self.fs)[0],
            sig.butter(4, [0.1, 0.5], btype='band', fs=self.fs)[1],
            phase
        )
        
        # 心跳带通 0.8-2.0 Hz
        b_heart, a_heart = sig.butter(4, [0.8, 2.0], btype='band', fs=self.fs)
        heart_signal = sig.filtfilt(b_heart, a_heart, residual)
        
        freqs = np.fft.rfftfreq(len(heart_signal), 1/self.fs)
        spectrum = np.abs(np.fft.rfft(heart_signal))
        heart_rate = freqs[np.argmax(spectrum)]
        
        return heart_rate, heart_signal


# 测试
if __name__ == "__main__":
    # 模型测试
    model = Radar4DCabinPerception()
    spatial = torch.randn(2, 1, 16, 32, 128)  # (B, C, E, A, R)
    velocity = torch.randn(2, 1, 128, 64)      # (B, C, R, D)
    
    outputs = model(spatial, velocity)
    print("=== 4D雷达座舱感知多任务 ===")
    for task, out in outputs.items():
        print(f"  {task}: {out.shape}")
    
    params = sum(p.numel() for p in model.parameters())
    print(f"\n参数量: {params/1e6:.2f}M")
    
    # 生命体征提取测试
    print("\n=== 生命体征提取 ===")
    rd_seq = np.random.randn(200, 128, 64) + 1j * np.random.randn(200, 128, 64)
    # 添加模拟呼吸信号
    t = np.arange(200) / 20
    breathing_mod = np.sin(2 * np.pi * 0.25 * t)  # 15 breaths/min
    rd_seq[:, 50, 0] += breathing_mod  # 注入到特定bin
    
    extractor = VitalSignsExtractor(fs=20)
    br_rate, br_signal = extractor.extract_breathing(rd_seq)
    print(f"  呼吸频率: {br_rate:.2f} Hz ({br_rate*60:.0f} breaths/min)")