# demo_final_beautiful.py
# 最终美化版：无图例，专注焊缝模型本身的高质量渲染。
# 运行：python demo_final_beautiful.py

import numpy as np
import open3d as o3d
from typing import Dict


def get_final_ideal_ground_truth() -> Dict:
    """
    使用你提供的最终版理想坐标。
    """
    print("--- Using your final provided ideal coordinates. ---")
    final_3_seams = {
        'bottom_left': {
            'start_3d': np.array([-142.2 + 8.9, 0.0, 0.0]),
            'end_3d': np.array([-2.7, 0.0, 0.0])
        },
        'middle': {
            'start_3d': np.array([0.0, 0.0, 0.0]),
            'end_3d': np.array([0.0, -50.3, 0.0])
        },
        'top_left': {
            'start_3d': np.array([0.0, 0.0, 5.2]),
            'end_3d': np.array([0.0, 0.0, 142.5])
        }
    }
    return final_3_seams


def rotation_matrix_from_vectors(vec_from: np.ndarray, vec_to: np.ndarray) -> np.ndarray:
    """计算从 vec_from 到 vec_to 的旋转矩阵。"""
    a = vec_from / (np.linalg.norm(vec_from) + 1e-12)
    b = vec_to / (np.linalg.norm(vec_to) + 1e-12)
    v = np.cross(a, b)
    c = np.clip(np.dot(a, b), -1.0, 1.0)
    s = np.linalg.norm(v)
    if s < 1e-12:
        return np.eye(3) if c > 0.0 else -np.eye(3)
    kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]], dtype=float)
    return np.eye(3) + kmat + kmat @ kmat * ((1.0 - c) / (s ** 2))


def visualize_beautiful_model(seams: Dict):
    """
    【最终美化版】可视化：
    - 使用PBR材质，增加金属质感。
    - 增加灯光，增强立体感。
    - 设置一个好的初始视角。
    """
    print("\n--- Visualizing Final Target Model (High Quality Render) ---")
    colors = {'bottom_left': [0.8, 0.1, 0.1], 'middle': [0.1, 0.8, 0.1], 'top_left': [0.1, 0.1, 0.8]}
    geoms = []

    # 1. 绘制焊缝主体 (圆柱体)
    radius = 0.5  # 适中的粗细
    for name, data in seams.items():
        start, end = np.asarray(data['start_3d']), np.asarray(data['end_3d'])
        direction = end - start
        length = np.linalg.norm(direction)
        if length < 1e-6: continue

        cyl = o3d.geometry.TriangleMesh.create_cylinder(radius=radius, height=length, resolution=64)
        cyl.compute_vertex_normals()  # 法线对于光照计算至关重要
        R = rotation_matrix_from_vectors(np.array([0.0, 0.0, 1.0]), direction)
        cyl.rotate(R, center=(0, 0, 0)).translate((start + end) / 2.0)
        cyl.paint_uniform_color(colors[name])
        geoms.append(cyl)

    # 2. 移动坐标轴到不遮挡的位置
    all_points = [p for data in seams.values() for p in data.values()]
    bbox = o3d.geometry.AxisAlignedBoundingBox.create_from_points(o3d.utility.Vector3dVector(all_points))
    axis_size = max(30.0, np.linalg.norm(bbox.get_extent()) * 0.2)
    axis_origin = bbox.get_min_bound() - np.array([axis_size * 1.5, axis_size * 0.5, 0])
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=axis_size, origin=axis_origin)
    geoms.append(frame)

    # 3. 可视化
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="Final Target Model (High Quality)", width=1280, height=720)

    for g in geoms:
        vis.add_geometry(g)

    # --- 渲染和视角设置 ---
    opt = vis.get_render_option()
    opt.background_color = np.asarray([0.1, 0.1, 0.1])  # 深灰色背景，突出主体
    opt.mesh_show_back_face = False
    opt.light_on = True  # 确保灯光已开启

    # 获取视图控制器并设置相机位置
    view_ctl = vis.get_view_control()
    # 这个函数会自动计算一个能看全所有物体的合适视角
    vis.reset_view_point(True)
    # 你可以进一步手动调整相机参数以获得特定角度
    # view_ctl.set_zoom(0.8)
    # view_ctl.rotate(x=100, y=100) # 旋转视角

    print("Visualization ready. You can rotate the view. Press 'Q' to close.")
    vis.run()
    vis.destroy_window()


if __name__ == '__main__':
    seams = get_final_ideal_ground_truth()

    print("--- Final Ideal Seam Coordinates ---")
    for n, p in seams.items():
        print(f"{n}: start={np.around(p['start_3d'], 1)}, end={np.around(p['end_3d'], 1)}")

    visualize_beautiful_model(seams)