"""Quantum tunnelling through a rectangular barrier, solved exactly.

Nothing here is stepped forward in time. The stationary Schrodinger equation is
solved in closed form in the three regions where the potential is constant, the
pieces are matched at the two walls, and the packet is rebuilt as a
superposition of those exact states:

    psi(x,t) = 1/sqrt(2 pi) INT dk  A(k) exp(-i k x0) phi_k(x) exp(-i k^2 t/2)

Every phi_k(x) in that integral is a formula. The only thing this program
approximates is the integral over k, which is a Fourier synthesis -- the same
act as summing a Fourier series to draw a square wave. psi at one time is
computed without computing psi at any earlier time.

Units: hbar = m = 1.

Run:  python make_figures.py      ->  writes the PNGs and the GIF into ../figures/
"""

from pathlib import Path

import numpy as np
import matplotlib

matplotlib.use("Agg")            # draw into files, never open a window

import matplotlib.pyplot as plt
from matplotlib.animation import PillowWriter
from matplotlib.gridspec import GridSpec


# =====================================================================
# 1. Where the pictures go
# =====================================================================

OUT = Path(__file__).resolve().parents[1] / "figures"
OUT.mkdir(parents=True, exist_ok=True)


# =====================================================================
# 2. How the pictures look
# =====================================================================
# Deliberately the palette of the numerical companion project: the two reports
# are meant to be read against each other, and a reader comparing two figures
# should not have to translate two colour schemes as well.

INK = "#0F172A"
MUTED = "#64748B"
GRID = "#CBD5E1"

SERIES = ["#1D4ED8", "#B45309", "#059669", "#7C3AED", "#BE123C"]

WAVE = SERIES[0]
WAVE_FILL = "#C7D6F7"
BARRIER_EDGE = "#334155"
BARRIER_FILL = "#EEF1F5"
ACCENT = SERIES[1]
GHOST = SERIES[3]

plt.rcParams.update({
    "mathtext.fontset": "cm",
    "font.family": "serif",
    "font.size": 17,
    "text.color": INK,
    "axes.labelcolor": INK,
    "axes.labelsize": 18,
    "axes.edgecolor": MUTED,
    "axes.titlesize": 19,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "xtick.color": MUTED,
    "ytick.color": MUTED,
    "xtick.labelsize": 15,
    "ytick.labelsize": 15,
    "legend.fontsize": 16,
    "legend.frameon": False,
    "grid.color": GRID,
    "figure.dpi": 120,
    "savefig.dpi": 190,
    "savefig.bbox": "tight",
    "savefig.facecolor": "white",
})


# =====================================================================
# 3. The barrier and the packet
# =====================================================================

hbar = 1.0
mass = 1.0

V0 = 2.0                          # barrier height
A_WIDTH = 1.0                     # barrier width -- the companion's, exactly
BAR_LEFT = -A_WIDTH / 2
BAR_RIGHT = +A_WIDTH / 2

X0 = -40.0                        # where the packet starts
SIGMA = 5.0                       # how wide it is
K0 = 1.6                          # its mean wavenumber
ENERGY = K0**2 / (2 * mass)       # 1.28, which is 0.64 V0



# =====================================================================
# 4. The exact stationary states
# =====================================================================
# In each of the three regions the potential is constant, so
#
#       phi'' + 2(E - V) phi = 0
#
# has constant coefficients and its solutions are exponentials. Outside the
# barrier they are e^{+-ikx} with k = sqrt(2E); inside they are e^{+-Kx} with
# K = sqrt(2(V0 - E)). Writing K as a complex square root covers both sides of
# the barrier top in one formula: above it K is imaginary, and the sinh in the
# transmission formula turns into a sin on its own.


def wavenumbers(k):
    """(k, K) for an incident wavenumber k. K is imaginary above the barrier."""
    E = k**2 / (2 * mass)
    K = np.sqrt(2 * mass * (V0 - E) + 0j) / hbar
    return k, K


