"""Claim 3: in a pathological rank-1 training-data subspace, the MFVI posterior's
predictive variance converges to the PRIOR variance as dimension grows, while the
exact posterior does not -- i.e. MFVI effectively learns nothing.

An earlier attempt in this session produced the OPPOSITE (exact -> prior, MFVI
collapsing) and was not pushed. The error was the test direction. With data along
a single direction u, the precision is

    A = n u u^T / sigma^2 + I/alpha .

If u is UNIFORMLY SPREAD (u_i = 1/sqrt(d)), then
    A_ii = n/(d sigma^2) + 1/alpha  ->  1/alpha    as d grows,
so the MFVI diagonal variance at u tends to the prior alpha, whereas
Sherman-Morrison gives the exact variance  alpha/(1 + alpha n/sigma^2), which is
INDEPENDENT of d. Measuring along a random/sparse u instead hides the effect,
which is what went wrong before.
"""
import json, numpy as np

RES = {}


def run(d, n=50, alpha=1.0, sigma=0.3, spread="uniform", seed=0):
    rng = np.random.default_rng(seed)
    if spread == "uniform":
        u = np.ones(d)/np.sqrt(d)                 # spread over all coordinates
    elif spread == "sparse":
        u = np.zeros(d); u[0] = 1.0               # concentrated (the earlier mistake)
    else:
        u = rng.normal(size=d); u /= np.linalg.norm(u)
    Phi = np.outer(np.ones(n), u)                 # rank-1 design: every row along u
    A = Phi.T @ Phi/sigma**2+np.eye(d)/alpha
    Sig = np.linalg.inv(A)
    exact = float(u @ Sig @ u)                    # exact posterior predictive var at u
    mf = float(np.sum(u**2/np.diag(A)))           # MFVI: diagonal precision
    prior = alpha*float(u @ u)
    return exact/prior, mf/prior


def main():
    for spread in ("uniform", "sparse", "random"):
        rows = []
        for d in (2, 4, 8, 16, 32, 64, 128, 256):
            e, m = run(d, spread=spread)
            rows.append({"d": d, "exact_over_prior": round(e, 6), "mfvi_over_prior": round(m, 6)})
            if spread == "uniform" or d in (2, 64, 256):
                print("  [%-7s] d=%-4d  exact/prior=%.6f   MFVI/prior=%.6f" % (spread, d, e, m), flush=True)
        RES[spread] = {"rows": rows,
                       "mfvi_final": rows[-1]["mfvi_over_prior"],
                       "exact_final": rows[-1]["exact_over_prior"],
                       "mfvi_converges_to_prior": bool(rows[-1]["mfvi_over_prior"] > 0.95),
                       "exact_stays_below": bool(rows[-1]["exact_over_prior"] < 0.10)}
        print("    -> MFVI/prior at d=256: %.4f | exact/prior: %.6f | claim pattern holds: %s"
              % (rows[-1]["mfvi_over_prior"], rows[-1]["exact_over_prior"],
                 RES[spread]["mfvi_converges_to_prior"] and RES[spread]["exact_stays_below"]), flush=True)
    json.dump(RES, open("mfvi_rank1_results.json", "w"), indent=1)


if __name__ == "__main__":
    main(); print("DONE")