import numpy as np
import pylab as pl


# recurrence relation is H_n(x) = 2*x*H_{n-1}(x) - 2*(n-1)*H_{n-2}(x)
# source https://en.wikipedia.org/wiki/Hermite_polynomials
def hermite(x: np.ndarray, n: int) -> np.ndarray:
    if n == 0:
        return 1.0 + 0.0 * x
    if n == 1:
        return 2.0 * x
    return 2.0 * x * hermite(x, n - 1) - 2.0 * (n - 1) * hermite(x, n - 2)


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


x_vals = np.linspace(-5, 5, 1000)
gaussian = np.exp(-x_vals ** 2)

psi_funcs = []
for n in range(4):
    psi = hermite(x_vals, n) * gaussian
    normalize(psi)
    psi_funcs.append(psi)

for i, psi in enumerate(psi_funcs):
    psi_sqrd = np.abs(psi) ** 2
    # integrate in the region outside of -sqrt(2):sqrt(2)
    prob = (np.trapz(psi_sqrd[x_vals < -np.sqrt(2)], x_vals[x_vals < -np.sqrt(2)])
            + np.trapz(psi_sqrd[x_vals > np.sqrt(2)], x_vals[x_vals > np.sqrt(2)]))
    print(f'prob_outside_{i} = {prob}')