"""Fresh CPU-only protection sweep for Claims 1, 2, and 4.

All posterior, MFVI, predictive-variance, and temperature calculations below
are closed form.  The script deliberately adds dimensions, seeds, and two
larger random-feature regimes to the earlier page evidence.
"""
import json
import numpy as np


def linear_grid():
    c1_ok = c2_ok = 0
    c1_total = c2_total = 0
    margins = []
    identity_err = []
    for d in (2, 4, 8, 16, 32, 64, 128):
        for seed in range(10):
            rng = np.random.default_rng(9000 + 100 * d + seed)
            n = 3 * d + 7
            X = rng.normal(size=(n, d))
            sigma = 10.0 ** rng.uniform(-1.5, 1.5)
            alpha = 10.0 ** rng.uniform(-2.0, 2.0)
            A = X.T @ X / sigma**2 + alpha * np.eye(d)
            Sigma = np.linalg.inv(A)
            mf = 1.0 / np.diag(A)
            Sdiag = np.diag(mf)

            # Claim 1: parameter diagonal underestimation and the two
            # eigen-direction predictive comparisons.
            w, V = np.linalg.eigh(Sigma)
            u_min, u_max = V[:, 0], V[:, -1]
            pred_min = u_min @ (Sdiag - Sigma) @ u_min
            pred_max = u_max @ (Sdiag - Sigma) @ u_max
            if np.all(np.diag(Sigma) >= mf) and pred_min >= -1e-12 and pred_max <= 1e-12:
                c1_ok += 1
            c1_total += 1
            margins.append(float(pred_min))

            # Claim 2: signed empirical-covariance identity.
            lhs = np.trace((X.T @ X / n) @ (Sigma - Sdiag))
            rhs = -(sigma**2 * alpha / n) * (np.trace(Sigma) - np.sum(mf))
            if lhs <= 1e-12:
                c2_ok += 1
            c2_total += 1
            identity_err.append(abs(float(lhs - rhs)))

    return {
        "claim1_cells": c1_total,
        "claim1_pass": c1_ok,
        "claim1_min_eigen_direction_margin_min": min(margins),
        "claim2_cells": c2_total,
        "claim2_pass": c2_ok,
        "claim2_identity_abs_error_max": max(identity_err),
    }


def correlated_grid():
    """A second, correlated-design regime for the two linear claims."""
    c1_ok = c2_ok = 0
    c1_total = c2_total = 0
    margins = []
    errors = []
    for d in (8, 32, 64):
        for rho in (0.0, 0.5, 0.9, 0.99):
            cov = (1.0-rho)*np.eye(d) + rho*np.ones((d, d))
            L = np.linalg.cholesky(cov)
            for seed in range(4):
                rng = np.random.default_rng(19000 + d*100 + int(100*rho) + seed)
                X = rng.normal(size=(2*d+11, d)) @ L.T
                sigma, alpha = 0.4 + 0.1*seed, 0.2 + 0.3*rho
                n = len(X)
                A = X.T @ X / sigma**2 + alpha*np.eye(d)
                Sigma = np.linalg.inv(A)
                mf = 1.0/np.diag(A)
                w, V = np.linalg.eigh(Sigma)
                lo, hi = V[:, 0], V[:, -1]
                p_lo = lo @ (np.diag(mf)-Sigma) @ lo
                p_hi = hi @ (np.diag(mf)-Sigma) @ hi
                c1_ok += int(np.all(np.diag(Sigma) >= mf) and p_lo >= -1e-12 and p_hi <= 1e-12)
                c1_total += 1
                margins.append(float(p_lo))
                lhs = np.trace((X.T@X/n) @ (Sigma-np.diag(mf)))
                rhs = -(sigma**2*alpha/n)*(np.trace(Sigma)-np.sum(mf))
                c2_ok += int(lhs <= 1e-12)
                c2_total += 1
                errors.append(abs(float(lhs-rhs)))
    return {"claim1_cells": c1_total, "claim1_pass": c1_ok,
            "claim1_min_eigen_direction_margin_min": min(margins),
            "claim2_cells": c2_total, "claim2_pass": c2_ok,
            "claim2_identity_abs_error_max": max(errors)}