def _pair(p, x):
    """The 2x2 matrix that reads (phi, phi') off the amplitudes of e^{+-px}."""
    return np.array([[np.exp(p * x), np.exp(-p * x)],
                     [p * np.exp(p * x), -p * np.exp(-p * x)]])


def scattering_amplitudes(k):
    """Exact r(k), t(k) and the two interior amplitudes, by matching.

    phi = e^{ikx} + r e^{-ikx}     to the left of the barrier
        = C e^{Kx} + D e^{-Kx}     inside it
        = t e^{ikx}                to the right

    Continuity of phi and phi' at each wall is four equations. Written as
    matrices the two walls become M(ik,x)(A,B)^T = M(K,x)(C,D)^T, so walking in
    from the right with a unit transmitted amplitude and rescaling at the end
    gives every constant at once -- and never divides by a large exponential.
    """
    k, K = wavenumbers(k)

    right = np.array([1.0 + 0j, 0.0 + 0j])                     # (t, 0), t = 1
    interior = np.linalg.solve(_pair(K, BAR_RIGHT),
                               _pair(1j * k, BAR_RIGHT) @ right)
    left = np.linalg.solve(_pair(1j * k, BAR_LEFT),
                           _pair(K, BAR_LEFT) @ interior)

    incident, reflected = left                                  # A and B
    return (reflected / incident,                               # r
            1.0 / incident,                                     # t
            interior[0] / incident,                             # C
            interior[1] / incident)                             # D


def phi(k, x):
    """The exact scattering state phi_k(x), on the whole line."""
    r, t, C, D = scattering_amplitudes(k)
    _, K = wavenumbers(k)

    out = np.empty_like(x, dtype=complex)
    L = x < BAR_LEFT
    M = (x >= BAR_LEFT) & (x <= BAR_RIGHT)
    R = x > BAR_RIGHT

    out[L] = np.exp(1j * k * x[L]) + r * np.exp(-1j * k * x[L])
    out[M] = C * np.exp(K * x[M]) + D * np.exp(-K * x[M])
    out[R] = t * np.exp(1j * k * x[R])
    return out


def transmission(energies, width=A_WIDTH):
    """T(E) in closed form, straight from the algebra of the report:

        T = [ 1 + V0^2 sinh^2(K a) / (4 E (V0 - E)) ]^{-1}

    One line covers both sides of the barrier top, because sinh(i q a) is
    i sin(q a) and the square puts the sign back.
    """
    E = np.asarray(energies, dtype=float) + 0j
    K = np.sqrt(2 * mass * (V0 - E)) / hbar
    with np.errstate(divide="ignore", invalid="ignore"):
        T = 1.0 / (1 + (V0**2 * np.sinh(K * width)**2) / (4 * E * (V0 - E)))
    # E = V0 is a removable singularity; fill it with its limit.
    return np.where(np.isfinite(T.real), T.real, 1.0 / (1 + width**2 * V0 / 2))


# =====================================================================
# 5. The packet, as a superposition of those exact states
# =====================================================================
# The Gaussian
#
#   psi(x,0) = (2 pi sigma^2)^{-1/4} exp(-(x-x0)^2/4sigma^2) exp(i k0 (x-x0))
#
# has the Fourier transform
#
#   A(k) = (2 sigma^2 / pi)^{1/4} exp(-sigma^2 (k-k0)^2),
#
# one integral done by hand. Sending the packet at the barrier means replacing
# each plane wave e^{ikx} in it by the exact phi_k(x), and letting each carry
# its own phase e^{-i k^2 t/2} forward in time.

K_HALF_RANGE = 0.75               # +- this many wavenumbers around k0
N_K = 1600                        # points in the k integral
k_grid = np.linspace(K0 - K_HALF_RANGE, K0 + K_HALF_RANGE, N_K)
dk = k_grid[1] - k_grid[0]

A_k = (2 * SIGMA**2 / np.pi)**0.25 * np.exp(-SIGMA**2 * (k_grid - K0)**2)

