"""
SRCDet稀疏查询融合框架

核心：基于DETR风格的查询机制
将摄像头和雷达特征通过学习的查询点进行稀疏采样
"""
import torch
import torch.nn as nn

class SparseQueryFusion(nn.Module):
    """
    稀疏查询融合模块
    
    原理：
    1. 生成N个可学习查询点（3D空间中的参考点）
    2. 在BEV和透视图空间分别采样摄像头+雷达特征
    3. 通过交叉注意力融合多模态特征
    """
    
    def __init__(self, n_queries=300, feat_dim=256, n_heads=8):
        super().__init__()
        self.n_queries = n_queries
        self.feat_dim = feat_dim
        
        # 可学习查询嵌入
        self.query_embeds = nn.Embedding(n_queries, feat_dim)
        
        # 3D参考点（从查询嵌入预测）
        self.ref_point_proj = nn.Linear(feat_dim, 3)  # (x, y, z)
        
        # BEV空间交叉注意力
        self.bev_cross_attn = nn.MultiheadAttention(
            embed_dim=feat_dim, num_heads=n_heads, batch_first=True
        )
        
        # 透视图空间交叉注意力
        self.persp_cross_attn = nn.MultiheadAttention(
            embed_dim=feat_dim, num_heads=n_heads, batch_first=True
        )
        
        # 跨模态注意力（摄像头←→雷达）
        self.camera_radar_attn = nn.MultiheadAttention(
            embed_dim=feat_dim, num_heads=n_heads, batch_first=True
        )
        
        # FFN
        self.ffn = nn.Sequential(
            nn.Linear(feat_dim, feat_dim * 4),
            nn.ReLU(),
            nn.Linear(feat_dim * 4, feat_dim),
        )
        
        # 输出头
        self.cls_head = nn.Linear(feat_dim, 1)  # 二分类
        self.box_head = nn.Linear(feat_dim, 7)   # (x,y,z,w,h,l,theta)
    
    def forward(self, camera_bev, radar_bev, camera_persp):
        """
        Args:
            camera_bev: (B, H_bev, W_bev, C) 摄像头BEV特征
            radar_bev: (B, H_bev, W_bev, C) 雷达BEV特征
            camera_persp: (B, H_p, W_p, C) 摄像头透视图特征
            
        Returns:
            boxes: (B, N, 7) 3D边界框
            scores: (B, N, 1) 置信度
        """
        B = camera_bev.shape[0]
        
        # 查询嵌入 → 3D参考点
        query = self.query_embeds.weight.unsqueeze(0).expand(B, -1, -1)
        ref_points = self.ref_point_proj(query)  # (B, N, 3)
        
        # BEV空间采样（双模态）
        bev_feat = camera_bev.flatten(1, 2)  # (B, H*W, C)
        radar_feat = radar_bev.flatten(1, 2)
        
        # 跨模态注意力
        radar_enhanced, _ = self.camera_radar_attn(
            query=radar_feat, key=bev_feat, value=bev_feat
        )
        
        # BEV交叉注意力
        fused_bev, _ = self.bev_cross_attn(
            query=query, key=radar_enhanced, value=radar_enhanced
        )
        
        # 透视图交叉注意力
        persp_feat = camera_persp.flatten(1, 2)
        fused_persp, _ = self.persp_cross_attn(
            query=query, key=persp_feat, value=persp_feat
        )
        
        # 双视角融合
        fused = fused_bev + fused_persp + self.ffn(fused_bev + fused_persp)
        
        # 输出
        scores = self.cls_head(fused)
        boxes = self.box_head(fused)
        
        return boxes, scores, ref_points

"""
座舱内DMS摄像头+CPD雷达稀疏融合

传统方案：摄像头和雷达各自处理后拼接
SRCDet启发：基于查询的稀疏融合，在BEV和透视图双空间交互
"""

class CabinSparseFusion(nn.Module):
    """
    座舱内多模态稀疏融合
    
    输入：
    - DMS摄像头图像（RGB/IR）
    - CPD雷达点云（UWB/60GHz）
    
    输出：
    - 乘员位置/分类
    - 姿态估计
    - 呼吸检测增强
    """
    
    def __init__(self, n_queries=50):
        super().__init__()
        self.n_queries = n_queries
        
        # 查询嵌入（座舱3D空间中的采样点）
        self.query_embeds = nn.Embedding(n_queries, 256)
        
        # 座舱坐标系（方向盘为原点）
        # x: 左右(-0.5m ~ 0.5m)
        # y: 前后(0 ~ 2m)
        # z: 高低(0 ~ 1.5m)
        self.ref_point_proj = nn.Linear(256, 3)
        
        # 摄像头特征→BEV投影
        self.camera_bev_proj = nn.Sequential(
            nn.Conv2d(512, 256, 1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((32, 32))
        )
        
        # 雷达点云→BEV栅格化
        self.radar_bev_proj = nn.Sequential(
            nn.Conv2d(1, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 256, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((32, 32))
        )
        
        # 跨模态注意力
        self.cross_modal_attn = nn.MultiheadAttention(
            embed_dim=256, num_heads=8, batch_first=True
        )
        
        # 输出头
        self.occupant_head = nn.Linear(256, 4)  # empty/adult/child/pet
        self.pose_head = nn.Linear(256, 13*2)   # 13关键点×2D
        self.breathing_head = nn.Linear(256, 1)  # 呼吸频率
    
    def forward(self, camera_feat, radar_points):
        """
        Args:
            camera_feat: (B, 512, H, W) DMS图像特征
            radar_points: (B, N_radar, 4) 雷达点(x,y,z,intensity)
        """
        # 投影到BEV
        cam_bev = self.camera_bev_proj(camera_feat)  # (B, 256, 32, 32)
        
        # 雷达栅格化到BEV
        radar_bev = self._points_to_bev(radar_points)  # (B, 1, 32, 32)
        radar_bev = self.radar_bev_proj(radar_bev)  # (B, 256, 32, 32)
        
        # 稀疏查询融合
        query = self.query_embeds.weight.unsqueeze(0)
        ref_pts = self.ref_point_proj(query)  # 座舱3D采样点
        
        cam_flat = cam_bev.flatten(2).permute(0, 2, 1)  # (B, 1024, 256)
        radar_flat = radar_bev.flatten(2).permute(0, 2, 1)
        
        # 跨模态：雷达用摄像头增强
        radar_enhanced, _ = self.cross_modal_attn(
            query=radar_flat, key=cam_flat, value=cam_flat
        )
        
        # 查询采样
        fused, _ = self.cross_modal_attn(
            query=query, key=radar_enhanced, value=radar_enhanced
        )
        
        # 多任务输出
        occupant = self.occupant_head(fused)
        pose = self.pose_head(fused)
        breathing = self.breathing_head(fused)
        
        return {
            'occupant_class': occupant,
            'pose_keypoints': pose.view(-1, 13, 2),
            'breathing_rate': breathing,
            'ref_points': ref_pts,
        }

IMS_FUSION_ROADMAP = {
    '2026_Q4': {
        'task': 'DMS摄像头+座椅ECG融合原型',
        'method': '信号质量门控+降级切换',
        'target': '疲劳检测准确率>90%',
    },
    '2027_Q1': {
        'task': 'DMS+CPD雷达稀疏融合',
        'method': 'SRCDet启发，BEV双空间',
        'target': 'CPD+OOP联合检测',
    },
    '2027_Q2': {
        'task': '全模态融合',
        'method': '摄像头+雷达+ECG+语音',
        'target': '驾驶员全状态评估',
    },
}