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)
dx = x_vals[1] - x_vals[0]

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)

print("==========")
expect_xs = []
for i, psi in enumerate(psi_funcs):
    expect_xs.append(np.sum(np.conjugate(psi) * x_vals * psi))
    print(f'<x>_{i} = {expect_xs[i]}')

print("==========")
expect_ps = []
for i, psi in enumerate(psi_funcs):
    expect_ps.append(np.sum(np.conjugate(psi) * -np.gradient(psi, dx)))
    print(f'<p>_{i} = {expect_ps[i]}')

print("==========")
expect_x_sqrds = []
for i, psi in enumerate(psi_funcs):
    expect_x_sqrds.append(np.sum(np.conjugate(psi) * x_vals * x_vals * psi))
    print(f'<x^2>_{i} = {expect_x_sqrds[i]}')

print("==========")
expect_p_sqrds = []
for i, psi in enumerate(psi_funcs):
    expect_p_sqrds.append(np.sum(np.conjugate(psi) * -np.gradient(np.gradient(psi, dx), dx)))
    print(f'<p^2>_{i} = {expect_p_sqrds[i]}')

print("==========")
sigma_xs = []
for i, psi in enumerate(psi_funcs):
    sigma_xs.append(np.sqrt(expect_x_sqrds[i] - expect_xs[i] ** 2))
    print(f'sigma_x_{i} = {sigma_xs[i]}')

print("==========")
sigma_ps = []
for i, psi in enumerate(psi_funcs):
    sigma_ps.append(np.sqrt(expect_p_sqrds[i] - expect_ps[i] ** 2))
    print(f'sigma_p_{i} = {sigma_ps[i]}')

print("==========")
print("hbar/2 = 0.5")
for i, psi in enumerate(psi_funcs):
    print(f'sigma_x_{i} * sigma_p_{i} = {sigma_xs[i] * sigma_ps[i]}')