import torch
import numpy as np
from gsplat import rasterization

class Incremental3DGS:
    def __init__(self, device='cuda'):
        self.device = device
        # 存储全局高斯点：位置、协方差、颜色、不透明度
        self.means = torch.empty((0, 3), device=device)
        self.covs = torch.empty((0, 3, 3), device=device)
        self.colors = torch.empty((0, 3), device=device)
        self.opacities = torch.empty((0, 1), device=device)
        
    def add_frame(self, rgb, depth, K, T):
        """
        rgb: (H, W, 3) RGB 图像
        depth: (H, W) 深度图
        K: (3, 3) 相机内参
        T: (4, 4) 相机位姿（世界坐标系到相机坐标系）
        """
        H, W = depth.shape
        
        # 1. 反投影到 3D（相机坐标系）
        u, v = np.meshgrid(np.arange(W), np.arange(H))
        points_cam = np.stack([
            (u - K[0, 2]) * depth / K[0, 0],
            (v - K[1, 2]) * depth / K[1, 1],
            depth
        ], axis=-1)  # (H, W, 3)
        
        # 2. 转到世界坐标系
        T_inv = np.linalg.inv(T)
        points_world = points_cam @ T_inv[:3, :3].T + T_inv[:3, 3]
        
        # 3. 下采样（每 4 个像素取一个）
        mask = (depth > 0) & (depth < 10.0)  # 有效深度范围
        mask[::4, ::4] = False  # 稀疏采样
        
        new_means = torch.from_numpy(points_world[mask]).float().to(self.device)
        new_colors = torch.from_numpy(rgb[mask]).float().to(self.device) / 255.0
        
        # 4. 初始化协方差（球形高斯，半径 = 0.01m）
        N = new_means.shape[0]
        new_covs = torch.eye(3).unsqueeze(0).repeat(N, 1, 1).to(self.device) * 0.01**2
        new_opacities = torch.ones((N, 1), device=self.device) * 0.9
        
        # 5. 去重（简化版：距离阈值）
        if self.means.shape[0] > 0:
            dist = torch.cdist(new_means, self.means)  # (N_new, N_old)
            keep_mask = (dist.min(dim=1)[0] > 0.02)  # 保留距离 > 2cm 的点
            new_means = new_means[keep_mask]
            new_colors = new_colors[keep_mask]
            new_covs = new_covs[keep_mask]
            new_opacities = new_opacities[keep_mask]
        
        # 6. 合并到全局
        self.means = torch.cat([self.means, new_means], dim=0)
        self.colors = torch.cat([self.colors, new_colors], dim=0)
        self.covs = torch.cat([self.covs, new_covs], dim=0)
        self.opacities = torch.cat([self.opacities, new_opacities], dim=0)
        
        # 7. 限制总数（内存管理）
        if self.means.shape[0] > 500000:
            self.prune_far_points()
    
    def prune_far_points(self):
        """保留不透明度高的点"""
        keep_idx = torch.argsort(self.opacities.squeeze(), descending=True)[:500000]
        self.means = self.means[keep_idx]
        self.colors = self.colors[keep_idx]
        self.covs = self.covs[keep_idx]
        self.opacities = self.opacities[keep_idx]

def render_novel_view(gs_model, K, T, H=480, W=640):
    """
    从任意位姿 T 渲染场景
    返回: (H, W, 3) RGB 图像
    """
    # 准备相机参数
    viewmat = torch.from_numpy(T).float().cuda()  # (4, 4)
    
    # gsplat 渲染（简化调用）
    rendered_rgb, _, _ = rasterization(
        means=gs_model.means,
        quats=cov_to_quat(gs_model.covs),  # 协方差 → 四元数
        scales=torch.ones_like(gs_model.means) * 0.01,
        opacities=gs_model.opacities,
        colors=gs_model.colors,
        viewmats=viewmat.unsqueeze(0),
        Ks=torch.from_numpy(K).float().cuda().unsqueeze(0),
        width=W,
        height=H,
    )
    
    return rendered_rgb[0].cpu().numpy()  # (H, W, 3)

def cov_to_quat(covs):
    """协方差矩阵 → 四元数（简化版：假设各向同性）"""
    # ... (省略特征值分解代码)
    return torch.tensor([1, 0, 0, 0]).repeat(covs.shape[0], 1).cuda()

import cv2

