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

class CabinSensorFusionSystem:
    """
    座舱多模态传感器融合系统
    
    三级融合架构：
    1. 早期融合（像素级）：ToF深度+RGB-IR对齐合并
    2. 中期融合（特征级）：摄像头特征+雷达特征交叉注意力
    3. 晚期融合（决策级）：各模态独立决策+置信度加权
    """
    
    SENSORS = {
        'dms_ir': {
            'type': 'near-infrared camera',
            'resolution': '1280x960',
            'fps': 30,
            'wavelength': '940nm',
            'primary_function': 'gaze_estimation, fatigue_detection',
            'fusion_role': 'feature_level',
        },
        'oms_wide': {
            'type': 'wide-angle camera',
            'resolution': '1920x1080',
            'fps': 30,
            'fov': '120°',
            'primary_function': 'occupant_classification, seatbelt',
            'fusion_role': 'feature_level',
        },
        'radar_60ghz': {
            'type': 'mmWave FMCW radar',
            'frequency': '60GHz',
            'channels': '2T4R',
            'range_resolution': '3.75cm',
            'fps': 20,
            'primary_function': 'CPD, vital_signs',
            'fusion_role': 'decision_level',
        },
        'tof_depth': {
            'type': 'time-of-flight depth camera',
            'resolution': '320x240',
            'fps': 30,
            'range': '0.3-3.0m',
            'primary_function': '3D_pose, OOP',
            'fusion_role': 'pixel_level',
        },
    }


class EarlyFusionModule(nn.Module):
    """
    早期融合：ToF深度图 + RGB-IR图像像素级合并
    
    适用：3D姿态估计 + OOP检测
    将深度信息作为RGB-IR的第4通道
    """
    
    def __init__(self):
        super().__init__()
        # 共享骨干网络处理4通道输入
        self.backbone = nn.Sequential(
            nn.Conv2d(4, 32, 3, stride=2, padding=1),  # 4通道: R, G, B/IR, Depth
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((20, 15)),
        )
        
    def forward(self, rgb_ir: torch.Tensor, depth: torch.Tensor) -> torch.Tensor:
        """
        Args:
            rgb_ir: (B, 3, H, W) RGB或IR图像
            depth: (B, 1, H, W) 深度图
        Returns:
            fused_feat: (B, 128, 20, 15) 融合特征
        """
        # 确保尺寸对齐
        if rgb_ir.shape[-2:] != depth.shape[-2:]:
            depth = F.interpolate(
                depth, size=rgb_ir.shape[-2:], 
                mode='bilinear', align_corners=False
            )
        
        # 像素级拼接
        fused_input = torch.cat([rgb_ir, depth], dim=1)  # (B, 4, H, W)
        
        # 共享特征提取
        fused_feat = self.backbone(fused_input)
        return fused_feat


class MidLevelFusionModule(nn.Module):
    """
    中期融合：摄像头特征 + 雷达特征交叉注意力
    
    适用：CPD + 乘员分类 + 生命体征
    摄像头提供空间细节，雷达提供穿透/微动
    """
    
    def __init__(self, cam_dim: int = 128, radar_dim: int = 64,
                 fusion_dim: int = 128):
        super().__init__()
        
        # 摄像头特征投影
        self.cam_proj = nn.Linear(cam_dim, fusion_dim)
        
        # 雷达特征投影
        self.radar_proj = nn.Linear(radar_dim, fusion_dim)
        
        # 双向交叉注意力
        self.cam_to_radar_attn = nn.MultiheadAttention(
            embed_dim=fusion_dim, num_heads=4, batch_first=True
        )
        self.radar_to_cam_attn = nn.MultiheadAttention(
            embed_dim=fusion_dim, num_heads=4, batch_first=True
        )
        
        # 层归一化
        self.norm1 = nn.LayerNorm(fusion_dim)
        self.norm2 = nn.LayerNorm(fusion_dim)
        self.norm3 = nn.LayerNorm(fusion_dim)
        self.norm4 = nn.LayerNorm(fusion_dim)
        
        # FFN
        self.ffn = nn.Sequential(
            nn.Linear(fusion_dim, fusion_dim * 4),
            nn.GELU(),
            nn.Linear(fusion_dim * 4, fusion_dim)
        )
        
    def forward(self, cam_feat: torch.Tensor, 
                radar_feat: torch.Tensor) -> torch.Tensor:
        """
        Args:
            cam_feat: (B, T_c, cam_dim) 摄像头时序特征
            radar_feat: (B, T_r, radar_dim) 雷达时序特征
        Returns:
            fused: (B, T_c, fusion_dim) 融合特征
        """
        # 投影到统一维度
        cam_q = self.cam_proj(cam_feat)  # (B, T_c, fusion_dim)
        radar_k = self.radar_proj(radar_feat)  # (B, T_r, fusion_dim)
        
        # 摄像头→雷达交叉注意力
        attn1, _ = self.cam_to_radar_attn(
            query=cam_q, key=radar_k, value=radar_k
        )
        cam_enhanced = self.norm1(cam_q + attn1)
        cam_enhanced = self.norm2(cam_enhanced + self.ffn(cam_enhanced))
        
        # 雷达→摄像头交叉注意力
        attn2, _ = self.radar_to_cam_attn(
            query=radar_k, key=cam_q, value=cam_q
        )
        radar_enhanced = self.norm3(radar_k + attn2)
        radar_enhanced = self.norm4(radar_enhanced + self.ffn(radar_enhanced))
        
        # 对齐到摄像头时间维度
        if radar_enhanced.shape[1] != cam_enhanced.shape[1]:
            radar_enhanced = F.interpolate(
                radar_enhanced.transpose(1, 2),
                size=cam_enhanced.shape[1],
                mode='linear', align_corners=True
            ).transpose(1, 2)
        
        # 元素级融合
        fused = (cam_enhanced + radar_enhanced) / 2
        return fused


