import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 简谐振子系统参数
k = 1.0        # 劲度系数
m = 1.0        # 质量
omega = np.sqrt(k/m)  # 角频率

def potential_energy(x):
    """势能: (1/2)*k*x^2"""
    return 0.5 * k * x**2

def kinetic_energy(v):
    """动能: (1/2)*m*v^2"""
    return 0.5 * m * v**2

def total_energy(x, v):
    """总能量"""
    return potential_energy(x) + kinetic_energy(v)

def acceleration(x):
    """加速度: a = -(k/m)*x"""
    return -omega**2 * x

# ============ 四种数值方法 ============

def euler_explicit(x0, v0, dt, n_steps):
    """欧拉显式 (前向欧拉)"""
    x = np.zeros(n_steps)
    v = np.zeros(n_steps)
    t = np.zeros(n_steps)
    x[0] = x0
    v[0] = v0
    t[0] = 0
    for i in range(n_steps - 1):
        x[i+1] = x[i] + v[i] * dt
        v[i+1] = v[i] + acceleration(x[i]) * dt
        t[i+1] = t[i] + dt
    return t, x, v


def euler_implicit(x0, v0, dt, n_steps):
    """欧拉隐式 (后向欧拉)"""
    x = np.zeros(n_steps)
    v = np.zeros(n_steps)
    t = np.zeros(n_steps)
    x[0] = x0
    v[0] = v0
    t[0] = 0
    for i in range(n_steps - 1):
        # 解析解：对简谐振子联立方程可直接得出
        x[i+1] = (x[i] + v[i] * dt) / (1 + omega**2 * dt**2)
        v[i+1] = v[i] + acceleration(x[i+1]) * dt
        t[i+1] = t[i] + dt
    return t, x, v


def midpoint_method(x0, v0, dt, n_steps):
    """中点差分法 (二阶 Runge-Kutta / 改进欧拉)"""
    x = np.zeros(n_steps)
    v = np.zeros(n_steps)
    t = np.zeros(n_steps)
    x[0] = x0
    v[0] = v0
    t[0] = 0
    for i in range(n_steps - 1):
        x_mid = x[i] + v[i] * dt / 2
        v_mid = v[i] + acceleration(x[i]) * dt / 2
        x[i+1] = x[i] + v_mid * dt
        v[i+1] = v[i] + acceleration(x_mid) * dt
        t[i+1] = t[i] + dt
    return t, x, v


def leapfrog_method(x0, v0, dt, n_steps):
    """蛙跳法 (Leapfrog / Störmer-Verlet)
    
    核心思想：位置和速度交错半步更新
        v_{n+1/2} = v_{n-1/2} + a(x_n) * dt    <- 速度在半整数时刻
        x_{n+1}   = x_n       + v_{n+1/2} * dt  <- 位置在整数时刻
    
    启动：用欧拉半步得到 v_{1/2}，之后进入交错循环。
    输出时将速度还原到整数时刻：v_n = (v_{n-1/2} + v_{n+1/2}) / 2
    """
    x = np.zeros(n_steps)
    v = np.zeros(n_steps)      # 整数时刻速度（仅用于输出/能量计算）
    t = np.zeros(n_steps)

    x[0] = x0
    v[0] = v0
    t[0] = 0

    # 用欧拉半步初始化 v_{1/2}
    v_half = v0 + acceleration(x0) * dt / 2

    for i in range(n_steps - 1):
        # 更新位置（整数时刻）
        x[i+1] = x[i] + v_half * dt
        t[i+1] = t[i] + dt

        # 更新半步速度（用新位置的加速度）
        v_half_new = v_half + acceleration(x[i+1]) * dt

        # 整数时刻速度（用于能量计算和输出）：取两个半步的平均
        v[i+1] = (v_half + v_half_new) / 2

        v_half = v_half_new

    return t, x, v


def exact_solution(x0, v0, t):
    """精确解"""
    A = np.sqrt(x0**2 + (v0 / omega)**2)
    phi = np.arctan2(-v0 / omega, x0)
    x_exact = A * np.cos(omega * t + phi)
    v_exact = -A * omega * np.sin(omega * t + phi)
    return x_exact, v_exact


# ============ 参数设置 ============
dt = 0.05           # 时间步长
T_total = 10.0      # 总时间
n_steps = int(T_total / dt) + 1

# 三个初始条件 (x0, v0)
initial_conditions = [
    (1.0, 0.0),      # 从最大位移释放
    (0.0, 1.0),      # 从平衡位置给初速度
    (0.8, 0.6)       # 混合初始条件
]

# 四种方法：(名称, 函数, 线型颜色, alpha)
methods = [
    ('欧拉显式',   euler_explicit,  'r-',  0.8),
    ('欧拉隐式',   euler_implicit,  'b-',  0.8),
    ('中点差分法', midpoint_method, 'g-',  0.8),
    ('蛙跳法',     leapfrog_method, 'm-',  0.8),
]

