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

class GazeDPTR(nn.Module):
    """
    GazeDPTR: Dual-stream Gaze Pyramid Transformer
    
    CVPR 2024论文核心方法实现
    双流处理: 原始图像 + 归一化图像
    金字塔特征: 多层级面部特征
    
    Args:
        feat_dim: 特征维度
        num_heads: 注意力头数
        num_layers: Transformer层数
    """
    
    def __init__(self, feat_dim=512, num_heads=8, num_layers=4):
        super().__init__()
        
        # 1. 原始图像流（捕获全局面部信息）
        self.raw_stream = PyramidBackbone(
            in_channels=3, 
            feat_dim=feat_dim,
            levels=3  # 金字塔3层
        )
        
        # 2. 归一化图像流（视角校正后，捕获细粒度眼部）
        self.normalized_stream = PyramidBackbone(
            in_channels=3,
            feat_dim=feat_dim,
            levels=3
        )
        
        # 3. 摄像头姿态引导（指导双流融合）
        self.camera_pose_embed = nn.Sequential(
            nn.Linear(6, feat_dim // 4),  # 6DoF摄像头姿态
            nn.GELU(),
            nn.Linear(feat_dim // 4, feat_dim)
        )
        
        # 4. 双流Transformer融合
        self.fusion_transformer = DualStreamFusion(
            d_model=feat_dim,
            num_heads=num_heads,
            num_layers=num_layers
        )
        
        # 5. 注视输出头
        self.gaze_head = nn.Sequential(
            nn.Linear(feat_dim, feat_dim // 2),
            nn.GELU(),
            nn.Linear(feat_dim // 2, 3)  # 3D注视向量
        )
    
    def forward(self, raw_img, normalized_img, camera_pose):
        """
        Args:
            raw_img: (B, 3, 224, 224) 原始驾驶员面部图像
            normalized_img: (B, 3, 224, 224) 视角归一化图像
            camera_pose: (B, 6) 6DoF摄像头姿态
        
        Returns:
            gaze_vector: (B, 3) 3D注视方向向量
        """
        # 1. 双流金字塔特征提取
        raw_feat = self.raw_stream(raw_img)  # (B, feat_dim)
        norm_feat = self.normalized_stream(normalized_img)  # (B, feat_dim)
        
        # 2. 摄像头姿态引导
        pose_embed = self.camera_pose_embed(camera_pose)  # (B, feat_dim)
        
        # 3. Transformer融合
        fused = self.fusion_transformer(
            raw_feat, norm_feat, pose_embed
        )  # (B, feat_dim)
        
        # 4. 注视向量预测
        gaze = self.gaze_head(fused)
        gaze = F.normalize(gaze, p=2, dim=-1)  # 单位向量
        
        return gaze


class PyramidBackbone(nn.Module):
    """金字塔特征提取器：多层级面部特征"""
    
    def __init__(self, in_channels=3, feat_dim=512, levels=3):
        super().__init__()
        self.levels = levels
        
        # 多尺度卷积特征
        self.convs = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(in_channels if i == 0 else 64*(2**i), 
                         64*(2**(i+1)), 3, 2, 1),
                nn.BatchNorm2d(64*(2**(i+1))),
                nn.ReLU()
            )
            for i in range(levels)
        ])
        
        # 各层级投影
        self.projections = nn.ModuleList([
            nn.Linear(64 * (2**(i+1)), feat_dim // levels)
            for i in range(levels)
        ])
        
        # 最终融合
        self.fusion = nn.Linear(feat_dim, feat_dim)
    
    def forward(self, x):
        features = []
        feat = x
        for i in range(self.levels):
            feat = self.convs[i](feat)
            # 全局平均池化
            pooled = F.adaptive_avg_pool2d(feat, 1).flatten(1)
            # 投影
            projected = self.projections[i](pooled)
            features.append(projected)
        
        # 拼接并融合
        combined = torch.cat(features, dim=-1)
        return self.fusion(combined)


class DualStreamFusion(nn.Module):
    """双流Transformer融合"""
    
    def __init__(self, d_model=512, num_heads=8, num_layers=4):
        super().__init__()
        
        # 交叉注意力：原始流query ↔ 归一化流key/value
        cross_attn_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=d_model * 4,
            dropout=0.1,
            activation='gelu',
            batch_first=True
        )
        self.cross_attn = nn.TransformerDecoder(cross_attn_layer, num_layers)
        
        # 摄像头姿态门控
        self.pose_gate = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.Sigmoid()
        )
    
    def forward(self, raw_feat, norm_feat, pose_embed):
        """
        双流融合 + 摄像头姿态引导
        """
        # 原始流作为query，归一化流作为key/value
        raw_query = raw_feat.unsqueeze(1)  # (B, 1, d_model)
        norm_kv = norm_feat.unsqueeze(1)  # (B, 1, d_model)
        
        # 交叉注意力
        fused = self.cross_attn(raw_query, norm_kv, norm_kv)
        fused = fused.squeeze(1)  # (B, d_model)
        
        # 摄像头姿态门控
        gate = self.pose_gate(pose_embed)
        fused = fused * gate + fused
        
        return fused


# 测试
if __name__ == "__main__":
    model = GazeDPTR(feat_dim=512, num_heads=8, num_layers=4)
    
    raw = torch.randn(4, 3, 224, 224)
    norm = torch.randn(4, 3, 224, 224)
    pose = torch.randn(4, 6)
    
    gaze = model(raw, norm, pose)
    print(f"Raw image: {raw.shape}")
    print(f"Normalized image: {norm.shape}")
    print(f"Camera pose: {pose.shape}")
    print(f"Gaze vector: {gaze.shape}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # 验证注视向量是单位向量
    norms = torch.norm(gaze, dim=-1)
    print(f"Gaze vector norms: {norms}")
    print("✅ 应接近1.0")