using Plots
using LinearAlgebra

m = 1
hbar = 1
xmin = -6.5
xmax = 6.5
N = 2000
x = collect(range(xmin, stop = xmax, length = N + 1)) # one extra, want 1000 spaces
dx = x[2] - x[1]
println("dx = ", dx)
p = 100.0
V0 = 2 * p^2 / (2.0 * m)
sig = 0.25
x0 = 2.0 # a bit away to not interfere with the walls
barrier_width = 0.025

# setup the potential barrier
V = 0 .* x
# println("V = ", length(V))
for i in 1:length(V)
	if x[i] > 0 && x[i] < barrier_width
		V[i] = V0
	end
end
plot(x, V, xlabel = "x", ylabel = "V(x)", title = "Potential V(x)", legend = :none)

# make the initial wavefunction
psi0 = exp.(-((x[2:end-1] .+ x0) .^ 2) / (2 * sig^2)) .* exp.(1.0im * p * x[2:end-1])
C = sum(abs2.(psi0) .* dx)
psi0 = psi0 / sqrt(C)
plot(x[2:end-1], abs.(psi0), xlabel = "x", ylabel = "Re(psi)", title = "Initial wavefunction", legend = :none)

# make the Hamiltonian
A = (hbar^2 / (2 * m * dx^2))
# main diagonal
d = A * ones(N - 1) + V[2:end-1] # add potential
# upper and lower diagonals
dl = (A / -2.0) * ones(N - 2)
du = (A / -2.0) * ones(N - 2)
H = Tridiagonal(dl, d, du)

# diagonalize matrix
E, psi = eigen(H)
psi1 = psi[:, 1]

# normalize
A = sum(abs2.(psi1) .* dx)
psi = psi ./ sqrt(A)
println("A = ", A)

c = zeros(ComplexF64, length(psi[:, 1]))
for i in 1:length(c)
	c[i] = sum(conj.(psi[:, i]) .* psi0 .* dx)
end

dt = 0.0001

anim = @animate for t in 0:dt:0.25
	tmp = zeros(ComplexF64, length(psi[:, 1]))
	for i in 1:length(c)
		tmp = tmp .+ c[i] .* exp.(-1.0im * E[i] * t / hbar) .* psi[:, i]
	end
	rounded_p = round(p, digits = 3)
	rounded_pot = round(V0, digits = 3)
	plot(x[2:end-1], abs.(tmp), xlabel = "x", ylabel = "abs(Ψ)", title = "Propagating Wavepacket (k_0 = $p, V0 = $rounded_pot)", label = "psi @ t = $t", ylims = (-2, 2), lw = 2)
	# plot the real
	plot!(x[2:end-1], real(tmp), xlabel = "x", label = "real")
	# plot the imaginary
	plot!(x[2:end-1], imag(tmp), xlabel = "x", label = "imag")
	endpoint_1 = 0.0
	endpoint_2 = barrier_width
	plot!([endpoint_1, endpoint_1], [0, 2], label = "", line = :dash, color = :red)
	plot!([endpoint_2, endpoint_2], [0, 2], label = "barrier (width = $barrier_width)", line = :dash, color = :red)
end

mp4(anim, "hw8/10-17-tunnelling-$barrier_width-$p.mp4", fps = 30)