# The x grid is only where the answer is drawn. It is not a grid anything is
# solved on: refining it costs plotting time and changes no result.
X_MIN, X_MAX = -110.0, 80.0
N_X = 2400
x = np.linspace(X_MIN, X_MAX, N_X)
dx = x[1] - x[0]

print("building the exact stationary states (%d of them) ..." % N_K)
PHI = np.empty((N_K, N_X), dtype=complex)
for i, k in enumerate(k_grid):
    PHI[i] = phi(k, x)

weight = A_k * np.exp(-1j * k_grid * X0) * dk / np.sqrt(2 * np.pi)
omega = k_grid**2 / (2 * mass)


def psi_at(t):
    """psi(x,t): the superposition integral, evaluated at one time."""
    return (weight * np.exp(-1j * omega * t)) @ PHI


def psi_free(x, t):
    """The same packet with no barrier -- and this one is closed form.

    Worth having twice over: it is the only honest way to see the delay, and
    checking it against the same superposition with phi_k = e^{ikx} tests both.
    """
    u = x - X0
    spread = 1 + 1j * hbar * t / (2 * mass * SIGMA**2)
    return ((2 * np.pi * SIGMA**2)**-0.25 / np.sqrt(spread)
            * np.exp((-u**2 / (4 * SIGMA**2) + 1j * K0 * u
                      - 1j * K0**2 * t / 2) / spread))


# =====================================================================
# 6. What the formulas say, before any picture is drawn
# =====================================================================

# Each plane wave in the packet is transmitted with probability T(E_k), so the
# packet is transmitted with the average of T over its own spectrum. No
# evolution required to know the answer.
spectrum = A_k**2
spectrum = spectrum / (spectrum.sum() * dk)
T_of_k = transmission(k_grid**2 / (2 * mass))
T_packet = float((spectrum * T_of_k).sum() * dk)

# The transmitted spectrum is the incident one filtered by T(k), which rises
# with k -- so what gets through is faster than what arrived.
transmitted_spectrum = spectrum * T_of_k
transmitted_spectrum = transmitted_spectrum / (transmitted_spectrum.sum() * dk)
k_mean_in = float((spectrum * k_grid).sum() * dk)
k_mean_out = float((transmitted_spectrum * k_grid).sum() * dk)

# R + T = 1 is an identity of the formulas here, not a diagnostic of a run.
unitarity = 0.0
for k in k_grid[::40]:
    r, t, _, _ = scattering_amplitudes(k)
    unitarity = max(unitarity, abs(abs(r)**2 + abs(t)**2 - 1.0))

# The closed-form free packet against the same superposition: a check on both.
free_by_sum = (A_k * np.exp(-1j * k_grid * X0) * dk / np.sqrt(2 * np.pi)
               * np.exp(-1j * omega * 30.0)) @ np.exp(1j * np.outer(k_grid, x))
free_check = float(np.abs(free_by_sum - psi_free(x, 30.0)).max())

# The decay constant inside the barrier, at the packet's mean energy.
KAPPA = float(np.sqrt(2 * mass * (V0 - ENERGY)) / hbar)

print()
print("V0 = %.1f,  a = %.1f,  E = %.3f = %.2f V0" % (V0, A_WIDTH, ENERGY, ENERGY / V0))
print("kappa = %.4f,  kappa a = %.3f" % (KAPPA, KAPPA * A_WIDTH))
print("|psi|^2 falls by e^{-2 kappa a} = %.4g across the barrier"
      % np.exp(-2 * KAPPA * A_WIDTH))
print("exact T at the packet's mean energy : %.5f" % transmission([ENERGY])[0])
print("exact T of the packet as a whole    : %.5f" % T_packet)
print("same formula at the companion's a=1 : %.5f"
      % transmission([ENERGY])[0])
print("mean k in = %.4f   mean k out = %.4f   (+%.1f%%)"
      % (k_mean_in, k_mean_out, 100 * (k_mean_out / k_mean_in - 1)))
