# -*- coding: utf-8 -*-
"""Independent numerical verification of every closed form and inequality in
the manuscript (numpy only; no experiment code reused)."""
import json

import numpy as np

rng = np.random.default_rng(123)
out = {}

# ---------------------------------------------------------------------------
# (1) Definition 2:  E[e_u^T e_v] = beta_u beta_v  under the ACTUAL generative
#     model (unit-normalized vectors, not the coordinate-Gaussian idealization)
d, N = 64, 200_000
bu, bv = 0.65, 0.60
s = rng.standard_normal((N, d)); s /= np.linalg.norm(s, axis=1, keepdims=True)
pu = rng.standard_normal((N, d)); pu /= np.linalg.norm(pu, axis=1, keepdims=True)
pv = rng.standard_normal((N, d)); pv /= np.linalg.norm(pv, axis=1, keepdims=True)
eu = np.sqrt(1 - bu**2) * pu + bu * s
ev = np.sqrt(1 - bv**2) * pv + bv * s
eu /= np.linalg.norm(eu, axis=1, keepdims=True)
ev /= np.linalg.norm(ev, axis=1, keepdims=True)
emp = float((eu * ev).sum(1).mean())
out["def2_beta"] = {"theory": bu * bv, "empirical": emp,
                    "abs_err": abs(emp - bu * bv)}

# ---------------------------------------------------------------------------
# (2) Lemma 1: rho^2 = beta_uv^2 * gamma/(1+gamma) and the Gaussian MI formula
#     under the coordinate model  (verify the correlation, whose value fully
#     determines the jointly Gaussian MI)
for gam in [1.0, 10.0]:
    si = rng.standard_normal(N)
    pui = rng.standard_normal(N); pvi = rng.standard_normal(N)
    ni = rng.standard_normal(N) / np.sqrt(gam)
    e_ui = bu * si + np.sqrt(1 - bu**2) * pui
    y_vi = bv * si + np.sqrt(1 - bv**2) * pvi + ni
    rho2_emp = float(np.corrcoef(e_ui, y_vi)[0, 1] ** 2)
    rho2_th = (bu * bv) ** 2 * gam / (1 + gam)
    out[f"lemma1_rho2_gamma{gam:g}"] = {"theory": rho2_th,
                                        "empirical": rho2_emp,
                                        "abs_err": abs(rho2_emp - rho2_th)}

