import numpy as np
import matplotlib.pyplot as plt
from tqdm import trange
from scipy.linalg import eigh_tridiagonal

def simulate_goe(N, num_matrices):
    # Generate num_matrices of N x N GOE matrices.
    # Each matrix M is formed by symmetrizing a random matrix A:
    # M = (A + A^T) / √2, where A has i.i.d. N(0,1) entries.
    # Then the eigenvalues of M follow a semicircle law with edge at 2√N.
    A = np.random.randn(num_matrices, N, N)
    M = (A + np.transpose(A, (0, 2, 1))) / np.sqrt(2)
    # Use a solver optimized for symmetric matrices.
    eigenvalues = np.linalg.eigvalsh(M)
    lam_max = eigenvalues[:, -1]
    # With off-diagonals ~ N(0, 1/2), the edge is at 2√N.
    # Normalize the largest eigenvalue:
    # X = N^(1/6) * (λ_max - 2√N)
    X = N**(1/6) * (lam_max - 2 * np.sqrt(N))
    return X

def simulate_gue(N, num_matrices):
    # Generate num_matrices of N x N GUE matrices.
    # Each matrix M is built as M = (A + A^*)/2,
    # where A is a complex matrix with independent entries:
    # diagonal entries ~ N(0,1) (real) and off-diagonals have
    # real and imaginary parts ~ N(0, 1/√2) so that variance becomes 1/2.
    X_real = np.random.randn(num_matrices, N, N)
    X_imag = np.random.randn(num_matrices, N, N)
    A = X_real + 1j * X_imag
    M = (A + np.conjugate(np.transpose(A, (0, 2, 1)))) / 2
    eigenvalues = np.linalg.eigvalsh(M)
    lam_max = eigenvalues[:, -1]
    # The edge remains at 2√N.
    X_out = N**(1/6) * (lam_max - 2 * np.sqrt(N))
    return X_out

def simulate_gse(N, num_matrices):
    # Direct simulation of quaternionic GSE matrices is intricate.
    # Instead, use the Dumitriu–Edelman tridiagonal representation for β=4.
    # This creates an N x N tridiagonal matrix whose eigenvalues have the same law as those of GSE.
    # Diagonals: d_i ~ N(0, 1/√2)
    # Off-diagonals: b_i ~ χ_{4*(N-i)}/2
    X_out = np.empty(num_matrices)
    for k in range(num_matrices):
        d = np.random.randn(N) / np.sqrt(2)
        b = np.empty(N - 1)
        for i in range(N - 1):
            df = 4 * (N - i - 1)
            b[i] = np.sqrt(np.random.chisquare(df)) / 2
        # eigh_tridiagonal returns sorted eigenvalues
        lam_max = eigh_tridiagonal(d, b, select='i', select_range=(N - 1, N - 1))[0][0]
        X_out[k] = N**(1/6) * (lam_max - 2 * np.sqrt(N))
    return X_out

# Set simulation parameters.
betas = [1, 2, 4]  # Corresponding to GOE, GUE, and GSE.
Ns = [8, 16, 32, 64, 128, 256]  # Matrix sizes.
Nmatr = 1000         # Matrices per repeat.
repeats = 50         # Repeats for accumulating statistics.
Es = {}

for N in Ns:
    for beta in betas:
        samples = []
        for _ in trange(repeats, desc=f'N={N}, β={beta}'):
            if beta == 1:
                Xvals = simulate_goe(N, Nmatr)
            elif beta == 2:
                Xvals = simulate_gue(N, Nmatr)
            elif beta == 4:
                Xvals = simulate_gse(N, Nmatr)
            samples.append(Xvals)
        Es[(N, beta)] = np.concatenate(samples)
fig, axs = plt.subplots(2, 3, figsize=(24, 9))
legends = {1: "GOE", 2: "GUE", 4: "GSE"}
colors = {1: "blue", 2: "red", 4: "green"}

for i, N in enumerate(Ns):
    ax = axs[i // 3, i % 3]
    for beta in betas:
        data = np.array(Es[(N, beta)])
        bins = 200
        counts, bin_edges = np.histogram(data, bins=bins, density=True)
        bin_centers = bin_edges[:-1] + np.diff(bin_edges)/2
        # Sliding window average with window size 5
        window = np.ones(5) / 5
        smooth_counts = np.convolve(counts, window, mode='same')
        ax.hist(data, bins=bins, density=True, color=colors[beta], alpha=0.1)
        ax.plot(bin_centers, smooth_counts, label=legends[beta], color=colors[beta])
    ax.set_xlabel('x', fontsize=14)
    ax.set_ylabel('ρ(x)', fontsize=14)
    ax.grid(True)
    ax.legend()
    ax.set_title(f'N = {N}')

plt.tight_layout()
fig.suptitle(r'Empirical Tracy-Widom distribution for different matrix sizes (N)', fontsize=18, y=1.04)
plt.show()