print("|r|^2 + |t|^2 - 1, worst case       : %.2e" % unitarity)
print("closed-form free packet vs the sum  : %.2e" % free_check)
print()


# =====================================================================
# 7. The frames
# =====================================================================

T_FINAL = 55.0
N_FRAMES = 140
times = np.linspace(0.0, T_FINAL, N_FRAMES)

print("evaluating the superposition at %d times ..." % N_FRAMES)
density = np.empty((N_FRAMES, N_X))
p_left = np.empty(N_FRAMES)
p_inside = np.empty(N_FRAMES)
p_right = np.empty(N_FRAMES)

inside = (x >= BAR_LEFT) & (x <= BAR_RIGHT)

for i, t in enumerate(times):
    psi = psi_at(t)
    density[i] = np.abs(psi)**2
    p_left[i] = np.sum(density[i][x < BAR_LEFT]) * dx
    p_inside[i] = np.sum(density[i][inside]) * dx
    p_right[i] = np.sum(density[i][x > BAR_RIGHT]) * dx

R_final, T_final = p_left[-1], p_right[-1]
print("from the evolved packet at t = %.0f:  R = %.4f   T = %.4f   R+T = %.4f"
      % (T_FINAL, R_final, T_final, R_final + T_final))
print()


PLOT_MIN, PLOT_MAX = -100.0, 80.0   # the companion project's window


PLOT_MIN, PLOT_MAX = -100.0, 80.0   # the companion project's window


# =====================================================================
# 9. Drawing conventions
# =====================================================================
# The palette, the type and the layouts are the numerical companion project's,
# unchanged. The two reports solve the same problem two ways and are meant to be
# read one after the other, so the thing that differs between them should be the
# method and nothing else -- a reader who has to relearn a colour or a layout
# between the two pays for something that carries no information.

def draw_barrier(ax, height):
    """The barrier as neutral structure behind the data, never as a series."""
    ax.fill_between([BAR_LEFT, BAR_RIGHT], 0, height, color=BARRIER_FILL,
                    hatch="///", edgecolor=BARRIER_EDGE, linewidth=0.0, zorder=1)
    ax.plot([BAR_LEFT, BAR_LEFT, BAR_RIGHT, BAR_RIGHT],
            [0, height, height, 0], color=BARRIER_EDGE, lw=2.0, zorder=2)


def draw_packet(ax, rho, lw=2.6):
    ax.fill_between(x, rho, color=WAVE_FILL, zorder=3)
    line, = ax.plot(x, rho, color=WAVE, lw=lw, zorder=4)
    return line


def tag(ax, position, text, colour=INK, size=16, ha="left", va="top"):
    """A label on a white plate, so it can sit anywhere without being read
    through the grid or the curve underneath it."""
    return ax.text(position[0], position[1], text, transform=ax.transAxes,
                   fontsize=size, color=colour, ha=ha, va=va, zorder=9,
                   bbox=dict(boxstyle="round,pad=0.32", facecolor="white",
                             alpha=0.92, edgecolor="none"))


def label_at(ax, xy, text, colour, size=15, ha="left", va="center"):
    """The same, in data coordinates: a label put beside its own mark."""
    return ax.text(xy[0], xy[1], text, fontsize=size, color=colour, ha=ha,
                   va=va, zorder=9,
                   bbox=dict(boxstyle="round,pad=0.28", facecolor="white",
                             alpha=0.93, edgecolor="none"))


# =====================================================================
# 10. Figure 1 -- the three regions and what phi is in each
# =====================================================================
# The setup figure of an analytic project is not a picture of a packet: it is a
# picture of the algebra. Where each formula holds, and where the two matching
# conditions are imposed.

fig, ax = plt.subplots(figsize=(11.0, 4.9), layout="constrained")

ax.fill_between([BAR_LEFT, BAR_RIGHT], 0, V0, color=BARRIER_FILL,
                hatch="///", edgecolor=BARRIER_EDGE, linewidth=0.0)