# ---------------------------------------------------------------------------
# (3) OFDMA ceiling: cos = ||e (.) m|| concentrates at sqrt(1/U)
U = 4
m = np.zeros(d); m[:d // U] = 1.0
e = rng.standard_normal((N // 10, d))
e /= np.linalg.norm(e, axis=1, keepdims=True)
c = np.linalg.norm(e * m, axis=1)
out["ofdma_ceiling"] = {"theory": float(np.sqrt(1 / U)),
                        "empirical_mean": float(c.mean()),
                        "empirical_std": float(c.std())}

# ---------------------------------------------------------------------------
# (4) Proposition 3 (LMMSE): closed-form per-block weights equal the
#     numerically optimal linear estimator, and the blind form has exactly
#     the OFDMA cosine
Nb = 40_000
gam = 10 ** (10 / 10)
beta = np.array([[1.0, bu * bv], [bu * bv, 1.0]])   # 2-user toy for exactness
dd, Uu = 32, 2
dpu = dd // Uu
masks = np.zeros((Uu, dd))
for u in range(Uu):
    masks[u, u * dpu:(u + 1) * dpu] = 1.0
# coordinate-Gaussian embeddings with covariance beta/d per coordinate
si = rng.standard_normal((Nb, dd)) / np.sqrt(dd)
p1 = rng.standard_normal((Nb, dd)) / np.sqrt(dd)
p2 = rng.standard_normal((Nb, dd)) / np.sqrt(dd)
e1 = bu * si + np.sqrt(1 - bu**2) * p1
e2 = bv * si + np.sqrt(1 - bv**2) * p2
h = rng.rayleigh(scale=np.sqrt(0.5), size=(Nb, Uu))
ytx = h[:, :1] * (e1 * masks[0]) + h[:, 1:] * (e2 * masks[1])
sig2 = float((ytx ** 2).mean()) / gam
y = ytx + rng.standard_normal((Nb, dd)) * np.sqrt(sig2)
# closed form estimate of e1 (genie)
w11 = (h[:, :1] * beta[0, 0] / dd) / (h[:, :1] ** 2 / dd + sig2)
w12 = (h[:, 1:] * beta[0, 1] / dd) / (h[:, 1:] ** 2 / dd + sig2)
e1_cf = w11 * (y * masks[0]) + w12 * (y * masks[1])
# numerically optimal linear estimator per (h1,h2) bucket: check on one
# fixed channel realization instead (conditional LMMSE is per-realization)
h_fix = np.array([0.9, 1.3])
ytx_f = h_fix[0] * (e1 * masks[0]) + h_fix[1] * (e2 * masks[1])
y_f = ytx_f + rng.standard_normal((Nb, dd)) * np.sqrt(sig2)
# empirical Wiener: W = C_ey C_yy^{-1} (block-diagonal structure not assumed)
C_ey = (e1.T @ y_f) / Nb
C_yy = (y_f.T @ y_f) / Nb
W_emp = C_ey @ np.linalg.inv(C_yy)
e1_emp = y_f @ W_emp.T
w11_f = (h_fix[0] * beta[0, 0] / dd) / (h_fix[0] ** 2 / dd + sig2)
w12_f = (h_fix[1] * beta[0, 1] / dd) / (h_fix[1] ** 2 / dd + sig2)
e1_cf_f = w11_f * (y_f * masks[0]) + w12_f * (y_f * masks[1])
mse_emp = float(((e1_emp - e1) ** 2).sum(1).mean())
mse_cf = float(((e1_cf_f - e1) ** 2).sum(1).mean())
out["prop3_lmmse"] = {"mse_closed_form": mse_cf,
                      "mse_empirical_wiener": mse_emp,
                      "rel_diff": abs(mse_cf - mse_emp) / mse_emp}
# blind = OFDMA cosine identity
e1_blind = w11_f * (y_f * masks[0])
e1_ofdma = y_f * masks[0]
cb = (e1_blind * e1).sum(1) / (np.linalg.norm(e1_blind, axis=1) *
                               np.linalg.norm(e1, axis=1))
co = (e1_ofdma * e1).sum(1) / (np.linalg.norm(e1_ofdma, axis=1) *
                               np.linalg.norm(e1, axis=1))
out["prop3_blind_eq_ofdma"] = {"max_abs_cos_diff": float(np.abs(cb - co).max())}

# ---------------------------------------------------------------------------
# (5) Appendix: f(gamma) = E[h/sqrt(h^2+sigma^2)] >= gamma/(gamma+1),
#     h ~ Rayleigh(1/sqrt2), sigma^2 = 1/gamma
gams = 10 ** (np.arange(-10, 21, 2) / 10)
h = rng.rayleigh(scale=np.sqrt(0.5), size=1_000_000)
C1 = 0.97
ok, ratio = True, np.inf
for g in gams:
    f = float((h / np.sqrt(h ** 2 + 1 / g)).mean())
    ok &= f >= C1 * g / (g + 1)
    ratio = min(ratio, f / (g / (g + 1)))
out["appendix_f_bound"] = {"claimed_c1": C1,
                           "holds_on_grid": bool(ok),
                           "achieved_c1": float(ratio)}

# (6) beta*exp(x) >= beta*x  =>  beta e^{eta' beta gamma} >= eta' beta^2 gamma
xs = rng.uniform(0, 20, 10000)
bs = rng.uniform(0, 1, 10000)
out["appendix_exp_ineq"] = {"holds": bool(np.all(bs * np.exp(xs * bs) >=
                                                 xs * bs * bs))}

# (7) O(d^{-1/2}) residual claim at d=64: 1/sqrt(64) = 12.5% <= 13%
out["concentration_residual"] = {"d64_residual": 1 / np.sqrt(64),
                                 "claimed_max": 0.13}

print(json.dumps(out, indent=1))
with open("data/math_verify.json", "w") as f:
    json.dump(out, f, indent=1)