#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Title: Quantum Tunneling Simulation
Author: Michael T. Power
Affiliation: Astrophysics PhD Student, University of Toronto
Date: Feb 1, 2025
Email: michael.power@astro.utoronto.ca

Description: This script animates the quantum tunneling of a Gaussian
wave packet through a square potential barrier. The time-dependent Schrödinger equation
is solved using a simple (naive, and possibly unstable in general) central difference
method in space, with a fourth-order Runge-Kutta method in time. The animation shows
the probability density |psi(x,t)|^2, as well as the real and imaginary parts of the
wave function psi(x,t). The potential barrier is also plotted for reference (though of
course you can't properly plot a potential on an amplitude plot, so it's just to show
the spatial location of the barrier). The simulation parameters can be adjusted to
explore different scenarios of quantum tunneling.
"""

import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
from numba import njit, prange

# -------------------------
# Simulation parameters.
# -------------------------
CONFIG = {
    "Nx": 1024,  # Number of spatial points.
    "x_min": -200.0,  # Minimum x value.
    "x_max": 200.0,  # Maximum x value.
    "t_final": 100.0,  # Total simulation time.
    "hbar": 1.0,  # Planck's constant (natural units).
    "m": 1.0,  # Particle mass (natural units).
    "V0": 1.8,  # Barrier height.
    "a": 15.0,  # Barrier half-width.
    "x0": -50.0,  # Initial center position of the packet.
    "k0": 2.0,  # Initial wave number (momentum).
    "sigma": 10.0,  # Width of the Gaussian.
}

# Spatial grid.
Nx = CONFIG["Nx"]
x = np.linspace(CONFIG["x_min"], CONFIG["x_max"], Nx)
dx = x[1] - x[0]  # Spatial step size.

# Physical constants.
hbar = CONFIG["hbar"]
m = CONFIG["m"]

# Time step is based on the Courant-Friedrichs-Lewy (CFL) condition. The Schrödinger
# equation is parabolic, so the CFL condition is similar to the heat equation. The
# CFL factor is set to 0.2 for stability, but this may need to be adjusted for
# different methods or simulations.
cfl_factor = 0.2  # Probably safe for RK4 here?
dt = cfl_factor * (m * dx**2) / hbar  # Parabolic CFL condition for Schrödinger eq.

# Compute total number of time steps based on simulation time.
t_final = CONFIG["t_final"]
Nt = int(np.ceil(t_final / dt))  # Ensure integer time steps.

# Fixed reference values for scaling (based on Nx = 1024 case).
# This just seemed to be nice for the live animation. Obviously as you increase Nx
# the live animation will slow down, so you may need to adjust the scaling.

Nt_ref = 10000  # Reference total number of time steps.
steps_ref = 10  # Reference steps per frame.

# Scale steps_per_frame to maintain a consistent animation speed.
steps_per_frame = int(np.ceil(steps_ref * (Nt / Nt_ref)))

# -------------------------
# Define the initial wave function (Gaussian wave packet).
# -------------------------
x0 = CONFIG["x0"]
k0 = CONFIG["k0"]
sigma = CONFIG["sigma"]

# Normalized Gaussian wave packet.
psi0 = (
    (1 / (sigma * np.sqrt(np.pi))) ** 0.5
    * np.exp(1j * k0 * x)
    * np.exp(-((x - x0) ** 2) / (2 * sigma**2))
)

# -------------------------
# Define the potential barrier.
# -------------------------
# Square barrier parameters.
V = np.zeros(Nx)  # Initialize potential to zero.
# Barrier centered at x = 0, set potential to V0 within [-a, a].
V[np.abs(x) < CONFIG["a"]] = CONFIG["V0"]


# -------------------------
# Define the RHS of the Schrödinger equation (central finite-difference version).
# -------------------------
def schrodinger_rhs(psi, V, dx, hbar, m):
    """
    Computes dpsi/dt for the time-dependent Schrödinger equation: dpsi/dt = -i/hbar *
    H psi = -i/hbar * [-(hbar^2/(2*m)) d^2 psi/dx^2 + V psi] The second derivative is
    computed via a central difference with periodic boundary conditions for simplicity.

    A second-order central finite-difference is:
    f_xx = (f(x+i) - 2f(x) + f(x-i)) / dx^2
    where dx is the step size.
    """
    # Use NumPy’s roll to handle periodic BCs in a very simple way. Note that
    # np.roll(psi, -1) is equivalent to psi(i + 1), and np.roll(psi, 1) is psi(i - 1)
    # so this just wraps around in a really nice way. It also prevents psi(i+1) from
    # going out of bounds, since it just wraps around to psi(0) in that case.
    psi_xx = (np.roll(psi, -1) - 2 * psi + np.roll(psi, 1)) / dx**2

    # Compute the Hamiltonian operator acting on psi.
    H_psi = -(hbar**2 / (2 * m)) * psi_xx + V * psi

    # Python handles the complex numbers simply, so...
    return -1j / hbar * H_psi


# -------------------------
# RK4 time evolution function.
# -------------------------
def RK4_step(psi, dt, V, dx, hbar, m):
    """
    Computes a fourth-order Runge-Kutta time evolution of the wave function psi, using
    the standard RK4 method.
    """
    k1 = schrodinger_rhs(psi, V, dx, hbar, m)
    k2 = schrodinger_rhs(psi + 0.5 * dt * k1, V, dx, hbar, m)
    k3 = schrodinger_rhs(psi + 0.5 * dt * k2, V, dx, hbar, m)
    k4 = schrodinger_rhs(psi + dt * k3, V, dx, hbar, m)
    psi_next = psi + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
    return psi_next


# -------------------------
# Animation update function.
# -------------------------
def update(frame, psi, line_prob, line_re, line_im):
    """
    Evolves the wave function for several time steps per frame to smooth the animation,
    and updates the probability density and real/imaginary parts of the wave function.
    """
    for _ in range(steps_per_frame):
        psi[:] = RK4_step(psi, dt, V, dx, hbar, m)

    # print((psi.real**2+psi.imag**2).sum())
    line_prob.set_ydata(np.abs(psi) ** 2)
    line_re.set_ydata(np.real(psi))
    line_im.set_ydata(np.imag(psi))
    return line_prob, line_re, line_im


def initialize_plot():
    """
    Sets up the figure and initial plots for the animation.
    """
    fig, (ax_prob, ax_reim) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)

    V_scale_prob = np.max(np.abs(psi0)) ** 2 * 1.5
    V_scale_reim = np.max(np.abs(psi0)) * 1.5

    ax_prob.plot(x, V / CONFIG["V0"] * V_scale_prob, "k--", lw=1.5, label="Potential")
    (line_prob,) = ax_prob.plot(
        x, np.abs(psi0) ** 2, "b-", lw=2, label=r"$|\psi(x,t)|^2$"
    )
    ax_prob.set_ylabel(r"$|\psi(x,t)|^2$")
    ax_prob.set_title("Quantum Tunneling")
    ax_prob.legend(loc="upper right")

    ax_reim.plot(x, V / CONFIG["V0"] * V_scale_reim, "k--", lw=1.5, label="Potential")
    (line_re,) = ax_reim.plot(x, np.real(psi0), "b-", lw=2, label=r"Re{$\psi$}")
    (line_im,) = ax_reim.plot(x, np.imag(psi0), "r-", lw=2, label=r"Im{$\psi$}")

    return fig, line_prob, line_re, line_im


# -------------------------
# Main function.
# -------------------------
def main():
    """
    Main function to set up the animation and run the simulation.
    """
    psi = psi0.copy()  # Set psi to initial condition at t = 0.

    fig, line_prob, line_re, line_im = initialize_plot()

    # Create animation for the live display.
    ani = animation.FuncAnimation(
        fig,
        update,
        fargs=(psi, line_prob, line_re, line_im),
        frames=Nt // steps_per_frame,
        interval=30,
        blit=False,
    )

    plt.tight_layout()
    plt.show()

    # Create a new animation for saving.
    psi = psi0.copy()  # Reset psi again before saving.
    fig, line_prob, line_re, line_im = initialize_plot()  # Reinitialize the figure.

    ani_save = animation.FuncAnimation(
        fig,
        update,
        fargs=(psi, line_prob, line_re, line_im),
        frames=Nt // steps_per_frame,
        interval=30,
        blit=False,
    )

    # Save the animation as an .mp4 file, but this requires ffmpeg.
    ani_save.save(
        "QuantumTunneling_prob_re_im.mp4", fps=30, extra_args=["-vcodec", "libx264"]
    )


# Run the main function if this is not being imported as a module.
if __name__ == "__main__":
    main()