from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False

output_dir = Path(__file__).resolve().parent / "figures"
output_dir.mkdir(exist_ok=True)

nx = 81
nt = 40
CFL = 0.4

x = np.linspace(0.0, 2.0, nx)
dx = x[1] - x[0]

u = np.ones(nx)
u[(x >= 0.5) & (x <= 1.0)] = 2.0
initial_u = u.copy()

dt = CFL * dx / np.max(u)
snapshots = {0: u.copy()}

for n in range(1, nt + 1):
    old_u = u.copy()
    u[1:] = old_u[1:] - old_u[1:] * dt / dx * (old_u[1:] - old_u[:-1])
    u[0] = 1.0

    if n in [10, 20, 40]:
        snapshots[n] = u.copy()

plt.figure(figsize=(7, 4))
for step, values in snapshots.items():
    plt.plot(x, values, linewidth=2, label=f"第 {step} 步")
plt.title("一维非线性对流")
plt.xlabel("x")
plt.ylabel("u")
plt.ylim(0.8, 2.2)
plt.grid(True, alpha=0.3)
plt.legend()
plt.tight_layout()
plt.savefig(output_dir / "result.png", dpi=160)
plt.close()

print("一维非线性对流程序运行完成。")
print(f"dx={dx:.6f}, dt={dt:.6f}, 最大初始 CFL={np.max(initial_u)*dt/dx:.3f}")
print(f"最终 u 范围=[{u.min():.3f}, {u.max():.3f}]")