#!/usr/bin/env python3
"""
三体问题求解器 - 快速示例
运行此文件以查看三体问题的基本模拟
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os

# 添加当前目录到路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from three_body_problem import ThreeBodySolver, Particle, ThreeBodyConfig, ThreeBodyVisualizer


def simple_example():
    """简单示例：自定义三体系统"""
    print("=" * 60)
    print("三体问题求解器 - 简单示例")
    print("=" * 60)
    
    # 创建三个质点
    print("\n1. 创建三个质点...")
    particles = [
        Particle(mass=1.0, position=[1.0, 0.0, 0.0], velocity=[0.0, 0.5, 0.0], name="Star A", color="red"),
        Particle(mass=1.0, position=[-1.0, 0.0, 0.0], velocity=[0.0, -0.5, 0.0], name="Star B", color="green"),
        Particle(mass=0.5, position=[0.0, 1.0, 0.0], velocity=[-0.3, 0.0, 0.0], name="Star C", color="blue")
    ]
    
    for i, p in enumerate(particles):
        print(f"   质点 {i+1}: {p.name}, 质量: {p.mass:.2f}, 位置: [{p.position[0]:.2f}, {p.position[1]:.2f}, {p.position[2]:.2f}]")
    
    # 创建求解器
    print("\n2. 创建求解器...")
    solver = ThreeBodySolver(particles, dt=0.001)
    
    # 模拟2年
    print("\n3. 模拟2年运动...")
    print("   开始积分...")
    solver.simulate(total_time=2.0, progress_interval=500)
    
    # 分析结果
    print("\n4. 分析结果...")
    trajectories = solver.get_trajectories()
    print(f"   模拟完成! 生成了 {len(trajectories[0])} 个轨迹点")
    
    com = solver.get_center_of_mass()
    print(f"   系统质心位置: [{com[0]:.4f}, {com[1]:.4f}, {com[2]:.4f}] AU")
    
    momentum_error, angular_momentum_error, energy_error = solver.get_conservation_errors()
    print(f"   守恒定律误差:")
    print(f"     - 动量误差: {momentum_error:.2e}")
    print(f"     - 角动量误差: {angular_momentum_error:.2e}")
    print(f"     - 能量相对误差: {energy_error:.2e}")
    
    # 可视化
    print("\n5. 生成可视化图形...")
    visualizer = ThreeBodyVisualizer(figsize=(12, 8))
    
    # 创建3D轨迹图
    visualizer.create_3d_plot()
    visualizer.plot_trajectories(solver, title="三体系统运动轨迹")
    
    # 保存图形
    output_file = "simple_example_3d.png"
    visualizer.save_figure(output_file, dpi=200)
    print(f"   3D轨迹图已保存到: {output_file}")
    
    # 创建2D投影图
    fig, ax = visualizer.plot_2d_projection(solver, projection='xy', title="XY平面投影")
    output_file_2d = "simple_example_2d.png"
    fig.savefig(output_file_2d, dpi=200, bbox_inches='tight')
    print(f"   2D投影图已保存到: {output_file_2d}")
    
    print("\n6. 显示图形...")
    plt.show()
    
    return solver


def figure8_example():
    """8字形轨道示例"""
    print("\n" + "=" * 60)
    print("8字形轨道示例")
    print("=" * 60)
    
    # 使用预置的8字形轨道配置
    particles = ThreeBodyConfig.create_figure8_config()
    
    # 打印配置信息
    ThreeBodyConfig.print_config_summary(particles)
    
    # 创建求解器
    solver = ThreeBodySolver(particles, dt=0.001)
    
    # 模拟5年
    print("\n模拟5年8字形轨道...")
    solver.simulate(total_time=5.0, progress_interval=1000)
    
    # 可视化
    visualizer = ThreeBodyVisualizer(figsize=(10, 8))
    visualizer.create_3d_plot()
    visualizer.plot_trajectories(solver, title="8字形三体轨道")
    
    output_file = "figure8_example.png"
    visualizer.save_figure(output_file, dpi=200)
    print(f"\n图形已保存到: {output_file}")
    
    plt.show()
    
    return solver


def quick_test():
    """快速测试所有模块"""
    print("=" * 60)
    print("快速测试所有模块")
    print("=" * 60)
    
    # 测试1: 创建质点
    print("\n1. 测试质点创建...")
    p1 = Particle(mass=1.0, position=[1, 0, 0], velocity=[0, 1, 0], name="Test Star")
    print(f"   创建质点: {p1.name}, 质量: {p1.mass}, 位置: {p1.position}")
    
    # 测试2: 创建配置
    print("\n2. 测试配置创建...")
    particles = ThreeBodyConfig.create_figure8_config()
    print(f"   创建了 {len(particles)} 个质点")
    
    # 测试3: 创建求解器
    print("\n3. 测试求解器创建...")
    solver = ThreeBodySolver(particles, dt=0.01)
    print(f"   求解器创建成功，时间步长: {solver.dt}")
    
    # 测试4: 单步积分
    print("\n4. 测试单步积分...")
    initial_positions = [p.position.copy() for p in particles]
    solver.step()
    for i, (old_pos, particle) in enumerate(zip(initial_positions, solver.particles)):
        moved = not np.allclose(old_pos, particle.position)
        print(f"   质点 {i+1} 位置变化: {'是' if moved else '否'}")
    
    # 测试5: 质心计算
    print("\n5. 测试质心计算...")
    com = solver.get_center_of_mass()
    print(f"   系统质心: [{com[0]:.4f}, {com[1]:.4f}, {com[2]:.4f}]")
    
    # 测试6: 能量计算
    print("\n6. 测试能量计算...")
    energy = solver._calculate_total_energy()
    print(f"   系统总能量: {energy:.6e}")
    
    print("\n所有测试通过!")
    return True


def main():
    """主函数"""
    print("三体问题求解器示例")
    print("=" * 60)
    print("选择要运行的示例:")
    print("  1. 简单示例 (自定义三体系统)")
    print("  2. 8字形轨道示例")
    print("  3. 快速测试所有模块")
    print("  4. 全部运行")
    
    try:
        choice = input("\n请输入选择 (1-4): ").strip()
        
        if choice == "1":
            simple_example()
        elif choice == "2":
            figure8_example()
        elif choice == "3":
            quick_test()
        elif choice == "4":
            quick_test()
            simple_example()
            figure8_example()
        else:
            print("无效选择，运行简单示例...")
            simple_example()
            
    except KeyboardInterrupt:
        print("\n\n用户中断")
    except Exception as e:
        print(f"\n错误: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n" + "=" * 60)
    print("示例运行完成!")
    print("=" * 60)


if __name__ == "__main__":
    main()