import torch
import torch.nn as nn
import torch.nn.functional as F

class CameraRadarFusion(nn.Module):
    """
    Murata+SmartEye风格的摄像头-雷达融合模型
    
    核心思路：
    1. 摄像头提供视觉上下文（姿态、行为、表情）
    2. 雷达提供深度和穿透检测（存在、距离、微动）
    3. 融合层学习两个模态的互补关系
    
    支持任务：DMS, OMS, CPD, OOP, 生命体征
    """
    
    def __init__(self, n_tasks=5, d_model=256):
        super().__init__()
        
        # 摄像头编码器
        self.camera_encoder = CameraEncoder(d_model=d_model)
        
        # 雷达编码器（60GHz range-Doppler图）
        self.radar_encoder = RadarEncoder(d_model=d_model)
        
        # 跨模态注意力
        self.cross_attention = nn.MultiheadAttention(
            d_model, num_heads=8, batch_first=True
        )
        
        # 多任务头
        self.task_heads = nn.ModuleDict({
            'dms_gaze': nn.Linear(d_model, 4),        # 视线4D
            'dms_fatigue': nn.Linear(d_model, 3),     # 疲劳等级
            'dms_distraction': nn.Linear(d_model, 3), # 分心等级
            'oms_occupant': nn.Linear(d_model, 4),    # 占位/分类
            'oms_posture': nn.Linear(d_model, 8),     # 姿态类别
            'cpd_presence': nn.Linear(d_model, 2),    # 儿童/无人
            'cpd_breathing': nn.Linear(d_model, 1),    # 呼吸率
            'oop_detection': nn.Linear(d_model, 5),   # OOP类别
            'airbag_deploy': nn.Linear(d_model, 3),   # 安全气囊策略
        })
        
    def forward(self, camera_img, radar_rd):
        """
        Args:
            camera_img: (B, 3, H, W) 座舱图像
            radar_rd: (B, 1, R, D) range-Doppler图
        
        Returns:
            dict of task outputs
        """
        # 编码
        cam_feat = self.camera_encoder(camera_img)  # (B, d_model)
        rad_feat = self.radar_encoder(radar_rd)     # (B, d_model)
        
        # 跨模态注意力
        cam_stack = cam_feat.unsqueeze(1)  # (B, 1, d)
        rad_stack = rad_feat.unsqueeze(1)
        
        # 摄像头用雷达增强（雷达提供深度/穿透信息）
        cam_enhanced, _ = self.cross_attention(cam_stack, rad_stack, rad_stack)
        cam_enhanced = cam_enhanced.squeeze(1)
        
        # 融合特征
        fused = cam_feat + rad_feat + cam_enhanced
        
        # 多任务输出
        outputs = {}
        for task_name, head in self.task_heads.items():
            outputs[task_name] = head(fused)
        
        return outputs


class CameraEncoder(nn.Module):
    """摄像头编码器（MobileNet风格轻量化）"""
    def __init__(self, d_model=256):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32), nn.ReLU6(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64), nn.ReLU6(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128), nn.ReLU6(),
            nn.Conv2d(128, d_model, 3, stride=2, padding=1),
            nn.BatchNorm2d(d_model), nn.ReLU6(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.proj = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        return self.proj(self.features(x).flatten(1))


class RadarEncoder(nn.Module):
    """60GHz雷达 range-Doppler 编码器"""
    def __init__(self, d_model=256):
        super().__init__()
        self.features = 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.MaxPool2d(2),
            nn.Conv2d(128, d_model, 3, padding=1),
            nn.BatchNorm2d(d_model), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.proj = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        return self.proj(self.features(x).flatten(1))


# 测试
if __name__ == "__main__":
    model = CameraRadarFusion(n_tasks=5, d_model=128)
    
    B = 2
    camera = torch.randn(B, 3, 224, 224)
    radar = torch.randn(B, 1, 64, 64)  # range-Doppler图
    
    outputs = model(camera, radar)
    
    print("=== 摄像头-雷达融合多任务输出 ===")
    for task, out in outputs.items():
        print(f"  {task}: {out.shape}")
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"\n参数量: {total_params/1e6:.2f}M")
    print(f"适合边缘部署: {'✅' if total_params < 50e6 else '⚠️ 需量化'}")