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

class TransGazeObject(nn.Module):
    """
    TransGaze-Object: 端到端注视目标预测
    
    架构:
    1. 面部分支: 面部特征 + 虹膜加权眼特征
    2. 场景分支: YOLOv8 检测对象 + 空间特征
    3. Transformer 交叉注意力: 计算面部-对象相似度
    4. 注视目标分类
    """
    
    def __init__(self, face_dim=512, eye_dim=256, 
                 object_dim=512, hidden_dim=256, num_heads=8):
        super().__init__()
        
        # 面部分支
        self.face_encoder = FaceEncoder(face_dim, hidden_dim)
        self.eye_encoder = IrisWeightedEyeEncoder(eye_dim, hidden_dim)
        
        # 场景对象分支
        self.object_encoder = ObjectEncoder(object_dim, hidden_dim)
        
        # 交叉注意力 Transformer
        self.cross_attn = nn.MultiheadAttention(
            hidden_dim, num_heads, batch_first=True
        )
        self.norm = nn.LayerNorm(hidden_dim)
        
        # 注视目标分类头
        self.gaze_classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, 1)  # 每个对象的注视概率
        )
        
    def forward(self, face_img, eye_img, scene_objects):
        """
        Args:
            face_img: (B, 3, 224, 224) 驾驶员面部
            eye_img: (B, 3, 112, 112) 眼部裁剪
            scene_objects: (B, N, 5) [x1,y1,x2,y2,class] 对象框
            
        Returns:
            gaze_scores: (B, N) 每个对象的注视概率
        """
        # 面部+眼部特征
        face_feat = self.face_encoder(face_img)     # (B, hidden_dim)
        eye_feat = self.eye_encoder(eye_img)          # (B, hidden_dim)
        driver_feat = face_feat + eye_feat            # (B, hidden_dim)
        
        # 场景对象特征
        obj_feat = self.object_encoder(scene_objects)  # (B, N, hidden_dim)
        
        # 交叉注意力: driver → objects
        driver_query = driver_feat.unsqueeze(1)  # (B, 1, hidden_dim)
        normed_q = self.norm(driver_query)
        normed_kv = self.norm(obj_feat)
        
        attn_out, attn_weights = self.cross_attn(
            normed_q, normed_kv, normed_kv
        )
        
        # 拼接 driver 和 attended objects
        driver_enhanced = attn_out.squeeze(1)  # (B, hidden_dim)
        
        # 对每个对象评分
        scores = []
        for i in range(scene_objects.shape[1]):
            pair = torch.cat([
                driver_enhanced,
                obj_feat[:, i, :]
            ], dim=-1)  # (B, hidden_dim*2)
            score = self.gaze_classifier(pair)  # (B, 1)
            scores.append(score)
        
        gaze_scores = torch.cat(scores, dim=-1)  # (B, N)
        gaze_scores = F.softmax(gaze_scores, dim=-1)
        
        return gaze_scores, attn_weights


class FaceEncoder(nn.Module):
    """面部特征编码器 (简化版 ResNet)"""
    
    def __init__(self, in_dim=3, out_dim=256):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_dim, 64, 7, stride=2, padding=3),
            nn.BatchNorm2d(64), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128), nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.BatchNorm2d(256), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, out_dim)
        )
        
    def forward(self, x):
        return self.encoder(x)


class IrisWeightedEyeEncoder(nn.Module):
    """虹膜加权眼部特征编码器"""
    
    def __init__(self, in_dim=3, out_dim=256):
        super().__init__()
        # 虹膜检测分支
        self.iris_detector = nn.Sequential(
            nn.Conv2d(in_dim, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 1, 1),
            nn.Sigmoid()  # 虹膜权重图
        )
        
        # 眼部特征编码
        self.eye_encoder = nn.Sequential(
            nn.Conv2d(in_dim, 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(1),
            nn.Flatten(),
            nn.Linear(128, out_dim)
        )
        
    def forward(self, x):
        # 虹膜权重
        iris_weight = self.iris_detector(x)  # (B, 1, H, W)
        
        # 加权眼部特征
        weighted_x = x * iris_weight.expand_as(x)
        
        # 编码
        feat = self.eye_encoder(weighted_x)
        return feat


class ObjectEncoder(nn.Module):
    """场景对象编码器"""
    
    def __init__(self, in_dim=5, out_dim=256):
        super().__init__()
        # [x1, y1, x2, y2, class] → 特征
        self.encoder = nn.Sequential(
            nn.Linear(in_dim, 128),
            nn.GELU(),
            nn.Linear(128, out_dim)
        )
        
    def forward(self, objects):
        """
        Args:
            objects: (B, N, 5) [x1, y1, x2, y2, class]
        """
        return self.encoder(objects)


if __name__ == "__main__":
    model = TransGazeObject()
    
    B, N = 2, 10  # 2 batch, 10 objects
    face = torch.randn(B, 3, 224, 224)
    eye = torch.randn(B, 3, 112, 112)
    objects = torch.rand(B, N, 5)  # 10 个对象
    
    scores, attn = model(face, eye, objects)
    print(f"Face: {face.shape}")
    print(f"Objects: {objects.shape}")
    print(f"Gaze scores: {scores.shape}")  # (2, 10)
    print(f"Attention: {attn.shape}")        # (2, 1, 10)
    print(f"Top object: {scores.argmax(dim=-1)}")
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")