def detect_frontiers(occupancy_map, robot_pos):
    """
    occupancy_map: (H, W) 0=未知, 1=空闲, 2=占据
    robot_pos: (x, y) 机器人当前位置（网格坐标）
    返回: 候选前沿方向列表 [(angle_1, score_1), ...]
    """
    # 1. 找到未知-空闲边界
    kernel = np.ones((5, 5), np.uint8)
    frontier = cv2.morphologyEx(
        (occupancy_map == 0).astype(np.uint8),
        cv2.MORPH_GRADIENT,
        kernel
    ) & (occupancy_map == 1)
    
    # 2. 聚类前沿点
    frontier_points = np.argwhere(frontier > 0)
    if len(frontier_points) == 0:
        return []
    
    # 3. 计算每个前沿方向的得分
    candidates = []
    for angle in np.linspace(0, 2*np.pi, 8, endpoint=False):
        direction = np.array([np.cos(angle), np.sin(angle)])
        # 统计该方向上的前沿点数
        scores = np.dot(frontier_points - robot_pos, direction)
        score = (scores > 0).sum()
        candidates.append((angle, score))
    
    # 4. 返回得分最高的 3 个方向
    candidates.sort(key=lambda x: x[1], reverse=True)
    return candidates[:3]

def generate_candidate_poses(robot_pose, frontier_angles, distance=1.0):
    """
    生成候选相机位姿
    robot_pose: (4, 4) 当前位姿
    frontier_angles: [(angle, score), ...]
    返回: [T1, T2, T3] 候选位姿列表
    """
    poses = []
    for angle, _ in frontier_angles:
        # 平移 distance 米
        T = robot_pose.copy()
        T[0, 3] += distance * np.cos(angle)
        T[1, 3] += distance * np.sin(angle)
        # 旋转朝向该方向
        T[:3, :3] = rotation_matrix_from_angle(angle)
        poses.append(T)
    return poses

def rotation_matrix_from_angle(angle):
    """2D 角度 → 3D 旋转矩阵（绕 Z 轴）"""
    c, s = np.cos(angle), np.sin(angle)
    return np.array([
        [c, -s, 0],
        [s,  c, 0],
        [0,  0, 1]
    ])

import base64
from io import BytesIO
from PIL import Image

def query_vlm(images, prompt, api_key):
    """
    调用 VLM API（以 OpenAI GPT-4V 为例）
    images: 图像列表
    prompt: 文本提示
    返回: VLM 的文本回复
    """
    import openai
    openai.api_key = api_key
    
    # 转换图像为 base64
    image_urls = []
    for img in images:
        buffered = BytesIO()
        Image.fromarray((img * 255).astype(np.uint8)).save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        image_urls.append(f"data:image/png;base64,{img_str}")
    
    # 构造消息
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                *[{"type": "image_url", "image_url": {"url": url}} for url in image_urls]
            ]
        }
    ]
    
    response = openai.ChatCompletion.create(
        model="gpt-4-vision-preview",
        messages=messages,
        max_tokens=300
    )
    
    return response.choices[0].message.content

import torch
import numpy as np
from gsplat import rasterization

class Incremental3DGS:
    def __init__(self, device='cuda'):
        self.device = device
        # 存储全局高斯点：位置、协方差、颜色、不透明度
        self.means = torch.empty((0, 3), device=device)
        self.covs = torch.empty((0, 3, 3), device=device)
        self.colors = torch.empty((0, 3), device=device)
        self.opacities = torch.empty((0, 1), device=device)
        
    def add_frame(self, rgb, depth, K, T):
        """增量添加新帧"""
        # 1. 反投影到 3D（相机坐标系）
        # ... (像素网格生成代码省略)
        points_cam = np.stack([...], axis=-1)
        
        # 2. 转到世界坐标系
        T_inv = np.linalg.inv(T)
        points_world = points_cam @ T_inv[:3, :3].T + T_inv[:3, 3]
        
        # 3. 下采样 + 颜色提取
        # ... (掩码过滤代码省略)
        new_means = torch.from_numpy(points_world[mask]).float().to(self.device)
        new_colors = torch.from_numpy(rgb[mask]).float().to(self.device) / 255.0
        
        # 4. 初始化协方差（球形高斯）
        N = new_means.shape[0]
        new_covs = torch.eye(3).unsqueeze(0).repeat(N, 1, 1).to(self.device) * 0.01**2
        new_opacities = torch.ones((N, 1), device=self.device) * 0.9
        
        # 5. 去重（距离阈值）
        if self.means.shape[0] > 0:
            dist = torch.cdist(new_means, self.means)
            keep_mask = (dist.min(dim=1)[0] > 0.02)
            new_means, new_colors, new_covs, new_opacities = \
                new_means[keep_mask], new_colors[keep_mask], new_covs[keep_mask], new_opacities[keep_mask]
        
        # 6. 合并到全局
        self.means = torch.cat([self.means, new_means], dim=0)
        self.colors = torch.cat([self.colors, new_colors], dim=0)
        self.covs = torch.cat([self.covs, new_covs], dim=0)
        self.opacities = torch.cat([self.opacities, new_opacities], dim=0)
        
        # 7. 内存管理（限制总数）
        if self.means.shape[0] > 500000:
            self.prune_far_points()
    
    def prune_far_points(self):
        """保留高不透明度点"""
        keep_idx = torch.argsort(self.opacities.squeeze(), descending=True)[:500000]
        self.means = self.means[keep_idx]
        # ... (其他属性同步裁剪省略)

