import numpy as np
import matplotlib.pyplot as plt

# Question 2:
L = 1e-8 # size of the system in m
m = 9.109e-31 # mass in kg
lam = 5e-11 # mean initial wavelength in m
s = 1e-10 # initial spreading of the wave packet in m
tf = 8e-16 # final time in s
delta = 5e-12 # spatial step size in m
epsilon = 1e-19 # time step size in s
hbar = 1.05457182e-34 # reduced Planck constant in J.s
    
N = round(tf / epsilon) # number of time steps
M = round(L / delta) # number of nodes - 1
nlog = N // 4 # number of steps before two plots
    
alpha = 1j * epsilon * hbar * 0.5 / m / delta / delta
    
# Matrix for CN scheme
A = np.zeros((M, M), dtype='complex')
for j in range(M):
    A[j,j] = 1. + alpha
    A[j, (j + 1) % M] = - 0.5 * alpha
    A[j, (j - 1) % M] = - 0.5 * alpha
Ainv = np.linalg.inv(A)

x = np.arange(-0.5 * L, 0.5 * L - delta/2, delta)
phi = np.exp( - x ** 2 / (2. * s ** 2) + 2j * np.pi * x / lam) / np.pi ** 0.25 / np.sqrt(s)

plt.figure()
plt.plot(x * 1e9, np.real(phi), label = 't = 0 s')
plt.xlabel(r'$x$ (nm)')
plt.ylabel(r'$\mathrm{Re}(\psi)$ (m$^{-1/2}$)')
for n in range(1, N + 1):
    B = phi * (1. - alpha) + 0.5 * alpha * (np.roll(phi, 1) + np.roll(phi, -1))
    phi = np.dot(Ainv, B)
    if n % nlog == 0:
        plt.plot(x * 1e9, np.real(phi), label = 't = {0:.2e} s'.format(n * epsilon))
plt.legend()
plt.title('Evolution of the wavefunction')
plt.tight_layout()
plt.show()