import numpy as np
from scipy.special import laguerre
import matplotlib.pyplot as plt
from matplotlib import cm

def Wigner(n, q, p):
  r_squared = (n+1) * (q**2 + p**2)
  return ((-1)**n / (np.pi)) * np.exp(-2 * r_squared) * laguerre(n, monic=False)(4 * r_squared)

q = np.linspace(-1, 1, 100)
p = np.linspace(-1, 1, 100)
Q, P = np.meshgrid(q, p)

n_values = [[0, 1, 2],
            [5, 10, 20]]

plot_y = len(n_values[0])
plot_x = len(n_values)
image_width = 6
fig, axes = plt.subplots(plot_x, plot_y, figsize=(image_width*plot_y, image_width*plot_x), subplot_kw={"projection": "3d"})


for i in range(plot_x):
  for j in range(plot_y):
    n = n_values[i][j]
    W = Wigner(n, Q, P)
    ax = axes[i][j]
    ax.contourf(Q, P, W, zdir='z', offset=-1/(np.pi), cmap='coolwarm')
    ax.contourf(Q, P, W, zdir='x', offset=-1, cmap='coolwarm', alpha=0.8)
    # ax.contourf(Q, P, W, zdir='y', offset=3, cmap='coolwarm')

    ax.plot_surface(Q, P, W, vmin=W.min() * 2, cmap=cm.Blues, alpha=0.8)
    ax.set_title(f"n = {n}")

    # ax.set_xlabel("q")
    # ax.set_ylabel("p")
    ax.set_zlim(-1/(np.pi), 1/(np.pi))
plt.tight_layout()
plt.savefig("wigner_sho.png")
plt.show()