#!/usr/bin/env python3
"""
光路可视化Blender脚本
读取miao_light_path_tsingtao.json文件并使用Blender绘制3D光路模型

使用方法:
1. 在Blender中运行此脚本
2. 或使用命令行: blender --background --python light_path_blender.py

输出文件:
- light_path_model.blend - Blender工程文件
- light_path_model.glb - Web友好格式
- light_path_render.png - 渲染图像
"""

import bpy
import json
import mathutils
import os
import bmesh
import numpy as np
from mathutils import Vector

# 清空场景
bpy.ops.wm.read_factory_settings(use_empty=True)

# 设置渲染引擎为Cycles（更好的材质渲染）
bpy.context.scene.render.engine = 'CYCLES'

def load_light_paths(filename):
    """加载光路数据"""
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            data = json.load(f)
        print(f"成功加载光路文件: {filename}")
        print(f"包含 {len(data)} 条光路")
        return data
    except Exception as e:
        print(f"加载光路文件失败: {e}")
        return []

def create_light_path_material(color=(1.0, 0.2, 0.2, 1.0), emission_strength=0.5):
    """创建光路材质"""
    mat = bpy.data.materials.new(name="light_path")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    
    # 清除默认节点
    nodes.clear()
    
    # 创建发光节点
    emission = nodes.new(type='ShaderNodeEmission')
    emission.inputs['Color'].default_value = color
    emission.inputs['Strength'].default_value = emission_strength
    
    # 创建输出节点
    output = nodes.new(type='ShaderNodeOutputMaterial')
    
    # 连接节点
    links.new(emission.outputs['Emission'], output.inputs['Surface'])
    
    return mat

def create_light_path_curve(points, name="light_path"):
    """创建光路曲线"""
    # 创建曲线数据
    curve_data = bpy.data.curves.new(name, type='CURVE')
    curve_data.dimensions = '3D'
    curve_data.resolution_u = 12
    curve_data.bevel_depth = 0.02  # 光路粗细
    curve_data.bevel_resolution = 4
    
    # 创建样条线
    spline = curve_data.splines.new('POLY')
    spline.points.add(len(points) - 1)
    
    # 设置点坐标
    for i, point in enumerate(points):
        spline.points[i].co = (point[0], point[1], point[2], 1.0)
    
    # 创建曲线对象
    curve_obj = bpy.data.objects.new(name, curve_data)
    bpy.context.collection.objects.link(curve_obj)
    
    return curve_obj

def create_light_path_mesh(points, name="light_path"):
    """创建光路网格（圆柱体连接）"""
    # 创建网格数据
    mesh_data = bpy.data.meshes.new(name)
    mesh_obj = bpy.data.objects.new(name, mesh_data)
    bpy.context.collection.objects.link(mesh_obj)
    
    # 创建bmesh
    bm = bmesh.new()
    
    # 为每段光路创建圆柱体
    for i in range(len(points) - 1):
        start_point = Vector(points[i])
        end_point = Vector(points[i + 1])
        
        # 计算圆柱体参数
        direction = (end_point - start_point).normalized()
        length = (end_point - start_point).length
        radius = 0.02  # 光路半径
        
        # 创建圆柱体
        bmesh.ops.create_cone(
            bm,
            segments=8,
            diameter1=radius,
            diameter2=radius,
            depth=length
        )
        
        # 获取刚创建的圆柱体
        cylinder_verts = [v for v in bm.verts if v.select]
        cylinder_faces = [f for f in bm.faces if f.select]
        
        # 取消选择
        for v in cylinder_verts:
            v.select = False
        for f in cylinder_faces:
            f.select = False
        
        # 计算旋转
        up = Vector((0, 0, 1))
        if abs(direction.dot(up)) > 0.99:
            up = Vector((0, 1, 0))
        
        rotation = direction.rotation_difference(up)
        
        # 应用变换
        center = (start_point + end_point) / 2
        for v in cylinder_verts:
            v.co = rotation @ v.co + center
    
    # 转换为网格
    bm.to_mesh(mesh_data)
    bm.free()
    
    return mesh_obj

def create_light_path_objects(light_paths):
    """创建所有光路对象"""
    objects = []
    
    # 创建材质
    base_colors = [
        (1.0, 0.2, 0.2, 1.0),  # 红色
        (0.2, 1.0, 0.2, 1.0),  # 绿色
        (0.2, 0.2, 1.0, 1.0),  # 蓝色
        (1.0, 1.0, 0.2, 1.0),  # 黄色
        (1.0, 0.2, 1.0, 1.0),  # 紫色
        (0.2, 1.0, 1.0, 1.0),  # 青色
        (1.0, 0.5, 0.2, 1.0),  # 橙色
        (0.5, 0.2, 1.0, 1.0),  # 紫罗兰
    ]
    
    for i, path in enumerate(light_paths):
        if len(path) < 2:
            continue
            
        # 选择颜色
        color = base_colors[i % len(base_colors)]
        
        # 创建材质
        material = create_light_path_material(color, emission_strength=0.3)
        
        # 创建光路对象（使用曲线）
        obj = create_light_path_curve(path, f"light_path_{i}")
        
        # 应用材质
        obj.data.materials.append(material)
        
        objects.append(obj)
        
        print(f"创建光路 {i+1}: {len(path)} 个点")
    
    return objects