colors_ic = ['#e63946', '#457b9d', '#2d6a4f']   # 三种初始条件的颜色
labels_ic = ['从最大位移释放', '平衡位置初速度', '混合条件']

# 四种方法对应的绘图颜色（用于位移/速度时序图）
method_colors = ['r', 'b', 'g', 'm']

print(f"时间步长 dt = {dt}s, 步数 = {n_steps}")

# ============ 计算结果 ============
results = {}
for method_name, method_func, color_style, alpha in methods:
    results[method_name] = {}
    for idx, (x0, v0) in enumerate(initial_conditions):
        t_arr, x_arr, v_arr = method_func(x0, v0, dt, n_steps)
        E = total_energy(x_arr, v_arr)
        results[method_name][idx] = {
            't': t_arr, 'x': x_arr, 'v': v_arr, 'E': E,
            'E_initial': E[0], 'E_final': E[-1],
            'E_drift': (E[-1] - E[0]) / E[0] * 100
        }

# 精确解（高密度时间轴）
exact_t = np.linspace(0, T_total, 1000)

# ============ 图形布局 ============
fig = plt.figure(figsize=(18, 12))
fig.suptitle('简谐振子数值方法四方比较', fontsize=16, fontweight='bold')

gs = fig.add_gridspec(3, 4, hspace=0.40, wspace=0.30)
ax_phase    = fig.add_subplot(gs[0, :])     # 第1行全宽：相空间
ax_energy   = fig.add_subplot(gs[1, :])     # 第2行全宽：能量演化
ax_traj     = fig.add_subplot(gs[2, 0])     # 第3行：位置-时间
ax_velocity = fig.add_subplot(gs[2, 1])     # 第3行：速度-时间
ax_err_bar  = fig.add_subplot(gs[2, 2])     # 第3行：能量误差柱状图
ax_animation= fig.add_subplot(gs[2, 3])     # 第3行：相空间动画

# ============ 1. 相空间轨迹 ============
ax_phase.set_title('相空间轨迹比较 (x vs v)', fontsize=13)
ax_phase.set_xlabel('位置 x', fontsize=11)
ax_phase.set_ylabel('速度 v', fontsize=11)
ax_phase.grid(True, alpha=0.3)
ax_phase.axhline(y=0, color='k', linestyle='-', alpha=0.3)
ax_phase.axvline(x=0, color='k', linestyle='-', alpha=0.3)

line_styles_ic = ['-', '--', ':']
method_line_colors = ['r', 'b', 'g', 'm']

for midx, (x0, v0) in enumerate(initial_conditions):
    x_exact, v_exact = exact_solution(x0, v0, exact_t)
    ax_phase.plot(x_exact, v_exact, 'k-', alpha=0.25, linewidth=1,
                  label='精确解' if midx == 0 else "")
    for mi, (method_name, _, _, _) in enumerate(methods):
        res = results[method_name][midx]
        ax_phase.plot(res['x'], res['v'],
                      linestyle=['-', '--', ':', '-.'][midx],
                      color=method_line_colors[mi], alpha=0.75, linewidth=1.4,
                      label=method_name if midx == 0 else "")

handles, labels_leg = ax_phase.get_legend_handles_labels()
by_label = dict(zip(labels_leg, handles))
ax_phase.legend(by_label.values(), by_label.keys(),
                loc='upper right', fontsize=8, ncol=2)

# ============ 2. 能量演化（归一化） ============
ax_energy.set_title('归一化能量随时间演化  E(t)/E(0)', fontsize=13)
ax_energy.set_xlabel('时间 t (s)', fontsize=11)
ax_energy.set_ylabel('E / E0', fontsize=11)
ax_energy.grid(True, alpha=0.3)

energy_linestyles = ['-', '--', ':', '-.']
for mi, (method_name, _, _, _) in enumerate(methods):
    res = results[method_name][0]   # 只用第一个初始条件展示能量
    E_norm = res['E'] / res['E_initial']
    ax_energy.plot(res['t'], E_norm,
                   linestyle=energy_linestyles[mi],
                   color=method_line_colors[mi],
                   linewidth=1.8, alpha=0.9, label=method_name)

ax_energy.axhline(y=1.0, color='k', linestyle='-', alpha=0.5, linewidth=0.8)
ax_energy.set_ylim(0.85, 1.70)
ax_energy.legend(loc='upper left', fontsize=9)

# ============ 3. 位置-时间 (第一初始条件) ============
idx_show = 0
ax_traj.set_title(f'位置-时间  ({labels_ic[idx_show]})', fontsize=11)
ax_traj.set_xlabel('时间 t (s)', fontsize=10)
ax_traj.set_ylabel('位置 x', fontsize=10)
ax_traj.grid(True, alpha=0.3)

