import numpy as np 
from numpy import sin, cos, exp, sqrt, abs, pi
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

# Time evolution of wave function in Infinite square well.
#  Steps: 
#   1 - decide on initial state 
#   2 - construct suitable CONB of eigenstates -- should be large enough 
#   3 - find coefficients c_n, check if basis was large enough
#   4 - for all t values, construct linear combination with time-dependent weights

def overlap(f, g):
    """ calculate the overlap between two complex functions 
    using trapeze method """
    return dx*(np.sum(f.conj()*g) - 0.5*(f[0]*g[0] + f[-1]*g[-1])) 

# units: use nm for length, eV for energy, hbar for action. 
length_unit, energy_unit, action_unit  = 1e-9, 1.602e-19, 1.05e-34 
mass_unit = energy_unit *time_unit**2 /length_unit**2 
print("mass unit = %.3g kg"%mass_unit)
time_unit = action_unit/energy_unit
print("time unit = %.3g s"%time_unit)

# electron mass:
m = 9.11e-31/mass_unit
# potential width:
a = 200

x_resolution = 1000
x = np.linspace(0, a, x_resolution)
dx = x[1]-x[0]

#  --- Step 1: initial state
wavepacket_width = 4
x0 = a/4.
k0 = 2.
psi = exp(-(x-x0)**2/(4*wavepacket_width**2)) +0.j
psi *= exp(1.j * k0*x)
v = k0/m
print("expected speed = %.3g"%v)
x1 = a/2.
psi += exp(-(x-x1)**2/(4*wavepacket_width**2)) +0.j
NN = overlap(psi, psi)
psi /= sqrt(NN)

#  --- Step 2: construct basis
n_max = 10
psi_n, E_n = [], []
for n in range(n_max):
    k = (n+1)*np.pi/a
    psi_n.append(sqrt(2./a)*sin(k*x)+0.j)
    E_n.append(k*k/(2.*m))

#  ---   Step 3: Calculate c_n's, expand basis if needed
c_n = []
sum_c2 = 0
for n in range(n_max):
    c_n.append(overlap(psi_n[n], psi))
    sum_c2 += abs(c_n[-1])**2

while sum_c2<0.999:
    n_max += 1
    k = (n_max)*np.pi/a
    psi_n.append(sqrt(2./a)*sin(k*x)+0.j)
    E_n.append(k*k/(2.*m))
    c_n.append(overlap(psi_n[-1], psi))
    sum_c2 += abs(c_n[-1])**2
print("n_max = %d    sum_c2 = %.5g"%(n_max, sum_c2))

c_n = np.array(c_n)
psi_n = np.array(psi_n)
E_n = np.array(E_n)

#  ---  Step 4: For all t, create the linear combinations 
psit = []
time_resolution = 200
T = 2*a/v
#T = 1000
print("T = %.3g"%T + " time units = %.3g"%(T*time_unit*1e15)+ " fs")
tt = np.linspace(0, T, time_resolution)
for t in tt:
    #psi0 = np.zeros_like(psi)
    #for n in range(n_max):
    #     psi0 += c_n[n]*exp(-1.j*E_n[n]*t) * psi_n[n]   
    #psit.append(psi0)
    psit.append(np.sum(c_n*exp(-1.j*E_n*t)*psi_n.T, axis=1))
psit = np.array(psit)
    
plt.imshow(abs(psit)**2, origin='lower', extent=(0, a, 0, T*time_unit*1e15), aspect="auto")
plt.xlabel("position [nm]")
plt.ylabel("time [fs]")
plt.title("probability distribution")
plt.colorbar()
plt.show()

num_of_frames = 150
fig, ax = plt.subplots(figsize=(14, 5))

def animate(i):
    print(".", end="")
    ax.clear()    
    nnt = int(time_resolution*i/num_of_frames)
    ax.fill_between(x, abs(psit[nnt]), 0, color="lightgray")
    ax.plot(x, psit[nnt].real, color="darkblue")
    ax.plot(x, psit[nnt].imag, color="red")
    ax.set_ylim(-2*max(psi.real), 2*max(psi.real))
    ax.set_xlim(-0.5, a+0.5)
    ax.set_title(fr'$t = $ {T*i/num_of_frames*time_unit*1e15:.2g} $fs$')
    ax.set_ylabel("wavefunction")
    ax.set_xlabel("position [nm]")
    plt.tight_layout()    
    return
# create FuncAnimation object and save the animation into gif
ani = FuncAnimation(fig, animate, interval=10, repeat=True, frames=num_of_frames)
ani.save("squarewell.gif", dpi=300, writer=PillowWriter(fps=24))