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

# ----------------------------
# Physical parameters
# ----------------------------
hbar = 1.0
m = 1.0
V0 = 1.0

E_target = 1.1 * V0   # mean packet energy
x0 = -80.0            # initial packet center
sigma = 12.0          # packet width

# Choose k0 from approximate Gaussian mean energy
spread_E = hbar**2 / (8 * m * sigma**2)
k0 = np.sqrt(2 * m * (E_target - spread_E)) / hbar

print(f"k0 = {k0:.4f}")
print(f"Mean energy ≈ {E_target:.4f}")
print(f"Step height V0 = {V0:.4f}")

# ----------------------------
# Spatial grid
# ----------------------------
Nx = 1600
x_min, x_max = -150, 150
x = np.linspace(x_min, x_max, Nx)

# ----------------------------
# k-space grid
# ----------------------------
Nk = 900
k_width = 5 / sigma
k_min = max(1e-4, k0 - 6 * k_width)
k_max = k0 + 6 * k_width
k = np.linspace(k_min, k_max, Nk)
dk = k[1] - k[0]

E = hbar**2 * k**2 / (2 * m)

# Momentum-space Gaussian amplitude
A = np.exp(-sigma**2 * (k - k0)**2) * np.exp(-1j * k * x0)

# ----------------------------
# Scattering coefficients
# ----------------------------
k_step = np.sqrt(2 * m * V0) / hbar

q = np.empty_like(k, dtype=complex)

above = k > k_step
below = ~above

q[above] = np.sqrt(k[above]**2 - k_step**2)
q[below] = 1j * np.sqrt(k_step**2 - k[below]**2)

r = (k - q) / (k + q)
t = 2 * k / (k + q)

# ----------------------------
# Build scattering eigenstates phi_k(x)
# Shape: Nk x Nx
# ----------------------------
X = x[None, :]
K = k[:, None]
Q = q[:, None]
R = r[:, None]
T = t[:, None]

phi = np.zeros((Nk, Nx), dtype=complex)

left_region = x < 0
right_region = x >= 0

phi[:, left_region] = (
    np.exp(1j * K * X[:, left_region])
    + R * np.exp(-1j * K * X[:, left_region])
)

phi[:, right_region] = (
    T * np.exp(1j * Q * X[:, right_region])
)

# ----------------------------
# Time-dependent wave function
# psi(x,t) = integral A(k) phi_k(x) exp(-i E_k t / hbar) dk
# ----------------------------
def psi_t(t):
    phase = np.exp(-1j * E * t / hbar)
    psi = np.sum((A * phase)[:, None] * phi, axis=0) * dk
    psi /= np.sqrt(np.trapz(np.abs(psi)**2, x))
    return psi

# ----------------------------
# Animation
# ----------------------------
fig, ax = plt.subplots(figsize=(9, 4.5))

density_line, = ax.plot([], [], lw=2, label=r"$|\psi(x,t)|^2$")

# scaled potential step
V_scaled = np.where(x >= 0, 0.04, 0.0)
potential_line, = ax.plot(x, V_scaled, "--", lw=1.5, label="step potential, scaled")

ax.axvline(0, lw=1, alpha=0.4)
ax.set_xlim(-120, 120)
ax.set_ylim(0, 0.09)
ax.set_xlabel("x")
ax.set_ylabel(r"$|\psi(x,t)|^2$")
ax.set_title(r"Wave packet scattering from a potential step, $\langle E\rangle = 1.1V_0$")
ax.legend(loc="upper right")

time_text = ax.text(0.02, 0.88, "", transform=ax.transAxes)

frames = 420
t_max = 120
times = np.linspace(0, t_max, frames)

def init():
    density_line.set_data([], [])
    time_text.set_text("")
    return density_line, time_text

def update(frame):
    t_now = times[frame]
    psi = psi_t(t_now)
    density = np.abs(psi)**2

    density_line.set_data(x, density)
    time_text.set_text(f"t = {t_now:.1f}")

    return density_line, time_text

ani = FuncAnimation(
    fig,
    update,
    frames=frames,
    init_func=init,
    interval=30,
    blit=True
)

plt.tight_layout()
plt.show()

# To save:
#ani.save("step_scattering_plane_wave_expansion.gif", writer="pillow", fps=30)
# ani.save("step_scattering_plane_wave_expansion.mp4", writer="ffmpeg", fps=30)