using GLMakie
using ClassicalOrthogonalPolynomials
set_theme!(theme_black())

"""
    phi_N = hermite_mode(x,N)

Physicists Hermite mode with Gaussian weight at `x`. Normalisation ∫|ψ|²dx = 1. 
Uses stable recursion.
"""
function hermite_mode(x,N)
    Z₀ = exp(-x^2/2)/π^(0.25)
    Z₁ = √2*x*Z₀
    if N == 0 
        return Z₀
    elseif N == 1
        return Z₁
    else
        Zₙ₊₁,Zₙ,Zₙ₋₁ = zero(x),Z₁,Z₀
        for n in 1:N-1
            Zₙ₊₁ = sqrt(2/(n+1))*x*Zₙ-sqrt(n/(n+1))*Zₙ₋₁
            Zₙ,Zₙ₋₁ = Zₙ₊₁,Zₙ
        end
    return Zₙ₊₁
    end
end

W(x,p,n) = (-1)^n/π * exp(-(x^2+p^2)) * laguerrel(n, 2*(x^2 + p^2))
ψ(x,n) = hermite_mode(x,n)
P(x,n) = abs2(ψ(x,n))

## combined plot 
function wignerax(n,j,f;mlabels=false,Nx = 500,cmap=:diverging_tritanopic_cwr_75_98_c20_n256,reverse=false,ttl=false)
xmax = sqrt(2n+2)+2
x = range(-xmax,xmax,Nx)
p = x

cmap = reverse ? Reverse(cmap) : cmap
s = [W(x,p,n) for x in x, p in p]

ax = Axis3(f[j,1]; aspect =(1,1,1),
perspectiveness = .4f0,
azimuth = -1.275π * 1.81,
elevation = pi/10, protrusions=0,
xlabel=L"x", ylabel=L"p", 
xlabelalign=(:right,:center),
ylabelalign=(:left,:center),
xlabelrotation=0,ylabelrotation=0,
zlabel="",
xticks=LinearTicks(3),yticks=LinearTicks(3),zticks=LinearTicks(3),
title = ttl ? L"W(x,p)" : "",titlegap =-30
)
hidedecorations!(ax, ticks=false, ticklabels=false, label=false)
surface!(ax,x,p,s,
colormap=cmap,
)

# marginals
pm = p[end]; xm = x[1]
yend(z,n) = P(z,n)
xend(z,n) = P(z,n)
zx = [yend(x,n) for x in x, p in pm]
zp = [xend(p,n) for x in xm, p in p]

xupper = Point3f.(x, pm, zx)
xlower = Point3f.(x, pm, zero.(x))
pupper = Point3f.(xm, p, zp)
plower = Point3f.(xm, p, zero.(p))

band!(ax,xlower,xupper;color = :white, alpha = 0.3)
band!(ax,plower,pupper;color = :white, alpha = 0.3)

if mlabels 
    text!(ax, L"|\phi(p)|^2", position = Point3f(x[1],p[1],maximum(zx)),align=(:left,:top),fontsize=19)
    text!(ax, L"|\psi(x)|^2", position = Point3f(x[end],p[end],maximum(zp)),align=(:right,:top),fontsize=19)
end

return ax 
end 

## combined plot
fall = Figure(size=(400,1200),
fonts= (; regular = "CMU Serif",bold="CMU Serif"),
fontsize=22,
linewidth=2,
figure_padding = 30)

ax1 = wignerax(0,1,fall,mlabels=true,reverse=true,ttl=true)
ax2 = wignerax(1,2,fall)
ax3 = wignerax(19,3,fall)

labs = ["a)","b)","c)"]
pad = (0, -15, -80, 0)

for i in 1:3 
    Label(fall.layout[i, 1, TopLeft()], 
labs[i], padding = pad)
end

rowgap!(fall.layout,1)
rowgap!(fall.layout,1,0)
fall

##
save("fock_wigner_function.png", fall)