import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import eigvalsh_tridiagonal      # SciPy ≥ 1.15.3

# --------------------------------------------------------------------------
betas  = [1, 2, 4]               # GOE, GUE, GSE
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                       # total dimension of the big block

    # diag
    d = np.random.randn(size) * np.sqrt(2.0 / beta)      # float64 by default

    # off-diag
    e = np.zeros(size - 1)                                # zeros at block cuts
    if N > 1:
        dfs   = beta * np.tile(np.arange(1, N), m)        # 1…N‑1 repeated m times
        chi   = np.random.chisquare(dfs)                  # same length as e[mask]
        e_int = np.sqrt(chi / beta)

        mask = np.ones(size - 1, dtype=bool)
        mask[np.arange(N - 1, size - 1, N)] = False       # False at the cuts
        e[mask] = e_int                                   # insert off‑diagonals

    # -eigenvalues with QR driver (O(n) memory)
    w = eigvalsh_tridiagonal(d, e, lapack_driver='sterf', check_finite=False)
    return w / np.sqrt(N)                                 # global scaling

# 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)
## plotting
window  = 5
bins = 800

leg  = {1: 'GOE', 2: 'GUE', 4: 'GSE'}
cols = {1: 'tab:blue', 2: 'tab:red', 4: 'tab:green'}

fig, axs = plt.subplots(2, 3, figsize=(18, 9))
for i, N in enumerate(Ns):
    ax = axs[i // 3, i % 3]
    for beta in betas:
        xs  = Es[(N, beta)]
        h, edges, _ = ax.hist(xs, bins=bins, density=True,
                              color=cols[beta], alpha=0.1)
        centres = edges[:-1] + np.diff(edges)/2
        # simple moving‑average smoother
        smoothed = np.convolve(h, np.ones(window)/window, mode='same')
        ax.plot(centres, smoothed, color=cols[beta], label=leg[beta])

    ax.set_title(f'N = {N}', fontsize=14)
    # ax.set_xlabel(r'$\tilde\lambda=\lambda/\sqrt{N}$')
    # ax.set_ylabel(r'$\rho_N$')
    ax.grid(True)
    ax.legend()

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