ax.plot([-5, BAR_LEFT, BAR_LEFT, BAR_RIGHT, BAR_RIGHT, 5],
        [0, 0, V0, V0, 0, 0], color=BARRIER_EDGE, lw=2.4)

ax.axhline(ENERGY, color=ACCENT, ls="--", lw=2.0, zorder=3)
label_at(ax, (-4.85, ENERGY), "$E = %.2f$" % ENERGY, ACCENT, ha="left")
ax.text(0.0, V0 + 0.08, "$V_0 = %.0f$,  width $a = %.0f$" % (V0, A_WIDTH),
        ha="center", va="bottom", fontsize=16, color=BARRIER_EDGE)

for wall, name in ((BAR_LEFT, "$x=-a/2$"), (BAR_RIGHT, "$x=+a/2$")):
    ax.plot([wall], [0], marker="o", ms=10, color=ACCENT, zorder=5,
            markeredgecolor="white", markeredgewidth=2)
ax.text(BAR_LEFT - 0.15, -0.30, "$x=-a/2$", ha="right", va="top", color=ACCENT,
        fontsize=15)
ax.text(BAR_RIGHT + 0.15, -0.30, "$x=+a/2$", ha="left", va="top", color=ACCENT,
        fontsize=15)
ax.text(0.0, -0.62, r"match $\varphi$ and $\varphi'$ at both walls",
        ha="center", va="top", color=ACCENT, fontsize=15)

ax.text(-3.1, 1.28, r"$e^{ikx} + r\,e^{-ikx}$", ha="center", fontsize=19,
        color=WAVE)
ax.text(-3.1, 1.00, "incident + reflected", ha="center", fontsize=14, color=MUTED)
ax.text(0.0, 0.92, r"$C e^{\kappa x} + D e^{-\kappa x}$", ha="center",
        fontsize=16, color=WAVE)
ax.text(0.0, 0.58, "no oscillation:\ntwo real exponentials", ha="center",
        va="top", fontsize=13, color=MUTED)
ax.text(3.1, 1.28, r"$t\,e^{ikx}$", ha="center", fontsize=19, color=WAVE)
ax.text(3.1, 1.00, "transmitted", ha="center", fontsize=14, color=MUTED)

ax.set_xlim(-5, 5)
ax.set_ylim(-1.0, V0 + 0.42)
ax.set_yticks([0, ENERGY, V0])
ax.set_yticklabels(["$0$", "$%.2f$" % ENERGY, "$%.0f$" % V0])
ax.set_xticks([])
ax.spines["bottom"].set_visible(False)
ax.set_ylabel("energy")
ax.set_title("Three regions, six constants, four matching conditions", pad=10)

fig.savefig(OUT / "fig_regions.png")
plt.close(fig)
print("wrote fig_regions.png")


# =====================================================================
# 11. Figure 2 -- the exact transmission curve
# =====================================================================
# The companion project measures one point of this curve. The formula draws all
# of it, including the part above the barrier, at no extra cost.

energies = np.linspace(0.004, 3.0, 1600) * V0
curve = transmission(energies)

fig, ax = plt.subplots(figsize=(11.0, 5.6), layout="constrained")

ax.axvspan(0, V0, facecolor=BARRIER_FILL, edgecolor="none", zorder=1)
ax.text(V0 * 0.5, 1.04, "classically forbidden,  $E < V_0$", ha="center",
        va="top", fontsize=15, color=MUTED)

ax.plot(energies, 1 - curve, color=MUTED, lw=2.0, ls=":", zorder=3)
ax.plot(energies, curve, color=WAVE, lw=3.0, zorder=4)

label_at(ax, (4.9, 0.80), "exact  $T(E)$", WAVE, size=17, ha="center")
label_at(ax, (4.9, 0.17), "exact  $R(E) = 1 - T(E)$", MUTED, size=15,
         ha="center")

T_here = transmission([ENERGY])[0]
ax.plot([ENERGY], [T_here], "o", color=ACCENT, ms=13, markeredgecolor="white",
        markeredgewidth=2.2, zorder=6)
