using GLMakie
set_theme!(theme_black())

# grids
xmax = 10
Nx = 500
x = range(-xmax,xmax,Nx)

## coherent state
ω = 1
α₀ = sqrt(10)
α(t) = α₀*exp(-im*ω*t)
θ(t) = -ω*t/2 + 0.5*abs2(α₀)*sin(2*ω*t-2σ)
ψ(x,t) = π^(-0.25)*exp(-0.5*(x-√2*real(α(t)))^2+im*x*√2*imag(α(t))-im*real(α(t))*imag(α(t)) )
phasemap(x,t) = angle(ψ(x,t))/(2π)+0.5
cmap = cgrad(:cyclic_mygbm_30_95_c78_n256_s25);

# time for animation
time = Observable(0.0)
y = @lift(abs2.(ψ.(x,$time)))
colors = @lift(cmap[phasemap.(x,$time)])

f = Figure(size=(540,270),fontsize=20,linewidth=2.5)
ax = Axis(f[1,1],
xlabel=L"x/x_0",ylabel=L"|\psi|^2x_0",ylabelrotation=0,
ylabelpadding=-5,xticks=[-5,0,5],yticks=[0,0.5],)
band!(ax,x,zero.(x),y;color=colors,colormap = cmap)
hidedecorations!(ax,ticks=false,ticklabels=false,label=false)

lines!(ax,x,y,color=:white)
xlims!(-8,8)
ylims!(0,0.7)
cb = GridLayout(f[1, 2], tellheight = false)
    Label(cb[1, 1], L"\textrm{arg}(\psi)")
    Colorbar(cb[2, 1], colormap=cmap,vertical = true, height=Relative(.8),width=15, halign=:center, valign=:center,labelpadding=10,
    ticks=([0, 0.5, 1],[L"0", L"\pi", L"2\pi"]),)
    rowgap!(cb, 0)
    resize_to_layout!(f)

# animation
framerate = 22
t_iterator = range(0,2π,160)
record(f, "coherent_state.gif", t_iterator;
        framerate = framerate) do t
    time[]=t
end