x_ex, v_ex = exact_solution(*initial_conditions[idx_show], exact_t)
ax_traj.plot(exact_t, x_ex, 'k-', alpha=0.5, linewidth=1.2, label='精确解')
for mi, (method_name, _, _, _) in enumerate(methods):
    res = results[method_name][idx_show]
    ax_traj.plot(res['t'], res['x'],
                 color=method_line_colors[mi], linewidth=1.4,
                 alpha=0.85, label=method_name)
ax_traj.legend(loc='upper right', fontsize=7)

# ============ 4. 速度-时间 (第一初始条件) ============
ax_velocity.set_title(f'速度-时间  ({labels_ic[idx_show]})', fontsize=11)
ax_velocity.set_xlabel('时间 t (s)', fontsize=10)
ax_velocity.set_ylabel('速度 v', fontsize=10)
ax_velocity.grid(True, alpha=0.3)

ax_velocity.plot(exact_t, v_ex, 'k-', alpha=0.5, linewidth=1.2, label='精确解')
for mi, (method_name, _, _, _) in enumerate(methods):
    res = results[method_name][idx_show]
    ax_velocity.plot(res['t'], res['v'],
                     color=method_line_colors[mi], linewidth=1.4,
                     alpha=0.85, label=method_name)
ax_velocity.legend(loc='upper right', fontsize=7)

# ============ 5. 能量误差柱状图 ============
ax_err_bar.set_title('能量漂移百分比  |dE/E0|*100%', fontsize=11)
ax_err_bar.set_ylabel('漂移 (%)', fontsize=10)
ax_err_bar.grid(True, alpha=0.3, axis='y')

method_names_short = ['显式', '隐式', '中点', '蛙跳']
bar_colors = method_line_colors
drifts = [abs(results[mn][0]['E_drift']) for mn, *_ in methods]
bars = ax_err_bar.bar(method_names_short, drifts, color=bar_colors, alpha=0.8, width=0.5)
for bar, val in zip(bars, drifts):
    ax_err_bar.text(bar.get_x() + bar.get_width() / 2,
                    bar.get_height() + 0.3,
                    f'{val:.2f}%', ha='center', va='bottom', fontsize=9)
ax_err_bar.set_ylim(0, max(drifts) * 1.2 + 2)

# ============ 6. 动画：四种方法相空间对比 (同一初始条件) ============
ax_animation.set_xlim(-1.5, 1.5)
ax_animation.set_ylim(-1.5, 1.5)
ax_animation.set_title('四种方法同步释放 (相空间)', fontsize=11)
ax_animation.set_xlabel('位置 x', fontsize=10)
ax_animation.set_ylabel('速度 v', fontsize=10)
ax_animation.grid(True, alpha=0.3)
ax_animation.axhline(y=0, color='k', linestyle='-', alpha=0.3)
ax_animation.axvline(x=0, color='k', linestyle='-', alpha=0.3)

anim_points = []
anim_trails = []
anim_idx = 0  # 选第一组初始条件，比较四种方法的同步演化
for mi, (method_name, _, _, _) in enumerate(methods):
    pt, = ax_animation.plot([], [], 'o', color=method_line_colors[mi],
                             markersize=7, label=method_name)
    tr, = ax_animation.plot([], [], '-', color=method_line_colors[mi], alpha=0.35, linewidth=1)
    anim_points.append(pt)
    anim_trails.append(tr)

x_anim_exact, v_anim_exact = exact_solution(*initial_conditions[anim_idx], exact_t)
ax_animation.plot(x_anim_exact, v_anim_exact, 'k-', alpha=0.2, linewidth=0.8, label='精确解')

ax_animation.legend(loc='upper right', fontsize=8)

def animate(frame):
    for mi, (method_name, _, _, _) in enumerate(methods):
        res = results[method_name][anim_idx]
        if frame < len(res['x']):
            anim_points[mi].set_data([res['x'][frame]], [res['v'][frame]])
            start = max(0, frame - 50)
            anim_trails[mi].set_data(res['x'][start:frame+1], res['v'][start:frame+1])
    return anim_points + anim_trails

anim = FuncAnimation(fig, animate, frames=n_steps, interval=20, blit=True)

# ============ 打印能量误差统计 ============
print("\n" + "=" * 65)
print("能量守恒误差分析 (相对漂移百分比，初始条件：从最大位移释放)")
print("=" * 65)
for method_name, _, _, _ in methods:
    res0 = results[method_name][0]
    print(f"  {method_name:10s}：{res0['E_drift']:+.4f}%")

print("\n" + "=" * 65)
print("所有初始条件下的能量漂移")
print("=" * 65)
for method_name, _, _, _ in methods:
    print(f"\n{method_name}:")
    for idx, (x0, v0) in enumerate(initial_conditions):
        res = results[method_name][idx]
        print(f"  {labels_ic[idx]}: {res['E_drift']:+.4f}%")

# 显示图形
plt.tight_layout()
plt.show()

# 可选：保存动画
# anim.save('harmonic_oscillator_4methods.gif', writer='pillow', fps=30)
