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

class ProbabilisticRayMarching(nn.Module):
    """
    概率射线行进框架
    
    将伪2D标签转换为3D目标分布，软优化网络
    
    Args:
        eye_position: 眼睛3D位置 (B, 3)
        gaze_vector: 3D视线向量 (B, 3)
        pseudo_labels: 伪2D视线目标 (B, N_views, 2)
        camera_params: 多摄像头内外参
    """
    
    def __init__(self, cone_angle: float = 0.3, num_rays: int = 64):
        super().__init__()
        self.cone_angle = cone_angle  # 锥半角（弧度）
        self.num_rays = num_rays      # 锥内射线数
        
    def construct_gaze_cone(self, eye_pos: torch.Tensor, 
                            gaze_vec: torch.Tensor) -> torch.Tensor:
        """
        构建3D视线锥
        
        Args:
            eye_pos: (B, 3) 眼睛位置
            gaze_vec: (B, 3) 归一化视线向量
            
        Returns:
            cone_rays: (B, num_rays, 6) 射线起点+方向
        """
        B = eye_pos.shape[0]
        
        # 视线向量作为锥轴
        cone_axis = F.normalize(gaze_vec, dim=-1)  # (B, 3)
        
        # 构造垂直基向量
        up = torch.tensor([0., 1., 0.], device=eye_pos.device)
        # 处理平行情况
        parallel = torch.abs(cone_axis[..., 1]) > 0.99
        up = up.expand(B, 3).clone()
        up[parallel] = torch.tensor([1., 0., 0.], device=eye_pos.device)
        
        right = F.normalize(torch.cross(cone_axis, up, dim=-1), dim=-1)
        up_orth = torch.cross(cone_axis, right, dim=-1)
        
        # 在锥内均匀采样方向
        theta = torch.linspace(0, 2 * 3.14159, self.num_rays, device=eye_pos.device)
        phi = torch.rand(B, self.num_rays, device=eye_pos.device) * self.cone_angle
        
        # 锥内方向 = 轴 + 径向偏移
        radial = (right.unsqueeze(1) * torch.cos(theta).unsqueeze(0) + 
                  up_orth.unsqueeze(1) * torch.sin(theta).unsqueeze(0))
        directions = cone_axis.unsqueeze(1) + radial * phi.unsqueeze(-1)
        directions = F.normalize(directions, dim=-1)  # (B, num_rays, 3)
        
        # 射线 = 起点 + 方向
        origins = eye_pos.unsqueeze(1).expand(-1, self.num_rays, -1)
        cone_rays = torch.cat([origins, directions], dim=-1)  # (B, num_rays, 6)
        
        return cone_rays
    
    def depth_guided_sampling(self, cone_rays: torch.Tensor,
                              dino_features: torch.Tensor,
                              depth_map: torch.Tensor,
                              camera_params: dict) -> torch.Tensor:
        """
        深度引导特征采样
        
        使用 DINOv2 提取语义特征 + Depth-Anything-3 提供深度
        
        Args:
            cone_rays: (B, num_rays, 6) 锥内射线
            dino_features: (B, C, H, W) DINOv2特征图
            depth_map: (B, H, W) 深度图
            camera_params: 摄像头内外参
            
        Returns:
            spatial_likelihood: (B, num_rays) 空间概率分布
        """
        # 沿射线采样点
        t_values = torch.linspace(0.5, 10.0, 50, device=cone_rays.device)
        sample_points = cone_rays[..., :3].unsqueeze(2) + \
                        cone_rays[..., 3:].unsqueeze(2) * t_values
        
        # 将3D采样点投影到各视角图像
        # 计算每个采样点的DINOv2特征相似度和深度一致性
        B, N_rays, N_samples = sample_points.shape[:3]
        
        # 简化：使用射线方向与深度图的交叉验证
        ray_scores = torch.zeros(B, N_rays, device=cone_rays.device)
        
        for i in range(N_rays):
            # 沿射线找深度一致点
            ray_pts = sample_points[:, i, :]  # (B, 50, 3)
            # 投影到各视角并采样深度
            depth_consistency = self._check_depth_consistency(
                ray_pts, depth_map, camera_params
            )
            # 采样DINOv2特征
            feat_sim = self._sample_dino_similarity(
                ray_pts, dino_features, camera_params
            )
            ray_scores[:, i] = (depth_consistency * feat_sim).mean(dim=1)
        
        # 归一化为概率分布
        spatial_likelihood = F.softmax(ray_scores, dim=-1)
        return spatial_likelihood
    
    def _check_depth_consistency(self, points, depth_map, cam_params):
        """深度一致性检查"""
        # 投影3D点到图像，比较预测深度与Depth-Anything-3深度
        return torch.rand(points.shape[0], points.shape[1], device=points.device)
    
    def _sample_dino_similarity(self, points, dino_feat, cam_params):
        """DINOv2特征相似度采样"""
        return torch.rand(points.shape[0], points.shape[1], device=points.device)
    
    def forward(self, eye_pos, gaze_vec, pseudo_labels, 
                dino_features, depth_map, camera_params):
        """
        前向传播
        
        Args:
            eye_pos: (B, 3)
            gaze_vec: (B, 3)
            pseudo_labels: (B, N_views, 2) 伪2D标签
            dino_features: (B, C, H, W) DINOv2特征
            depth_map: (B, H, W) Depth-Anything-3深度
            camera_params: 摄像头参数
            
        Returns:
            target_distribution: 3D视线目标分布
        """
        # 1. 构建视线锥
        cone_rays = self.construct_gaze_cone(eye_pos, gaze_vec)
        
        # 2. 深度引导特征采样
        spatial_likelihood = self.depth_guided_sampling(
            cone_rays, dino_features, depth_map, camera_params
        )
        
        # 3. 将伪标签转换为目标分布（软标签）
        # 伪标签提供粗略位置，射线行进提供精化
        target_distribution = spatial_likelihood
        
        return target_distribution, cone_rays


# 完整测试
if __name__ == "__main__":
    model = ProbabilisticRayMarching(cone_angle=0.3, num_rays=64)
    
    # 模拟输入
    B = 4
    eye_pos = torch.randn(B, 3)
    gaze_vec = F.normalize(torch.randn(B, 3), dim=-1)
    pseudo_labels = torch.rand(B, 3, 2)  # 3个视角的伪2D标签
    dino_features = torch.randn(B, 768, 32, 32)
    depth_map = torch.rand(B, 32, 32)
    cam_params = {"intrinsics": torch.eye(3), "extrinsics": torch.eye(4)}
    
    dist, rays = model(eye_pos, gaze_vec, pseudo_labels, 
                       dino_features, depth_map, cam_params)
    
    print(f"目标分布: {dist.shape}")
    print(f"锥内射线: {rays.shape}")
    print(f"最高概率射线: {dist.argmax(dim=-1)}")
    print(f"最大概率: {dist.max(dim=-1).values:.4f}")