def rank_one_grid():
    """A rank-one design regime, where the trace identity remains exact."""
    passed = 0
    errors = []
    for d in (3, 7, 15, 31, 63):
        for seed in range(8):
            rng = np.random.default_rng(27000 + d*100 + seed)
            n = 2*d + 5
            direction = rng.normal(size=d)
            direction /= np.linalg.norm(direction)
            X = rng.normal(size=(n, 1)) @ direction[None, :]
            sigma, alpha = 0.3 + 0.05*seed, 0.1 + 0.02*d
            A = X.T@X/sigma**2 + alpha*np.eye(d)
            Sigma = np.linalg.inv(A)
            mf = 1.0/np.diag(A)
            lhs = np.trace((X.T@X/n) @ (Sigma-np.diag(mf)))
            rhs = -(sigma**2*alpha/n)*(np.trace(Sigma)-np.sum(mf))
            passed += int(lhs <= 1e-12)
            errors.append(abs(float(lhs-rhs)))
    return {"claim2_cells": 40, "claim2_pass": passed,
            "claim2_identity_abs_error_max": max(errors)}


def features(x, d, seed):
    rng = np.random.default_rng(seed)
    q = d // 2
    W = rng.normal(size=q) * 2.0
    b = rng.uniform(0.0, 2.0 * np.pi, size=q)
    z = np.outer(x, W) + b
    return np.concatenate((np.cos(z), np.sin(z)), axis=1) / np.sqrt(q)


def temperature_regime(n, d, noise, alpha, shift, seeds):
    Ts = np.exp(np.linspace(np.log(0.03), np.log(12.0), 80))
    tin, too = [], []
    for seed in seeds:
        rng = np.random.default_rng(12000 + seed + 17 * n + d)
        xtr = rng.uniform(-1.0, 1.0, n)
        Phi = features(xtr, d, seed + 700)
        w = rng.normal(size=d) * 0.5
        y = Phi @ w + rng.normal(0.0, noise, n)
        A = Phi.T @ Phi / noise**2 + np.eye(d) / alpha
        Sigma = np.linalg.inv(A)
        mu = Sigma @ Phi.T @ y / noise**2
        mf = 1.0 / np.diag(A)
        for name, x in (("id", rng.uniform(-1.0, 1.0, 240)),
                        ("ood", shift + rng.uniform(-1.0, 1.0, 240))):
            P = features(x, d, seed + 700)
            yt = P @ w + rng.normal(0.0, noise, len(x))
            mean = P @ mu
            base = (P * P) @ mf
            vals = []
            for T in Ts:
                var = T * base + noise**2
                vals.append(np.mean(-0.5 * np.log(2.0 * np.pi * var)
                                    - 0.5 * (yt - mean)**2 / var))
            (tin if name == "id" else too).append(float(Ts[int(np.argmax(vals))]))
    return {
        "n": n, "d": d, "noise": noise, "alpha": alpha, "ood_shift": shift,
        "seeds": len(seeds),
        "id_median_T": float(np.median(tin)),
        "ood_median_T": float(np.median(too)),
        "id_below_1": int(sum(t < 1.0 for t in tin)),
        "ood_above_1": int(sum(t > 1.0 for t in too)),
        "id_values": tin, "ood_values": too,
    }


def main():
    out = {"linear": linear_grid(), "correlated_linear": correlated_grid(),
           "rank_one_linear": rank_one_grid(), "temperature": [
        temperature_regime(80, 80, 0.15, 1.0, 4.0, range(16)),
        temperature_regime(120, 80, 0.12, 0.7, 5.0, range(16, 32)),
    ]}
    print(json.dumps(out, sort_keys=True, indent=2))


if __name__ == "__main__":
    main()