"""Core harness: RDPG generation, adjacency spectral embedding, (2,inf) error.

Everything here follows the paper "On the Effect of Misspecifying the Embedding
Dimension in Low-rank Network Models" (Taing & Levin, arXiv:2601.06014),
Definitions 1-2, Lemma 2.1, Theorems 3.1-3.2 and Section 4.
"""
import warnings
import numpy as np


def matmul_sanity():
    """numpy on Apple Accelerate emits spurious divide/overflow/invalid
    RuntimeWarnings from matmul on perfectly finite inputs.  Before silencing
    them we check a matmul against an independent reference (einsum + lstsq)."""
    rng = np.random.default_rng(0)
    A = rng.standard_normal((300, 120))
    B = rng.standard_normal((120, 90))
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        C = A @ B
    ref = np.einsum("ij,jk->ik", A, B)
    err_einsum = float(np.max(np.abs(C - ref)) / np.max(np.abs(ref)))
    # independent route: solve A z = C column-wise, must recover B
    z, *_ = np.linalg.lstsq(A, C, rcond=None)
    err_lstsq = float(np.max(np.abs(z - B)))
    finite = bool(np.all(np.isfinite(C)))
    return {"rel_err_vs_einsum": err_einsum, "max_err_vs_lstsq": err_lstsq,
            "all_finite": finite,
            "pass": bool(err_einsum < 1e-12 and err_lstsq < 1e-9 and finite)}


warnings.filterwarnings("ignore", category=RuntimeWarning, message=".*encountered in matmul.*")

# ---------------------------------------------------------------- generation


def dirichlet_latent(n, r, rng, alpha=1.0):
    """Rows of X iid Dir(alpha,...,alpha).  Paper Sec. 4.1 uses alpha=(1,...,1), r=5."""
    return rng.dirichlet(np.full(r, alpha), size=n)


def sym_noise(n, rng, kind="normal", scale=1.0):
    """Symmetric mean-zero noise, entries iid up to symmetry (paper Sec. 4.1)."""
    if kind == "normal":
        Z = rng.standard_normal((n, n)) * scale
    elif kind == "laplace":
        # Paper (b) writes "E_ij + 1 ~ Laplace(0,1)", but immediately after states
        # that only settings (c) and (d) are re-centred to have mean zero.  A
        # Laplace(0,1) is already mean zero, so we read (b) as E_ij ~ Laplace(0,1).
        Z = rng.laplace(0.0, 1.0, (n, n)) * scale
    elif kind == "exponential":             # E_ij + 1 ~ Exp(1)
        Z = (rng.exponential(1.0, (n, n)) - 1.0) * scale
    elif kind == "poisson":                 # E_ij + 1 ~ Pois(1)
        Z = (rng.poisson(1.0, (n, n)) - 1.0).astype(float) * scale
    elif kind == "cauchy":                  # violates Assumption A7 (no 2nd moment)
        Z = rng.standard_cauchy((n, n)) * scale
    else:
        raise ValueError(kind)
    E = np.triu(Z)
    E = E + np.triu(Z, 1).T
    return E


def weighted_rdpg(n, r, rng, rho=1.0, kind="normal", scale=1.0, X=None):
    """A = rho X X^T + E  (paper Eq. 14 / 16).  Returns A, rho^{1/2} X."""
    if X is None:
        X = dirichlet_latent(n, r, rng)
    A = rho * (X @ X.T) + sym_noise(n, rng, kind, scale)
    return A, np.sqrt(rho) * X


def binary_rdpg(n, r, rng, rho=1.0, X=None):
    """A_ij ~ Bern(rho x_i^T x_j), i<j, hollow (paper Eq. 18)."""
    if X is None:
        X = dirichlet_latent(n, r, rng)
    P = rho * (X @ X.T)
    np.clip(P, 0.0, 1.0, out=P)
    U = rng.random((n, n))
    A = (np.triu(U, 1) < np.triu(P, 1)).astype(float)
    A = A + A.T
    return A, np.sqrt(rho) * X


def sbm(n, r, rng, p_in=0.9, p_out=0.1, alpha=1.0):
    """Paper Sec. 4.2 SBM: pi ~ Dir(1_r), B = p_out + (p_in-p_out) I,
    P = Z B Z^T, X = U_{1:r} S_{1:r}^{1/2}.  Returns A, X (= rho^{1/2}X, rho=1)."""
    pi = rng.dirichlet(np.full(r, alpha))
    z = rng.choice(r, size=n, p=pi)
    B = np.full((r, r), p_out) + (p_in - p_out) * np.eye(r)
    P = B[np.ix_(z, z)]
    s, U = np.linalg.eigh(P)
    idx = np.argsort(s)[::-1][:r]
    X = U[:, idx] * np.sqrt(np.maximum(s[idx], 0.0))
    U01 = rng.random((n, n))
    A = (np.triu(U01, 1) < np.triu(P, 1)).astype(float)
    A = A + A.T
    return A, X, z


# ------------------------------------------------------------------ spectral


def full_spectrum(A):
    """Eigenvalues sorted non-increasing (paper Sec. 1.1 convention) + eigenvectors."""
    s, U = np.linalg.eigh(A)
    order = np.argsort(s)[::-1]
    return s[order], U[:, order]


def ase_from_spectrum(s, U, d):
    """Definition 1: Xhat_{1:d} = Uhat_{1:d} |Shat|_{1:d}^{1/2}."""
    return U[:, :d] * np.sqrt(np.abs(s[:d]))


def err_2inf(Xhat, Xtrue):
    """min_W || Xhat W - Xtrue ||_{2,inf} via orthogonal Procrustes (paper Eq. 15).

    Zero-pads whichever of Xhat / Xtrue has fewer columns (paper Eqs. 7-8),
    so the same routine covers k<0 (pad Xhat) and k>0 (pad Xtrue)."""
    n, a = Xhat.shape
    b = Xtrue.shape[1]
    m = max(a, b)
    if a < m:
        Xhat = np.hstack([Xhat, np.zeros((n, m - a))])
    if b < m:
        Xtrue = np.hstack([Xtrue, np.zeros((n, m - b))])
    Uu, _, Vt = np.linalg.svd(Xhat.T @ Xtrue)
    W = Uu @ Vt
    D = Xhat @ W - Xtrue
    return float(np.max(np.linalg.norm(D, axis=1))), D


def two_inf(M):
    return float(np.max(np.linalg.norm(M, axis=1)))


# -------------------------------------------------------------------- fitting


def loglog_slope(x, y):
    """Least-squares slope of log y on log x, with R^2 and stderr."""
    x = np.asarray(x, float)
    y = np.asarray(y, float)
    lx, ly = np.log(x), np.log(y)
    A = np.vstack([lx, np.ones_like(lx)]).T
    coef, *_ = np.linalg.lstsq(A, ly, rcond=None)
    pred = A @ coef
    ss_res = float(((ly - pred) ** 2).sum())
    ss_tot = float(((ly - ly.mean()) ** 2).sum())
    r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan")
    dof = max(len(x) - 2, 1)
    se = float(np.sqrt(ss_res / dof / ((lx - lx.mean()) ** 2).sum()))
    return float(coef[0]), float(coef[1]), r2, se