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]

C = sum(abs2.(psi0) .* dx)
psi0 = psi0 / sqrt(C)

function get_energy_levels_for_n(n, pot_plot)
	V = 0 .* x
	for i in 1:length(V)
		V[i] = abs(x[i])^n
	end
	# plot the potential
	println("Plotting potential for n = ", n)
	plot!(pot_plot, x, V, xlabel = "x", ylabel = "V(x)", title = "Potential V(x) = |x|^n", label = "n = $n")
	println("Calculating energy levels for n = ", n)
	# 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)

	println("Diagonalizing matrix for n = ", n)

	# diagonalize matrix
	E, psi = eigen(H)

	println("Energy levels for n = ", n, ": ", E[1:10])
	return E
end

po = plot!()
energy_levels = []
n_max = 19
for n in 1:n_max
	E = get_energy_levels_for_n(n, po)
	push!(energy_levels, E)
end
savefig(po, "hw8/10-3-potentials.png")

# plot the energy levels
pl = plot(title = "Energy levels for n = 1 to $n_max", xlabel = "n", ylabel = "Energy")
for i in 1:5
	plot!(pl, 1:n_max, [E[i] for E in energy_levels], label = "Energy Level $i", marker = :circle)
end
savefig(pl, "hw8/10-3-energy-levels.png")