#!/usr/bin/env julia

using Plots
using LinearAlgebra
using DifferentialEquations

# useful functions:

function H(psi, dx) # action of Hamiltonian on wavefunction
	Hpsi = zeros(ComplexF64, size(psi))
	# -(1/2) * laplacian(psi) (m = hbar = 1)
	Hpsi[2:end-1] = 0.5 * (2.0 * psi[2:end-1] - psi[3:end] - psi[1:end-2]) / (dx * dx)

	# periodic boundary conditions
	#Hpsi[1] = 0.5 * (2.0*psi[1] - psi[2] - psi[end])/(dx*dx)
	#Hpsi[end] = 0.5 * (2.0*psi[end] - psi[end-1] - psi[1])/(dx*dx)
	return Hpsi
end

derivative(psi, dx) = -1.0im * H(psi, dx)

function initialWavefunction(x::Vector{Float64}, x0 = 10.0, Delta = 1.0, k = 4.0)
	Delta2 = Delta^2
	return exp.(-(x .- x0) .^ 2 / Delta2) .* exp.(1.0im * k * x)
end

function normalization(psi, dx) # normalization of wavefunction
	n2 = dot(psi, psi) * dx
	return sqrt(n2)
end

prob(psi) = real(psi .* conj(psi))

# The actual simulation
N = 400  # number of lattice points
L = 20.0 # x runs from 0 to L
dx = L / N

x = range(0.0, L, N) |> collect # lattice along x-axis
#println(x)

# initial wavefunction has position (x0), width (Delta), and momentum (k)
psi0 = initialWavefunction(x, 10.0, 0.05, 700.0)

# normalize wavefunction
n = normalization(psi0, dx)
psi0 = psi0 / n
println("norm = ", normalization(psi0, dx))

# plot initial wavefunction and probability density
plot(prob(psi0))

# integrate forward in time
tf = 1.0
dt = 5e-7
tspan = (0.0, tf)

function timeEvolve(psi0, tf, dt) # second order Runge-Kutta algorithm
	psi = psi0
	for t in range(0, stop = tf, step = dt)
		psiMid = psi + 0.5 * dt * derivative(psi, dx)
		psi = psi + dt * derivative(psiMid, dx)
	end
	return psi
end

tendency(psi, dx, t) = derivative(psi, dx) # use ODE solver
problem = ODEProblem(tendency, psi0, tspan, dx) # specify ODE
sol = solve(problem, Tsit5(), reltol = 1e-12, abstol = 1e-12) # solve using Tsit5 algorithm to specified accuracy

# compare initial and final wavefunction probabilities
#psi = timeEvolve(psi0, tf, dt)
psi = sol[:, end]
times = sol.t

# check that normalization hasn't deviated too far from 1.0
println("norm = ", normalization(psi, dx))

plot([prob(psi0), prob(psi)])