class LateFusionModule(nn.Module):
    """
    晚期融合：各模态独立决策 + 置信度加权
    
    适用：最终决策层（警告等级/气囊部署）
    """
    
    def __init__(self, num_classes: int = 5):
        super().__init__()
        
        # 各模态置信度估计
        self.camera_confidence = ConfidenceEstimator()
        self.radar_confidence = ConfidenceEstimator()
        self.tof_confidence = ConfidenceEstimator()
        
        # 自适应权重学习
        self.weight_predictor = nn.Sequential(
            nn.Linear(6, 32),  # 3模态×2(置信度+值)
            nn.ReLU(),
            nn.Linear(32, 3),
            nn.Softmax(dim=-1)
        )
        
    def forward(self, cam_result: dict, radar_result: dict,
                tof_result: dict) -> dict:
        """
        Args:
            cam_result: {prediction, confidence}
            radar_result: {prediction, confidence}
            tof_result: {prediction, confidence}
        Returns:
            fused_result
        """
        # 自适应权重
        input_vec = torch.stack([
            cam_result['confidence'],
            radar_result['confidence'],
            tof_result['confidence'],
            cam_result['prediction'].max(),
            radar_result['prediction'].max(),
            tof_result['prediction'].max(),
        ], dim=-1)
        
        weights = self.weight_predictor(input_vec)  # (B, 3)
        
        # 加权融合
        fused_pred = (weights[:, 0:1] * cam_result['prediction'] +
                      weights[:, 1:2] * radar_result['prediction'] +
                      weights[:, 2:3] * tof_result['prediction'])
        
        return {
            'prediction': fused_pred,
            'weights': weights,
            'cam_confidence': cam_result['confidence'],
            'radar_confidence': radar_result['confidence'],
            'tof_confidence': tof_result['confidence'],
        }


class ConfidenceEstimator(nn.Module):
    """基于预测熵的置信度估计"""
    
    def __init__(self):
        super().__init__()
        self.proj = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(128, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        return self.proj(x).squeeze(-1)


# 测试
if __name__ == "__main__":
    print("座舱多模态传感器融合系统:")
    print("=" * 50)
    
    system = CabinSensorFusionSystem()
    for name, spec in system.SENSORS.items():
        print(f"\n{name}:")
        for k, v in spec.items():
            print(f"  {k}: {v}")
    
    # 测试中期融合
    fusion = MidLevelFusionModule()
    cam = torch.randn(2, 30, 128)  # 30帧摄像头特征
    radar = torch.randn(2, 20, 64)  # 20帧雷达特征
    
    with torch.no_grad():
        out = fusion(cam, radar)
    
    print(f"\n中期融合测试:")
    print(f"  摄像头输入: {cam.shape}")
    print(f"  雷达输入: {radar.shape}")
    print(f"  融合输出: {out.shape}")
    print(f"  参数量: {sum(p.numel() for p in fusion.parameters())/1e3:.1f}K")