def setup_scene():
    """设置场景（相机、灯光等）"""
    # 添加相机
    bpy.ops.object.camera_add(location=(10, -10, 8))
    camera = bpy.context.active_object
    camera.rotation_euler = (1.0, 0, 0.785)
    bpy.context.scene.camera = camera
    
    # 添加主灯光
    bpy.ops.object.light_add(type='SUN', location=(5, 5, 10))
    sun = bpy.context.active_object
    sun.data.energy = 2.0
    
    # 添加环境光
    bpy.ops.object.light_add(type='AREA', location=(-5, -5, 8))
    area = bpy.context.active_object
    area.data.energy = 30.0
    area.data.size = 8.0
    
    # 设置世界背景为深色
    world = bpy.context.scene.world
    world.use_nodes = True
    bg_node = world.node_tree.nodes['Background']
    bg_node.inputs['Color'].default_value = (0.05, 0.05, 0.1, 1.0)
    bg_node.inputs['Strength'].default_value = 0.5

def calculate_bounds(light_paths):
    """计算光路的边界框"""
    if not light_paths:
        return None
    
    all_points = []
    for path in light_paths:
        all_points.extend(path)
    
    all_points = np.array(all_points)
    
    min_coords = all_points.min(axis=0)
    max_coords = all_points.max(axis=0)
    
    return min_coords, max_coords

def adjust_camera_to_bounds(min_coords, max_coords):
    """根据边界调整相机位置"""
    center = (min_coords + max_coords) / 2
    size = max_coords - min_coords
    max_size = max(size)
    
    # 计算相机距离
    distance = max_size * 2.5
    
    # 设置相机位置
    camera = bpy.context.scene.camera
    camera.location = center + Vector((distance, -distance, distance * 0.8))
    
    # 让相机看向中心
    direction = center - camera.location
    rot_quat = direction.to_track_quat('-Z', 'Y')
    camera.rotation_euler = rot_quat.to_euler()

def export_model(objects):
    """导出模型"""
    if not objects:
        print("警告：没有创建任何对象，跳过导出")
        return
    
    # 选择所有光路对象
    bpy.ops.object.select_all(action='DESELECT')
    for obj in objects:
        if obj and obj.name in bpy.data.objects:
            obj.select_set(True)
    
    try:
        # 导出为GLB格式
        glb_path = os.path.abspath("./light_path_model.glb")
        bpy.ops.export_scene.gltf(
            filepath=glb_path, 
            use_selection=True, 
            export_materials='EXPORT',
            export_lights=True
        )
        print(f"模型已导出为: {glb_path}")
    except Exception as e:
        print(f"GLB导出失败: {e}")
    
    try:
        # 保存Blender文件
        blend_path = os.path.abspath("./light_path_model.blend")
        bpy.ops.wm.save_as_mainfile(filepath=blend_path)
        print(f"Blender文件已保存为: {blend_path}")
    except Exception as e:
        print(f"Blender文件保存失败: {e}")

def render_image():
    """渲染图像"""
    try:
        # 设置渲染参数
        scene = bpy.context.scene
        scene.render.resolution_x = 1920
        scene.render.resolution_y = 1080
        scene.render.image_settings.file_format = 'PNG'
        scene.render.film_transparent = False
        
        render_path = os.path.abspath("./light_path_render.png")
        scene.render.filepath = render_path
        print(f"开始渲染到: {render_path}")
        bpy.ops.render.render(write_still=True)
        print(f"渲染图像已保存为: {render_path}")
    except Exception as e:
        print(f"渲染失败: {e}")
        import traceback
        traceback.print_exc()

def main():
    """主函数"""
    print("🎨 光学系统光路可视化")
    print("=" * 40)
    
    # 加载光路数据
    light_paths = load_light_paths("miao_light_path_tsingtao.json")
    
    if not light_paths:
        print("❌ 无法加载光路数据")
        return
    
    # 分析数据
    print(f"\n=== 光路数据分析 ===")
    print(f"总光路数量: {len(light_paths)}")
    
    # 统计每条光路的点数
    path_lengths = [len(path) for path in light_paths]
    print(f"光路点数统计:")
    print(f"  最少点数: {min(path_lengths)}")
    print(f"  最多点数: {max(path_lengths)}")
    print(f"  平均点数: {np.mean(path_lengths):.2f}")
    
    # 计算边界
    bounds = calculate_bounds(light_paths)
    if bounds:
        min_coords, max_coords = bounds
        print(f"\n坐标范围:")
        print(f"  X: [{min_coords[0]:.3f}, {max_coords[0]:.3f}]")
        print(f"  Y: [{min_coords[1]:.3f}, {max_coords[1]:.3f}]")
        print(f"  Z: [{min_coords[2]:.3f}, {max_coords[2]:.3f}]")
    
    # 创建光路对象
    print(f"\n正在创建光路对象...")
    light_path_objects = create_light_path_objects(light_paths)
    
    # 设置场景
    print("正在设置场景...")
    setup_scene()
    
    # 调整相机
    if bounds:
        adjust_camera_to_bounds(bounds[0], bounds[1])
    
    # 导出模型
    print("正在导出模型...")
    export_model(light_path_objects)
    
    # 渲染图像
    print("正在渲染图像...")
    render_image()
    
    print("\n🎉 任务完成！")
    print("\n生成的文件:")
    print("- light_path_model.blend: Blender工程文件")
    print("- light_path_model.glb: Web友好格式")
    print("- light_path_render.png: 渲染图像")
    print(f"\n成功创建了 {len(light_path_objects)} 条光路")

if __name__ == "__main__":
    main()
