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

class OmniGazeRewardModel(nn.Module):
    """
    OmniGaze 奖励模型
    
    NeurIPS 2025, arXiv:2510.13660
    代码：https://github.com/quhongyu/OmniGaze
    
    三通道置信度评分：
    1. 3D方向向量置信度（伪标签一致性）
    2. 视觉嵌入置信度（外观质量）
    3. 语义线索置信度（MLLM视线描述）
    """
    
    def __init__(self, visual_dim: int = 512, 
                 text_dim: int = 768,
                 gaze_dim: int = 3):
        super().__init__()
        
        # 通道1：3D方向向量置信度
        self.gaze_confidence = nn.Sequential(
            nn.Linear(gaze_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
        
        # 通道2：视觉嵌入置信度
        self.visual_confidence = nn.Sequential(
            nn.Linear(visual_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
        # 通道3：语义线索置信度
        self.semantic_confidence = nn.Sequential(
            nn.Linear(text_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
        # 融合权重学习
        self.fusion_weights = nn.Parameter(torch.ones(3) / 3)
        
    def forward(self, gaze_vec: torch.Tensor,
                visual_embed: torch.Tensor,
                semantic_embed: torch.Tensor) -> dict:
        """
        Args:
            gaze_vec: (B, 3) 伪标签3D方向向量
            visual_embed: (B, 512) 视觉编码器嵌入
            semantic_embed: (B, 768) MLLM语义嵌入
            
        Returns:
            result: dict with confidence score and individual scores
        """
        # 各通道置信度
        gaze_conf = self.gaze_confidence(gaze_vec)  # (B, 1)
        visual_conf = self.visual_confidence(visual_embed)  # (B, 1)
        semantic_conf = self.semantic_confidence(semantic_embed)  # (B, 1)
        
        # 加权融合
        weights = F.softmax(self.fusion_weights, dim=0)
        total_conf = (weights[0] * gaze_conf + 
                     weights[1] * visual_conf + 
                     weights[2] * semantic_conf)
        
        return {
            'confidence': total_conf.squeeze(-1),  # (B,)
            'gaze_conf': gaze_conf.squeeze(-1),
            'visual_conf': visual_conf.squeeze(-1),
            'semantic_conf': semantic_conf.squeeze(-1),
            'weights': weights.detach(),
        }


class OmniGazeTrainer:
    """
    OmniGaze 半监督训练管道
    
    Step 1: 教师模型在标注数据上训练
    Step 2: 教师模型在无标签数据上生成伪标签
    Step 3: 奖励模型评估伪标签质量
    Step 4: 加权损失训练学生模型
    """
    
    def __init__(self, teacher_model: nn.Module,
                 student_model: nn.Module,
                 reward_model: OmniGazeRewardModel):
        self.teacher = teacher_model
        self.student = student_model
        self.reward = reward_model
        
    def compute_loss(self, labeled_images, labeled_gaze,
                     unlabeled_images, 
                     visual_encoder, mllm_model):
        """
        计算OmniGaze训练损失
        
        Args:
            labeled_images: (B_l, 3, 224, 224) 标注图像
            labeled_gaze: (B_l, 3) 标注3D方向
            unlabeled_images: (B_u, 3, 224, 224) 无标签图像
            visual_encoder: 视觉特征提取器
            mllm_model: 多模态大语言模型
        """
        # === 标注数据损失 ===
        pred_gaze = self.student(labeled_images)
        supervised_loss = F.mse_loss(pred_gaze, labeled_gaze)
        
        # === 无标签数据处理 ===
        with torch.no_grad():
            # 1. 教师生成伪标签
            pseudo_gaze = self.teacher(unlabeled_images)
            
            # 2. 视觉嵌入
            visual_feats = visual_encoder(unlabeled_images)
            
            # 3. MLLM语义描述
            semantic_feats = mllm_model.encode_gaze_description(
                unlabeled_images
            )
            
            # 4. 奖励模型评分
            reward_out = self.reward(
                pseudo_gaze, visual_feats, semantic_feats
            )
        
        # 5. 置信度加权损失
        confidences = reward_out['confidence']  # (B_u,)
        
        # 高/低置信度分离
        high_mask = confidences > 0.5
        low_mask = ~high_mask
        
        # 高置信度样本加权训练
        if high_mask.sum() > 0:
            high_conf = confidences[high_mask]
            high_pseudo = pseudo_gaze[high_mask]
            high_images = unlabeled_images[high_mask]
            
            student_pred = self.student(high_images)
            weighted_loss = (high_conf * 
                           F.mse_loss(student_pred, high_pseudo, 
                                     reduction='none').mean(dim=1)).mean()
        else:
            weighted_loss = torch.tensor(0.0, device=labeled_images.device)
        
        # 低置信度样本一致性正则
        if low_mask.sum() > 0:
            low_images = unlabeled_images[low_mask]
            with torch.no_grad():
                aug_low = torch.clamp(low_images + 
                    torch.randn_like(low_images) * 0.1, 0, 1)
            pred_low = self.student(low_images)
            pred_aug = self.student(aug_low)
            consistency_loss = F.mse_loss(pred_low, pred_aug)
        else:
            consistency_loss = torch.tensor(0.0, device=labeled_images.device)
        
        # 总损失
        total_loss = (supervised_loss + 
                     0.5 * weighted_loss + 
                     0.3 * consistency_loss)
        
        return {
            'total_loss': total_loss,
            'supervised_loss': supervised_loss.item(),
            'weighted_loss': weighted_loss.item(),
            'consistency_loss': consistency_loss.item(),
            'high_conf_ratio': high_mask.float().mean().item(),
        }


# 测试
if __name__ == "__main__":
    # 简化测试
    reward_model = OmniGazeRewardModel()
    
    batch = 8
    gaze = torch.randn(batch, 3)
    visual = torch.randn(batch, 512)
    semantic = torch.randn(batch, 768)
    
    with torch.no_grad():
        out = reward_model(gaze, visual, semantic)
    
    print("OmniGaze 奖励模型测试:")
    print(f"  总置信度: {out['confidence']}")
    print(f"  视线置信度: {out['gaze_conf']}")
    print(f"  视觉置信度: {out['visual_conf']}")
    print(f"  语义置信度: {out['semantic_conf']}")
    print(f"  融合权重: {out['weights']}")
    print(f"  参数量: {sum(p.numel() for p in reward_model.parameters())/1e3:.1f}K")