import numpy as np
import pylab as pl


def ladder_plus(xs: np.ndarray, psi: np.ndarray) -> np.ndarray:
    dx = xs[1] - xs[0]
    return -np.gradient(psi, dx) + xs * psi


def ladder_minus(xs: np.ndarray, psi: np.ndarray) -> np.ndarray:
    dx = xs[1] - xs[0]
    return np.gradient(psi, dx) + xs * psi


def normalize(psi: np.ndarray) -> None:
    psi /= np.sqrt(np.sum(psi * np.conjugate(psi)))


x_vals = np.linspace(-5, 5, 1000)

psi_0 = np.exp(-x_vals ** 2)
normalize(psi_0)

psi_funcs = [psi_0]
for i in range(1, 4):
    psi_next = ladder_plus(x_vals, psi_funcs[i - 1])
    normalize(psi_next)
    psi_funcs.append(psi_next)

pl.rcParams['figure.dpi'] = 300
fig, axs = pl.subplots(2, 2, tight_layout=True)
fig.tight_layout(pad=2.0)
axs[0, 0].plot(x_vals, psi_funcs[0])
axs[0, 0].set_title('ψ0')
axs[0, 0].grid()
axs[0, 1].plot(x_vals, psi_funcs[1], 'tab:orange')
axs[0, 1].set_title('ψ1')
axs[0, 1].grid()
axs[1, 0].plot(x_vals, psi_funcs[2], 'tab:green')
axs[1, 0].set_title('ψ2')
axs[1, 0].grid()
axs[1, 1].plot(x_vals, psi_funcs[3], 'tab:red')
axs[1, 1].set_title('ψ3')
axs[1, 1].grid()
pl.show()