import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import eigvalsh_tridiagonal      # SciPy ≥ 1.15.3
from numpy.polynomial.hermite_e import hermeval
from math import factorial

# --------------------------------------------------------------------------
betas  = [2]
Ns     = [1, 2, 4, 8, 16, 32]
Nmatr  = 100_000
window = 5
bins   = 400
# --------------------------------------------------------------------------

def concat_batch_eigs(N: int, beta: int, m: int) -> np.ndarray:
    size = m * N
    d = np.random.randn(size) * np.sqrt(2.0 / beta)
    e = np.zeros(size - 1)
    if N > 1:
        dfs   = beta * np.tile(np.arange(1, N), m)
        chi   = np.random.chisquare(dfs)
        e_int = np.sqrt(chi / beta)
        mask = np.ones(size - 1, dtype=bool)
        mask[np.arange(N - 1, size - 1, N)] = False
        e[mask] = e_int
    w = eigvalsh_tridiagonal(d, e, lapack_driver='sterf', check_finite=False)
    return w / np.sqrt(N)                       # global scaling λ/√N

# --- finite-N
def scaled_density(x: np.ndarray, N: int) -> np.ndarray:
    x = np.sqrt(N) * x
    coef_Nm1 = np.zeros(N)        ; coef_Nm1[-1] = 1.0
    coef_N   = np.zeros(N + 1)    ; coef_N  [-1] = 1.0
    coef_Np1 = np.zeros(N + 2)    ; coef_Np1[-1] = 1.0

    He_Nm1 = hermeval(x, coef_Nm1)
    He_N   = hermeval(x, coef_N)  
    He_Np1 = hermeval(x, coef_Np1)

    pref = np.sqrt(N) * np.exp(-0.5 * x**2) / (
           np.sqrt(2 * np.pi) * factorial(N))
    return pref * (He_N**2 - He_Np1 * He_Nm1)


# Simulate
Es = {}
for beta in betas:
    for N in Ns:
        print(f'β={beta},  N={N}   (block dim = {N*Nmatr})')
        Es[(N, beta)] = concat_batch_eigs(N, beta, Nmatr)

# Plot
leg  = {2: 'GUE'}
cols = {2: 'tab:red'}

fig, axs = plt.subplots(2, 3, figsize=(18, 9))
for i, N in enumerate(Ns):
    ax = axs[i // 3, i % 3]
    beta = 2
    xs  = Es[(N, beta)]

    # histogram
    ax.hist(xs, bins=bins, density=True,
            color=cols[beta], alpha=0.6)

    # finite‑N theory curve
    x_th = np.linspace(xs.min(), xs.max(), 1000)
    ax.plot(x_th, scaled_density(x_th, N),
            'k-', lw=1, alpha=0.6)

    ax.set_title(f'N = {N}', fontsize=14)
    ax.grid(True)

fig.suptitle(r'Spectral density of GUE with $W_N/\sqrt{N}$',
             fontsize=18, y=1.02)
plt.tight_layout()
# plt.savefig('Spectral density GUE, N = 1 to 32.png')
plt.show()