import torch
import torch.nn as nn

class GazeToTextModel(nn.Module):
    """
    视线→自然语言描述生成模型
    
    范式：从gaze轨迹生成描述性文本
    输入：gaze轨迹序列 + 视觉上下文
    输出：自然语言描述
    """
    
    def __init__(self, gaze_dim=4, d_model=256, vocab_size=10000, 
                 max_seq_len=50, n_heads=8, n_layers=4):
        super().__init__()
        
        # Gaze轨迹编码
        self.gaze_encoder = nn.Sequential(
            nn.Linear(gaze_dim, d_model),
            nn.LayerNorm(d_model),
            nn.GELU(),
            nn.Dropout(0.1)
        )
        
        # 位置编码
        self.pos_encoding = nn.Parameter(
            torch.randn(1, max_seq_len, d_model) * 0.02
        )
        
        # 视觉上下文编码（来自场景理解）
        self.visual_context_encoder = nn.Linear(512, d_model)  # 假设512维场景特征
        
        # Transformer解码器（生成文本）
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=0.1,
            batch_first=True
        )
        self.decoder = nn.TransformerDecoder(
            decoder_layer, num_layers=n_layers
        )
        
        # 语言模型头
        self.lm_head = nn.Linear(d_model, vocab_size)
        
    def forward(self, gaze_seq, visual_context, target_tokens=None,
                teacher_forcing_ratio=0.8):
        """
        Args:
            gaze_seq: (B, T_gaze, 4) gaze轨迹 [x, y, timestamp, velocity]
            visual_context: (B, 512) 场景理解特征
            target_tokens: (B, T_text) 目标文本token (训练时)
            teacher_forcing_ratio: teacher forcing比例
        
        Returns:
            logits: (B, T_text, vocab_size)
        """
        B = gaze_seq.size(0)
        
        # 编码gaze
        gaze_emb = self.gaze_encoder(gaze_seq)  # (B, T_gaze, d)
        gaze_emb = gaze_emb + self.pos_encoding[:, :gaze_seq.size(1)]
        
        # 编码视觉上下文
        ctx_emb = self.visual_context_encoder(visual_context)  # (B, d)
        ctx_emb = ctx_emb.unsqueeze(1)  # (B, 1, d)
        
        # 记忆 = gaze编码 + 视觉上下文
        memory = torch.cat([ctx_emb, gaze_emb], dim=1)  # (B, 1+T_gaze, d)
        
        # 文本生成
        if target_tokens is not None:
            # 训练模式：teacher forcing
            T_text = target_tokens.size(1)
            tgt_emb = nn.functional.embedding(
                target_tokens, 
                torch.randn(self.lm_head.out_features, 
                           self.lm_head.in_features, 
                           device=target_tokens.device)
            )
            tgt_emb = tgt_emb + self.pos_encoding[:, :T_text]
            
            # 因果掩码
            causal_mask = torch.triu(
                torch.ones(T_text, T_text, device=target_tokens.device) 
                * float('-inf'), diagonal=1
            )
            
            decoded = self.decoder(
                tgt_emb, memory,
                tgt_mask=causal_mask
            )
            logits = self.lm_head(decoded)
            return logits
        else:
            # 推理模式：自回归生成
            generated = torch.ones(B, 1, dtype=torch.long, 
                                   device=gaze_seq.device)  # <BOS>
            
            for _ in range(50):  # 最多生成50个token
                tgt_emb = nn.functional.embedding(
                    generated,
                    torch.randn(self.lm_head.out_features,
                               self.lm_head.in_features,
                               device=generated.device)
                )
                tgt_emb = tgt_emb + self.pos_encoding[:, :generated.size(1)]
                
                decoded = self.decoder(tgt_emb, memory)
                next_logits = self.lm_head(decoded[:, -1])
                next_token = next_logits.argmax(dim=-1, keepdim=True)
                generated = torch.cat([generated, next_token], dim=1)
                
                if (next_token == 2).all():  # <EOS>
                    break
            
            return generated


# === DMS 集成方案 ===
class GazeDescriptionDMS:
    """将gaze-to-text集成到DMS的方案"""
    
    # 描述模板映射（用于DMS决策）
    DESCRIPTION_PATTERNS = {
        'phone_check': r'phone|mobile|screen.*hand',
        'fatigue': r'drowsy|tired|heavy.*eyes|slow.*blink',
        'distraction': r'away.*road|off.*road|looking.*side',
        'mirror_check': r'mirror|rear|behind',
        'dashboard': r'speed|odometer|fuel|dashboard',
        'navigation': r'nav|route|direction|GPS',
        'passenger_interaction': r'passenger|side.*seat',
    }
    
    def gaze_to_dms_state(self, description):
        """将自然语言描述转为DMS状态"""
        import re
        states = []
        for state, pattern in self.DESCRIPTION_PATTERNS.items():
            if re.search(pattern, description, re.IGNORECASE):
                states.append(state)
        
        if not states:
            return 'normal'
        if 'phone_check' in states:
            return 'phone_use_warning'
        if 'fatigue' in states:
            return 'fatigue_warning'
        if 'distraction' in states:
            return 'distraction_warning'
        if len(states) >= 2:
            return 'multi_task_warning'
        return states[0]


# 测试
if __name__ == "__main__":
    model = GazeToTextModel(gaze_dim=4, d_model=128, vocab_size=5000)
    
    B = 2
    T_gaze = 30  # 1秒@30fps
    gaze = torch.randn(B, T_gaze, 4)
    ctx = torch.randn(B, 512)
    
    # 模拟文本目标
    target = torch.randint(0, 5000, (B, 20))
    
    logits = model(gaze, ctx, target_tokens=target)
    print(f"Gaze输入: {gaze.shape}")
    print(f"视觉上下文: {ctx.shape}")
    print(f"文本logits: {logits.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters())/1e6:.2f}M")
    
    # DMS集成示例
    dms = GazeDescriptionDMS()
    test_descs = [
        "Driver looks at phone screen, possibly reading message",
        "Driver eyes heavy, slow blinks detected",
        "Driver checks rear mirror frequently",
        "Driver gazes at navigation screen",
    ]
    print("\n=== Gaze描述→DMS状态 ===")
    for desc in test_descs:
        state = dms.gaze_to_dms_state(desc)
        print(f"  '{desc}' → {state}")