ax.annotate("the packet sits here:  $T = %.3f$" % T_here,
            xy=(ENERGY + 0.05, T_here + 0.01), xytext=(2.05, 0.55),
            color=ACCENT, fontsize=16, va="center",
            arrowprops=dict(arrowstyle="->", color=ACCENT, lw=2.0,
                            connectionstyle="arc3,rad=0.22"))

# Perfect transmission above the barrier, wherever it is a whole number of half
# wavelengths thick. The formula hands these over for nothing.
res = [V0 + (n * np.pi / A_WIDTH)**2 / 2 for n in (1, 2, 3, 4)]
res = [E for E in res if E < 3 * V0]
if res:
    ax.plot(res, [1.0] * len(res), "o", color=SERIES[2], ms=9, zorder=6,
            markeredgecolor="white", markeredgewidth=1.8)
    ax.annotate(r"$T = 1$ exactly, whenever $qa = n\pi$", xy=(res[0], 1.0),
                xytext=(res[0], 0.62), ha="center", color=SERIES[2],
                fontsize=15,
                arrowprops=dict(arrowstyle="->", color=SERIES[2], lw=1.8))

ax.set_xlim(0, 3 * V0)
ax.set_ylim(0, 1.09)
ax.set_xlabel("incident energy  $E$")
ax.set_ylabel("probability")
ax.grid(ls="--", alpha=0.45)
ax.set_title("One formula, every energy at once", pad=10)

fig.savefig(OUT / "fig_transmission.png")
plt.close(fig)
print("wrote fig_transmission.png")


# =====================================================================
# 12. Figure 3 -- the barrier as a filter on the packet's spectrum
# =====================================================================
# This figure exists only because the problem was solved analytically. Stepping
# the equation forward gives one number, T. The formula gives T for every plane
# wave in the packet at once, so multiplying the packet's own spectrum by it
# shows the barrier selecting the fast components out of what arrives -- and the
# transmitted packet leaving faster than the incident one.

fig, (ax_t, ax_s) = plt.subplots(2, 1, figsize=(11.0, 7.0), sharex=True,
                                 height_ratios=[1.0, 1.4], layout="constrained")

ax_t.plot(k_grid, T_of_k, color=WAVE, lw=3.0)
ax_t.set_ylabel("$T(k)$")
ax_t.grid(ls="--", alpha=0.45)
ax_t.set_title("Why the transmitted packet is the faster one", pad=10)
label_at(ax_t, (k_grid[0] + 0.05, T_of_k[-1] * 0.72),
         "transmission of each plane wave in the packet:\nover this range it "
         r"rises by a factor of $%.1f$" % (T_of_k[-1] / T_of_k[0]),
         INK, size=15, ha="left", va="top")

ax_s.fill_between(k_grid, spectrum, color=WAVE_FILL, zorder=2)
ax_s.plot(k_grid, spectrum, color=WAVE, lw=2.6, zorder=3)
ax_s.plot(k_grid, transmitted_spectrum, color=ACCENT, lw=2.8, zorder=4)

peak = max(spectrum)
label_at(ax_s, (k_mean_in - 0.11, peak * 0.60), "incident\n$|A(k)|^2$", WAVE,
         ha="right")
label_at(ax_s, (k_mean_out + 0.12, peak * 0.60),
         "transmitted\n" + r"$|A(k)|^2\,T(k)$," + "\nrenormalised", ACCENT,
         ha="left")

for centre, colour in ((k_mean_in, WAVE), (k_mean_out, ACCENT)):
    ax_s.plot([centre, centre], [0, peak * 1.04], color=colour, ls="--", lw=1.6,
              zorder=1)
ax_s.annotate("", xy=(k_mean_out, peak * 1.10), xytext=(k_mean_in, peak * 1.10),
              arrowprops=dict(arrowstyle="-|>", color=INK, lw=2.2,
                              mutation_scale=22))
