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

使用方法:
blender --background --python light_path_simple.py

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

import bpy
import json
import os
import numpy as np

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

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)):
    """创建光路材质"""
    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 = 1.0
    
    # 创建输出节点
    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_paths(light_paths):
    """创建所有光路"""
    objects = []
    
    # 颜色方案
    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),  # 紫色
    ]
    
    for i, path in enumerate(light_paths[len(light_paths):]):
        if len(path) < 2:
            continue
            
        # 选择颜色
        color = colors[i % len(colors)]
        
        # 创建材质
        material = create_light_path_material(color)
        
        # 创建光路对象
        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
    
    # 设置世界背景
    world = bpy.context.scene.world
    if world is None:
        world = bpy.data.worlds.new("World")
        bpy.context.scene.world = 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 save_files():
    """保存文件"""
    try:
        # 保存Blender文件
        blend_path = os.path.abspath("./light_path_simple.blend")
        bpy.ops.wm.save_as_mainfile(filepath=blend_path)
        print(f"Blender文件已保存为: {blend_path}")
        
        # 导出GLB文件
        glb_path = os.path.abspath("./light_path_simple.glb")
        bpy.ops.export_scene.gltf(
            filepath=glb_path, 
            export_materials='EXPORT'
        )
        print(f"GLB文件已导出为: {glb_path}")
        
        # 渲染图像
        scene = bpy.context.scene
        scene.render.resolution_x = 1280
        scene.render.resolution_y = 720
        scene.render.image_settings.file_format = 'PNG'
        
        render_path = os.path.abspath("./light_path_simple_render.png")
        scene.render.filepath = render_path
        bpy.ops.render.render(write_still=True)
        print(f"渲染图像已保存为: {render_path}")
        
        return True
    except Exception as e:
        print(f"保存文件时出错: {e}")
        return False

def main():
    """主函数"""
    print("🎨 光学系统光路可视化 - 简化版")
    print("=" * 40)
    
    # 加载光路数据
    light_paths = load_light_paths("miao_light_path_tsingtao.json")
    
    if not light_paths:
        print("❌ 无法加载光路数据")
        return
    
    print(f"处理前10条光路...")
    
    # 创建光路对象
    light_path_objects = create_light_paths(light_paths)
    
    # 设置场景
    setup_scene()
    
    # 保存文件
    if save_files():
        print("\n🎉 任务完成！")
        print(f"成功创建了 {len(light_path_objects)} 条光路")
    else:
        print("\n❌ 保存文件失败")

if __name__ == "__main__":
    main()