import numpy as np
import pylab as pl

a = 5
F = 1
D = 1
hbar = 1
m = 1
z = 6.933
l = z / a
V0 = 2

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

psi_conditions = [
    (0 < x_pos) & (x_pos <= a),
    (x_pos > a)
]
psi_funcs = [
    lambda x: D * np.cos(l * x),
    lambda x: F * np.exp(-k * x),
]
psi_half = np.piecewise(x_pos, psi_conditions, psi_funcs)
psi = np.concatenate((np.flip(psi_half), psi_half))

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

pl.plot(x, psi)
pl.plot(x, V0_graph)
pl.show()