label_at(ax_s, ((k_mean_in + k_mean_out) / 2, peak * 1.20),
         r"$\langle k\rangle:\ %.3f \to %.3f$,  $+%.1f\%%$"
         % (k_mean_in, k_mean_out, 100 * (k_mean_out / k_mean_in - 1)),
         INK, size=16, ha="center")

ax_s.set_xlim(k_grid[0], k_grid[-1])
ax_s.set_ylim(0, peak * 1.36)
ax_s.set_yticks([])
ax_s.set_xlabel("wavenumber  $k$")
ax_s.set_ylabel("spectral weight")
ax_s.grid(axis="x", ls="--", alpha=0.45)

fig.savefig(OUT / "fig_filter.png")
plt.close(fig)
print("wrote fig_filter.png")


# =====================================================================
# 13. Figure 4 -- four instants of the collision
# =====================================================================
# The same four-panel figure the companion project draws from its simulation,
# drawn here from the superposition instead. That it comes out the same is the
# result: two methods, one answer.

wanted = [0.0, 20.0, 35.0, T_FINAL]
notes = ["the packet approaches",
         "incident and reflected waves interfere",
         "reflected and transmitted parts separate",
         "final state"]

chosen = [int(np.argmin(np.abs(times - w))) for w in wanted]
in_view = (x > PLOT_MIN) & (x < PLOT_MAX)
y_max = max(density[i][in_view].max() for i in chosen) * 1.34

fig, axes = plt.subplots(4, 1, figsize=(11.0, 10.4), sharex=True,
                         layout="constrained")

for panel, i in enumerate(chosen):
    ax = axes[panel]
    draw_barrier(ax, y_max * 0.88)
    draw_packet(ax, density[i])
    ax.set_ylim(0, y_max)
    ax.set_yticks([])
    ax.grid(axis="x", ls="--", alpha=0.5)
    tag(ax, (0.015, 0.94), "$t = %.0f$   %s" % (times[i], notes[panel]))

tag(axes[3], (0.985, 0.94), "$R = %.2f$,  $T = %.2f$" % (R_final, T_final),
    colour=ACCENT, ha="right")
axes[3].set_xlim(PLOT_MIN, PLOT_MAX)
axes[3].set_xlabel("position  $x$")
axes[1].set_ylabel(r"$|\psi(x,t)|^2$  (same scale in every panel)")

fig.suptitle("Built out of formulas, not out of time steps", fontsize=20)
fig.savefig(OUT / "fig_snapshots.png")
plt.close(fig)
print("wrote fig_snapshots.png")


# =====================================================================
# 14. The animation, and the still frame the PDF shows
# =====================================================================
# The companion project's animation, to the letter: the packet across the whole
# line on top, the running probability on each side of the barrier underneath,
# the same barrier, the same packet, the same axes, the same colours. The two
# are meant to be watched one after the other, and the whole point is that they
# are indistinguishable. What differs is behind the picture -- one is stepped
# forward 13750 times, the other is a formula evaluated at each frame from
# scratch -- and if the two pictures did not agree, one of the two reports would
# be wrong.

def commentary(t):
    if t < 14:
        return "incident packet, still free"
    if t < 26:
        return "incident and reflected waves interfere"
    if t < 40:
        return "the packet splits at the barrier"
    return "reflected (left) and transmitted (right)"


animation_y_max = max(density[i][in_view].max() for i in range(N_FRAMES)) * 1.30


