import numpy as np
import pylab as pl

a = 10
D = 1
hbar = 1
m = 1
zs = np.array([2.991, 5.98, 8.96, 11.927])
l = zs / a
V0 = 2


def cot(x):
    return 1 / np.tan(x)


x = np.linspace(-20, 20, 1000)
x_pos = x[x > 0]
z0 = (a / hbar) * np.sqrt(2 * m * V0)
print(f'z0 = {z0}')
k = -l * cot(l * a)
F = np.sin(l * a) / np.exp(-k * a)

psis = []
for (i, z) in enumerate(zs):
    psi_conditions = [
        (0 < x_pos) & (x_pos <= a),
        (x_pos > a)
    ]
    psi_funcs = [
        lambda x: D * np.sin(l[i] * x),
        lambda x: F[i] * np.exp(-k[i] * x),
    ]
    psi_half = np.piecewise(x_pos, psi_conditions, psi_funcs)
    psi = np.concatenate((-np.flip(psi_half), psi_half))
    psis.append(psi)

V0_graph = np.piecewise(x, [(x < -a), (-a <= x) & (x <= a), (x > a)], [0, -V0, 0])

pl.rcParams['figure.dpi'] = 300
pl.plot(x, V0_graph)
for (i, psi) in enumerate(psis):
    pl.plot(x, psi, label=f'psi_odd_{i} z={zs[i]}')
pl.legend()
pl.show()