#!/usr/bin/env python3
"""
测试三体问题求解器导入和基本功能
"""

import sys
import os

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

print("测试三体问题求解器导入...")
print("=" * 60)

try:
    # 测试导入
    from three_body_problem import Particle, ThreeBodySolver, ThreeBodyConfig, ThreeBodyVisualizer
    print("✓ 成功导入核心模块")
    
    # 测试创建质点
    p = Particle(mass=1.0, position=[1, 0, 0], velocity=[0, 1, 0], name='Test Star')
    print(f"✓ 成功创建质点: {p.name}, 质量: {p.mass}, 位置: {p.position}")
    
    # 测试能量计算
    energy = p.get_energy()
    print(f"✓ 质点动能: {energy:.6f}")
    
    # 测试创建配置
    particles = ThreeBodyConfig.create_figure8_config()
    print(f"✓ 成功创建8字形轨道配置: {len(particles)}个质点")
    for i, particle in enumerate(particles):
        print(f"  质点{i+1}: {particle.name}, 质量: {particle.mass:.3f}")
    
    # 测试创建求解器
    solver = ThreeBodySolver(particles, dt=0.001)
    print(f"✓ 成功创建求解器，时间步长: {solver.dt}")
    
    # 测试单步积分
    initial_positions = [particle.position.copy() for particle in particles]
    new_particles = solver.step()
    print(f"✓ 单步积分完成，时间: {solver.time:.4f}年")
    
    # 检查位置是否变化
    for i, (old_pos, new_particle) in enumerate(zip(initial_positions, new_particles)):
        moved = not all(abs(old_pos[j] - new_particle.position[j]) < 1e-10 for j in range(3))
        print(f"  质点{i+1} 位置变化: {'是' if moved else '否'}")
    
    # 测试质心计算
    com = solver.get_center_of_mass()
    print(f"✓ 系统质心: [{com[0]:.6f}, {com[1]:.6f}, {com[2]:.6f}]")
    
    # 测试能量计算
    energy = solver._calculate_total_energy()
    print(f"✓ 系统总能量: {energy:.6e}")
    
    # 测试守恒误差计算
    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("\n测试配置管理...")
    random_particles = ThreeBodyConfig.create_random_config()
    print(f"✓ 成功创建随机配置: {len(random_particles)}个质点")
    
    lagrange_particles = ThreeBodyConfig.create_lagrange_point_config(lagrange_point=4)
    print(f"✓ 成功创建拉格朗日点L4配置: {len(lagrange_particles)}个质点")
    
    # 测试可视化器创建
    visualizer = ThreeBodyVisualizer()
    print("✓ 成功创建可视化器")
    
    print("\n" + "=" * 60)
    print("所有测试通过! 三体问题求解器工作正常。")
    print("=" * 60)
    
    # 显示使用示例
    print("\n使用示例:")
    print("1. 运行简单示例: python three_body_problem/run_example.py")
    print("2. 运行8字形轨道: python three_body_problem/examples/figure8.py")
    print("3. 运行拉格朗日点示例: python three_body_problem/examples/lagrange.py")
    print("4. 运行随机系统示例: python three_body_problem/examples/random.py")
    print("5. 运行测试: python three_body_problem/tests/test_solver.py")
    
except ImportError as e:
    print(f"✗ 导入失败: {e}")
    print("请确保在项目根目录下运行此测试")
    
except Exception as e:
    print(f"✗ 测试失败: {e}")
    import traceback
    traceback.print_exc()