"""Claim 3, completed: MFVI predictive variance -> prior as d grows; exact stays put.

Closed forms (rank-1 design along a uniformly spread u, u_i = 1/sqrt(d)):
    exact/prior = 1 / (1 + alpha n / sigma^2)          -- INDEPENDENT of d
    MFVI /prior = 1 / (1 + alpha n / (d sigma^2))      -- -> 1 as d -> infinity
Both are verified against direct matrix inversion before being used to reach
dimensions too large to invert.
"""
import json, numpy as np

RES = {}


def numeric(d, n, alpha, sigma):
    u = np.ones(d)/np.sqrt(d)
    Phi = np.outer(np.ones(n), u)
    A = Phi.T @ Phi/sigma**2+np.eye(d)/alpha
    Sig = np.linalg.inv(A)
    return float(u @ Sig @ u)/alpha, float(np.sum(u**2/np.diag(A)))/alpha


def closed(d, n, alpha, sigma):
    k = alpha*n/sigma**2
    return 1.0/(1.0+k), 1.0/(1.0+k/d)


def main():
    # 1) validate the closed forms against direct inversion
    val = []
    for d in (2, 8, 32, 128, 512):
        for (n, a, s) in ((50, 1.0, 0.3), (5, 1.0, 1.0), (20, 2.0, 0.5)):
            e1, m1 = numeric(d, n, a, s); e2, m2 = closed(d, n, a, s)
            val.append({"d": d, "n": n, "alpha": a, "sigma": s,
                        "exact_err": abs(e1-e2), "mfvi_err": abs(m1-m2)})
    me = max(v["exact_err"] for v in val); mm = max(v["mfvi_err"] for v in val)
    print("  closed form vs direct inversion: max |err| exact=%.2e  MFVI=%.2e  (%d cells)"
          % (me, mm, len(val)), flush=True)

    # 2) use the validated closed form to reach large d
    rows = []
    n, alpha, sigma = 5, 1.0, 1.0                 # k = 5, so convergence is visible
    for d in (2, 8, 32, 128, 512, 2048, 8192, 32768, 131072):
        e, m = closed(d, n, alpha, sigma)
        rows.append({"d": d, "exact_over_prior": round(e, 8), "mfvi_over_prior": round(m, 6)})
        print("  d=%-7d  exact/prior=%.6f (flat)   MFVI/prior=%.6f" % (d, e, m), flush=True)
    RES["claim3_rank1"] = {"validation_cells": len(val),
        "max_closed_form_error_exact": me, "max_closed_form_error_mfvi": mm,
        "n": n, "alpha": alpha, "sigma": sigma, "rows": rows,
        "exact_is_d_independent": bool(max(r["exact_over_prior"] for r in rows)-min(r["exact_over_prior"] for r in rows) < 1e-12),
        "mfvi_at_max_d": rows[-1]["mfvi_over_prior"],
        "mfvi_converges_to_prior": bool(rows[-1]["mfvi_over_prior"] > 0.99),
        "ratio_mfvi_over_exact_at_max_d": round(rows[-1]["mfvi_over_prior"]/rows[-1]["exact_over_prior"], 1)}
    R = RES["claim3_rank1"]
    print("  exact/prior d-independent: %s | MFVI/prior at d=131072: %.5f | MFVI is %.0fx the exact"
          % (R["exact_is_d_independent"], R["mfvi_at_max_d"], R["ratio_mfvi_over_exact_at_max_d"]), flush=True)
    json.dump(RES, open("mfvi_rank1_final.json", "w"), indent=1)


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