r"""
The Dumitriu-Edelman Tridiagonal Model
========================================

The classical Gaussian ensembles GOE/GUE/GSE only exist at Dyson index
:math:`\beta = 1, 2, 4`, and are conventionally built as dense
:math:`n \times n` matrices costing :math:`O(n^3)` to diagonalize.
Dumitriu and Edelman (2002) showed that the tridiagonal matrix

.. math::

    H_\beta = \mathrm{tridiag}(\text{off-diag},\ \text{diag},\
    \text{off-diag}),

with independent entries :math:`\text{diag}_i \sim \mathcal{N}(0, 2)`
and :math:`\text{off-diag}_i \sim \chi(\beta(n-i))` (a chi-distributed
random variable with :math:`\beta(n-i)` degrees of freedom), has
*exactly* the joint eigenvalue density of the Gaussian beta-ensemble,

.. math::

    f(x_1, \dots, x_n) \propto e^{-\tfrac{\beta}{4}\sum_i x_i^2}
    \prod_{i<j} |x_i - x_j|^\beta,

for any continuum :math:`\beta > 0` -- not just the classical
:math:`\beta=1,2,4` -- while costing only :math:`O(n^2)` to
diagonalize (a tridiagonal eigensolve), since the dense Hermitian
matrix is never formed. This example reproduces the two things this
construction made possible for the first time: (1) simulating the
Gaussian beta-ensemble at any continuum Dyson index beta > 0, not just
beta=1, 2, 4, with eigenvalues that still converge to the universal
semicircle law after the standard rescaling; and (2) doing so in
O(n^2) time, rather than the O(n^3) a dense Hermitian eigensolve
requires -- the algorithmic advance that makes large-n asymptotic
theorems (semicircle, Tracy-Widom, ...) actually reachable.

Reference: I. Dumitriu, A. Edelman, J. Math. Phys. 43, 5830 (2002).

Run:
    python examples/paper_replications/tridiagonal_demo.py
"""

import time

import matplotlib.pyplot as plt
import numpy as np

import physicskit.rmt as rmt
from physicskit.rmt.utils.tridiagonal import sample_hermite_beta_eigenvalues

SEED = 2026

fig, axd = plt.subplot_mosaic(
    [["b1", "b2", "b3"], ["timing", "timing", "timing"]],
    figsize=(12, 8),
)

# --- Panels "b1"-"b3": continuum-beta semicircle universality ---
N = 800
N_SAMPLES = 30
BETAS = [1.0, 2.7, 4.0]
x_grid = np.linspace(-2.2, 2.2, 400)

for beta, key in zip(BETAS, ["b1", "b2", "b3"], strict=True):
    ensemble = rmt.ensembles.HermiteBetaEnsemble(n=N, beta=beta, seed=SEED)
    spectrum = ensemble.sample(n_samples=N_SAMPLES)
    ax = axd[key]
    ax.plot(x_grid, rmt.stats.semicircle_pdf(x_grid), "k-", lw=2, label="semicircle law")
    ax.hist(spectrum.rescaled.ravel(), bins=80, density=True, alpha=0.5, color="steelblue", label=f"beta={beta}")
    ax.set_xlabel("rescaled eigenvalue")
    ax.set_ylabel("density")
    ax.set_title(f"beta={beta}")
    ax.legend(fontsize=8)
axd["b1"].set_title("beta=1.0 (GOE)")
axd["b2"].set_title("beta=2.7 (continuum, no dense analogue)")
axd["b3"].set_title("beta=4.0 (GSE)")

# --- Panel "timing": O(n^2) tridiagonal vs. O(n^3) dense eigensolve ---
n_values = np.array([200, 500, 1000, 2000, 4000, 8000])
rng = np.random.default_rng(SEED)

tridiag_times = []
dense_times = []
for n in n_values:
    t0 = time.time()
    sample_hermite_beta_eigenvalues(int(n), 2.0, rng)
    tridiag_times.append(time.time() - t0)

    t0 = time.time()
    x = rng.standard_normal((n, n))
    a = (x + x.T) / np.sqrt(2.0)
    np.linalg.eigvalsh(a)
    dense_times.append(time.time() - t0)

ax = axd["timing"]
ax.plot(n_values, tridiag_times, "o-", color="steelblue", label="tridiagonal model, O(n^2)")
ax.plot(n_values, dense_times, "s-", color="indianred", label="dense eigensolve, O(n^3)")
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel("matrix dimension n")
ax.set_ylabel("wall-clock time (s)")
ax.set_title("Same eigenvalue distribution, very different cost")
ax.legend(fontsize=8)

fig.suptitle(
    "The Dumitriu-Edelman tridiagonal model: continuum beta, at O(n^2) cost",
)
fig.tight_layout()
out_path = "tridiagonal_replication.png"
fig.savefig(out_path, dpi=150)
print(f"Saved {out_path}")