import open3d as o3d

def visualize_3dgs(gs_model):
    """可视化当前 3DGS 点云"""
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(gs_model.means.cpu().numpy())
    pcd.colors = o3d.utility.Vector3dVector(gs_model.colors.cpu().numpy())
    
    # 添加坐标系
    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5)
    
    o3d.visualization.draw_geometries([pcd, coordinate_frame])

def detect_frontiers(occupancy_map, robot_pos):
    """检测未知-空闲边界，返回候选前沿方向"""
    # 1. 形态学梯度找边界
    kernel = np.ones((5, 5), np.uint8)
    frontier = cv2.morphologyEx(
        (occupancy_map == 0).astype(np.uint8),
        cv2.MORPH_GRADIENT, kernel
    ) & (occupancy_map == 1)
    
    frontier_points = np.argwhere(frontier > 0)
    if len(frontier_points) == 0:
        return []
    
    # 2. 8方向评分
    candidates = []
    for angle in np.linspace(0, 2*np.pi, 8, endpoint=False):
        direction = np.array([np.cos(angle), np.sin(angle)])
        scores = np.dot(frontier_points - robot_pos, direction)
        score = (scores > 0).sum()
        candidates.append((angle, score))
    
    return sorted(candidates, key=lambda x: x[1], reverse=True)[:3]

def generate_candidate_poses(robot_pose, frontier_angles, distance=1.0):
    """根据前沿角度生成候选位姿"""
    poses = []
    for angle, _ in frontier_angles:
        T = robot_pose.copy()
        T[:2, 3] += distance * np.array([np.cos(angle), np.sin(angle)])
        T[:3, :3] = np.array([[np.cos(angle), -np.sin(angle), 0],
                               [np.sin(angle),  np.cos(angle), 0],
                               [0, 0, 1]])
        poses.append(T)
    return poses

def query_vlm(images, prompt, api_key):
    """调用 VLM API"""
    import openai
    openai.api_key = api_key
    
    # 转换图像为 base64
    image_urls = []
    for img in images:
        # ... (图像编码代码省略)
        image_urls.append(f"data:image/png;base64,{img_str}")
    
    # 构造消息
    messages = [{
        "role": "user",
        "content": [
            {"type": "text", "text": prompt},
            *[{"type": "image_url", "image_url": {"url": url}} for url in image_urls]
        ]
    }]
    
    response = openai.ChatCompletion.create(model="gpt-4-vision-preview", messages=messages)
    return response.choices[0].message.content

def navigate(target_object, gs_model, robot, vlm_api_key):
    """基于VLM和3DGS的主动导航"""
    max_steps = 100
    
    for step in range(max_steps):
        # 1. 获取观察并更新3DGS
        rgb, depth, K, T = robot.get_observation()
        gs_model.add_frame(rgb, depth, K, T)
        
        # 2. 检测前沿
        frontiers = detect_frontiers(robot.get_occupancy_map(), robot.get_position())
        if len(frontiers) == 0:
            break
        
        # 3. 渲染候选视角
        candidate_poses = generate_candidate_poses(T, frontiers)
        candidate_views = [render_novel_view(gs_model, K, pose) for pose in candidate_poses]
        
        # 4. VLM推理
        prompt = f"任务：找到 {target_object}\n当前视角+候选视角A/B/C，选择最可能看到目标的视角"
        response = query_vlm([rgb] + candidate_views, prompt, vlm_api_key)
        
        # 5. 执行决策
        if "Found" in response:
            return True
        elif "A" in response:
            robot.move_to(candidate_poses[0])
        # ... (B/C分支类似)
    
    return False