def build_figure(frame):
    fig = plt.figure(figsize=(11.0, 7.0))
    grid = GridSpec(2, 1, figure=fig, height_ratios=[2.4, 1.0], hspace=0.42,
                    left=0.08, right=0.97, top=0.86, bottom=0.10)
    ax_top = fig.add_subplot(grid[0])
    ax_bot = fig.add_subplot(grid[1])

    draw_barrier(ax_top, animation_y_max * 0.88)
    line = draw_packet(ax_top, density[frame])
    fill = ax_top.collections[-1]

    clock = tag(ax_top, (0.015, 0.95), "")
    note = tag(ax_top, (0.985, 0.95), "", colour=MUTED, size=15, ha="right")

    ax_top.set_xlim(PLOT_MIN, PLOT_MAX)
    ax_top.set_ylim(0, animation_y_max)
    ax_top.set_yticks([])
    ax_top.set_ylabel(r"$|\psi(x,t)|^2$")
    ax_top.set_xlabel("position  $x$")
    ax_top.grid(axis="x", ls="--", alpha=0.5)

    ax_bot.plot(times, p_left, color=WAVE, lw=2.6)
    ax_bot.plot(times, p_right, color=ACCENT, lw=2.6)
    ax_bot.text(times[-1] * 0.99, 0.92,
                r"left of barrier $\to R = %.2f$" % p_left[-1],
                ha="right", va="center", color=WAVE, fontsize=15)
    ax_bot.text(times[-1] * 0.99, 0.13,
                r"right of barrier $\to T = %.2f$" % p_right[-1],
                ha="right", va="center", color=ACCENT, fontsize=15)
    cursor = ax_bot.axvline(times[frame], color=INK, lw=2.0)
    ax_bot.set_xlim(times[0], times[-1])
    ax_bot.set_ylim(0, 1.18)
    ax_bot.set_yticks([0, 0.5, 1.0])
    ax_bot.set_yticklabels(["$0$", "$0.5$", "$1$"])
    ax_bot.set_xlabel("time  $t$")
    ax_bot.set_ylabel("probability", fontsize=15)
    ax_bot.grid(ls="--", alpha=0.5)

    fig.suptitle("Tunnelling through a rectangular barrier, solved exactly",
                 y=0.975, fontsize=21)
    fig.text(0.5, 0.905,
             "no time steps: every frame is the superposition integral, "
             "evaluated from scratch",
             ha="center", fontsize=16, color=MUTED)

    return fig, line, fill, clock, note, cursor


def show_frame(i, line, fill, clock, note, cursor):
    line.set_ydata(density[i])
    corners = np.vstack([[x[0], 0.0],
                         np.column_stack([x, density[i]]),
                         [x[-1], 0.0]])
    fill.set_paths([corners])
    clock.set_text("$t = %.1f$" % times[i])
    note.set_text(commentary(times[i]))
    cursor.set_xdata([times[i], times[i]])


poster_frame = int(np.argmin(np.abs(times - 35.0)))
handles = build_figure(poster_frame)
show_frame(poster_frame, *handles[1:])
handles[0].savefig(OUT / "anim_poster.png", dpi=140)
plt.close(handles[0])
print("wrote anim_poster.png")

print("rendering the animation ...")
handles = build_figure(0)
writer = PillowWriter(fps=12)
with writer.saving(handles[0], str(OUT / "anim.gif"), dpi=72):
    for i in range(len(times)):
        show_frame(i, *handles[1:])
        writer.grab_frame()
plt.close(handles[0])
print("wrote anim.gif")


# =====================================================================
# 15. What the run found
# =====================================================================

print()
print("--- numbers quoted in the report ---")
print("T at the mean energy E = %.2f          : %.5f"
      % (ENERGY, transmission([ENERGY])[0]))
print("T of the packet, from its spectrum     : %.5f" % T_packet)
print("T of the packet, from the evolved psi  : %.5f" % T_final)
print("R of the packet, from the evolved psi  : %.5f" % R_final)
print("R + T                                  : %.5f" % (R_final + T_final))
print("T(k) rises by a factor                 : %.2f over the packet"
      % (T_of_k[-1] / T_of_k[0]))
print("mean k, incident / transmitted         : %.4f / %.4f  (+%.1f%%)"
      % (k_mean_in, k_mean_out, 100 * (k_mean_out / k_mean_in - 1)))
print("worst |r|^2+|t|^2-1 over the packet    : %.2e" % unitarity)
print("free-packet formula vs superposition   : %.2e" % free_check)

print()
print("files written to %s:" % OUT)
for path in sorted(OUT.iterdir()):
    print("  %-20s %8.0f kB" % (path.name, path.stat().st_size / 1024))
