"""可视化工具 - 用于显示点云、平面、ROI等"""
import numpy as np
import open3d as o3d
from pathlib import Path
from typing import List, Tuple, Optional


class Visualizer:
    """点云和测量结果可视化工具"""

    @staticmethod
    def visualize_roi_in_context(
        full_cloud: np.ndarray,
        roi_cloud: np.ndarray,
        bbox_min: List[float],
        bbox_max: List[float],
        title: str = "ROI in Context",
        downsample_full: int = 10,
        save_path: Optional[str] = None
    ):
        """
        在完整点云上下文中可视化ROI区域

        Args:
            full_cloud: 完整点云 (N, 3)
            roi_cloud: ROI区域点云 (M, 3)
            bbox_min: 包围盒最小坐标 [x, y, z]
            bbox_max: 包围盒最大坐标 [x, y, z]
            title: 窗口标题
            downsample_full: 完整点云降采样倍数（每N个点取1个）
            save_path: 保存截图路径（可选）
        """
        # 降采样完整点云以提高显示性能
        full_cloud_downsampled = full_cloud[::downsample_full]

        # 创建完整点云对象（灰色）
        pcd_full = o3d.geometry.PointCloud()
        pcd_full.points = o3d.utility.Vector3dVector(full_cloud_downsampled)
        pcd_full.paint_uniform_color([0.7, 0.7, 0.7])  # 灰色

        # 创建ROI点云对象（红色）
        pcd_roi = o3d.geometry.PointCloud()
        pcd_roi.points = o3d.utility.Vector3dVector(roi_cloud)
        pcd_roi.paint_uniform_color([1.0, 0.0, 0.0])  # 红色

        # 创建包围盒线框
        bbox = Visualizer._create_bbox_lineset(bbox_min, bbox_max, color=[0, 1, 0])

        # 可视化
        geometries = [pcd_full, pcd_roi, bbox]
        o3d.visualization.draw_geometries(
            geometries,
            window_name=title,
            width=1200,
            height=800,
            point_show_normal=False
        )

        # 保存截图（如果指定）
        if save_path:
            vis = o3d.visualization.Visualizer()
            vis.create_window(window_name=title, width=1200, height=800, visible=False)
            for geom in geometries:
                vis.add_geometry(geom)
            vis.update_renderer()
            vis.capture_screen_image(save_path)
            vis.destroy_window()
            print(f"截图已保存: {save_path}")

    @staticmethod
    def visualize_plane_measurement(
        source_cloud: np.ndarray,
        target_cloud: np.ndarray,
        source_plane_normal: np.ndarray,
        source_plane_point: np.ndarray,
        target_plane_normal: np.ndarray,
        target_plane_point: np.ndarray,
        sample_points: np.ndarray,
        distances: np.ndarray,
        title: str = "Plane Measurement",
        save_path: Optional[str] = None
    ):
        """
        可视化平面到平面的距离测量

        Args:
            source_cloud: 源平面点云
            target_cloud: 目标平面点云
            source_plane_normal: 源平面法向量
            source_plane_point: 源平面上的点
            target_plane_normal: 目标平面法向量
            target_plane_point: 目标平面上的点
            sample_points: 采样点
            distances: 每个采样点到目标平面的距离
            title: 窗口标题
            save_path: 保存截图路径
        """
        # 创建源平面点云（蓝色）
        pcd_source = o3d.geometry.PointCloud()
        pcd_source.points = o3d.utility.Vector3dVector(source_cloud)
        pcd_source.paint_uniform_color([0.0, 0.5, 1.0])  # 蓝色

        # 创建目标平面点云（绿色）
        pcd_target = o3d.geometry.PointCloud()
        pcd_target.points = o3d.utility.Vector3dVector(target_cloud)
        pcd_target.paint_uniform_color([0.0, 1.0, 0.0])  # 绿色

        # 创建采样点（红色）
        pcd_samples = o3d.geometry.PointCloud()
        pcd_samples.points = o3d.utility.Vector3dVector(sample_points)
        pcd_samples.paint_uniform_color([1.0, 0.0, 0.0])  # 红色

        # 创建距离线（从采样点到目标平面的投影）
        lines = []
        for i, (point, dist) in enumerate(zip(sample_points, distances)):
            # 计算投影点
            projection = point - dist * target_plane_normal
            lines.append([point, projection])

        # 创建线集
        line_points = []
        line_indices = []
        for i, (start, end) in enumerate(lines):
            line_points.append(start)
            line_points.append(end)
            line_indices.append([2*i, 2*i+1])

        lineset = o3d.geometry.LineSet()
        lineset.points = o3d.utility.Vector3dVector(line_points)
        lineset.lines = o3d.utility.Vector2iVector(line_indices)
        lineset.paint_uniform_color([1.0, 0.5, 0.0])  # 橙色

        # 可视化
        geometries = [pcd_source, pcd_target, pcd_samples, lineset]
        o3d.visualization.draw_geometries(
            geometries,
            window_name=title,
            width=1200,
            height=800,
            point_show_normal=False
        )

        if save_path:
            vis = o3d.visualization.Visualizer()
            vis.create_window(window_name=title, width=1200, height=800, visible=False)
            for geom in geometries:
                vis.add_geometry(geom)
            vis.update_renderer()
            vis.capture_screen_image(save_path)
            vis.destroy_window()
            print(f"截图已保存: {save_path}")

    @staticmethod
    def _create_bbox_lineset(
        bbox_min: List[float],
        bbox_max: List[float],
        color: List[float] = [0, 1, 0]
    ) -> o3d.geometry.LineSet:
        """创建包围盒线框"""
        # 8个顶点
        points = [
            [bbox_min[0], bbox_min[1], bbox_min[2]],  # 0
            [bbox_max[0], bbox_min[1], bbox_min[2]],  # 1
            [bbox_max[0], bbox_max[1], bbox_min[2]],  # 2
            [bbox_min[0], bbox_max[1], bbox_min[2]],  # 3
            [bbox_min[0], bbox_min[1], bbox_max[2]],  # 4
            [bbox_max[0], bbox_min[1], bbox_max[2]],  # 5
            [bbox_max[0], bbox_max[1], bbox_max[2]],  # 6
            [bbox_min[0], bbox_max[1], bbox_max[2]],  # 7
        ]

        # 12条边
        lines = [
            [0, 1], [1, 2], [2, 3], [3, 0],  # 底面
            [4, 5], [5, 6], [6, 7], [7, 4],  # 顶面
            [0, 4], [1, 5], [2, 6], [3, 7],  # 竖边
        ]

        lineset = o3d.geometry.LineSet()
        lineset.points = o3d.utility.Vector3dVector(points)
        lineset.lines = o3d.utility.Vector2iVector(lines)
        lineset.paint_uniform_color(color)

        return lineset

    @staticmethod
    def save_roi_pointcloud(
        roi_cloud: np.ndarray,
        save_path: str
    ):
        """保存ROI点云为PLY文件"""
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(roi_cloud)
        o3d.io.write_point_cloud(save_path, pcd)
        print(f"ROI点云已保存: {save_path}")