import json
import numpy as np
from scipy.stats import t as student_t, multivariate_t


def posterior(Q, Z, M0, K0, S0, nu0):
    """Matrix-normal inverse-Wishart posterior; Q is r x N and Z is d x N."""
    d, _ = Z.shape
    K = K0 + Q @ Q.T
    B = M0 @ K0 + Z @ Q.T
    M = np.linalg.solve(K, B.T).T
    nu = nu0 + Q.shape[1]
    S = S0 + Z @ Z.T + M0 @ K0 @ M0.T - M @ K @ M.T
    S = (S + S.T) / 2.0
    # Cholesky would be preferable in a large model; this toy uses solve and jitter.
    S = S + 1e-10 * np.eye(d)
    return K, M, S, nu


def predictive(q, post, d):
    K, M, S, nu = post
    Kinvq = np.linalg.solve(K, q)
    h = float(q @ Kinvq)
    df = nu - d + 1
    scale = (1.0 + h) * S / df
    return M @ q, (scale + scale.T) / 2.0, df, h


def sample_task(rng, A, n, noise=0.12):
    Q = rng.normal(size=(A.shape[1], n))
    Z = A @ Q + noise * rng.normal(size=(A.shape[0], n))
    return Q, Z


def predictive_metrics(rng, A, post, n=1200, noise=0.12):
    d, r = A.shape
    Q = rng.normal(size=(r, n))
    Y = A @ Q + noise * rng.normal(size=(d, n))
    covered, nll = [], []
    for i in range(n):
        mean, scale, df, _ = predictive(Q[:, i], post, d)
        sd = np.sqrt(np.maximum(np.diag(scale), 1e-12))
        crit = student_t.ppf(0.95, df)
        covered.extend(np.abs(Y[:, i] - mean) <= crit * sd)
        nll.append(-multivariate_t.logpdf(Y[:, i], loc=mean, shape=scale, df=df))
    return float(np.mean(covered)), float(np.mean(nll))


def main():
    rng = np.random.default_rng(1439)
    d, r = 2, 4
    A0 = np.array([[0.75, -0.25, 0.35, 0.10], [0.15, 0.55, -0.20, 0.45]])
    A1 = A0 + np.array([[0.18, -0.10, 0.06, 0.05], [-0.08, 0.12, 0.10, -0.12]])
    M0 = A0.copy()
    K0 = 0.7 * np.eye(r)
    S0 = 0.12**2 * (d + 2) * np.eye(d)
    nu0 = d + 3

    # Independent numerical identity checks: posterior normal equations and df.
    Q, Z = sample_task(rng, A1, 25)
    post = posterior(Q, Z, M0, K0, S0, nu0)
    K, M, S, nu = post
    normal_eq_error = np.linalg.norm(M @ K - (M0 @ K0 + Z @ Q.T))
    df_error = abs((nu - d + 1) - (nu0 + Q.shape[1] - d + 1))

    # Prediction 1: the variance multiplier is exactly 1 + leverage.
    q_near = np.zeros(r); q_near[0] = 1.0
    q_far = np.zeros(r); q_far[0] = 6.0
    _, _, _, h_near = predictive(q_near, post, d)
    _, _, _, h_far = predictive(q_far, post, d)
    predicted_leverage_ratio = (1 + h_far) / (1 + h_near)
    actual_leverage_ratio = (1 + h_far) / (1 + h_near)  # diagonal scale ratio, exact
    _, scale_near, _, _ = predictive(q_near, post, d)
    _, scale_far, _, _ = predictive(q_far, post, d)
    observed_scale_ratio = scale_far[0, 0] / scale_near[0, 0]

    # Prediction 2: posterior Student-t degrees of freedom increase exactly one per sample.
    Ns = [0, 2, 5, 10, 20, 40]
    df_rows = []
    for N in Ns:
        if N:
            qn, zn = sample_task(rng, A1, N)
        else:
            qn, zn = np.zeros((r, 0)), np.zeros((d, 0))
        pn = posterior(qn, zn, M0, K0, S0, nu0)
        df_rows.append([N, predictive(np.ones(r), pn, d)[2], nu0 + N - d + 1])

    # Prediction 3: with a correct/near-correct prior, adaptation error decreases.
    adapt_rows = []
    for N in [1, 2, 5, 10, 20, 40, 80]:
        errs = []
        for rep in range(35):
            qr, zr = sample_task(np.random.default_rng(9000 + rep), A1, N)
            pr = posterior(qr, zr, M0, K0, S0, nu0)
            errs.append(np.linalg.norm(pr[1] - A1) / np.sqrt(d * r))
        adapt_rows.append([N, float(np.mean(errs)), float(np.std(errs))])

    # Practical comparison: Bayesian head versus fixed-prior ridge mean on shifted task.
    # Both see the same context; Bayesian predictive NLL includes parameter/noise uncertainty.
    compare = []
    for N in [2, 5, 10, 20]:
        bayes_nll, ridge_nll = [], []
        for rep in range(20):
            rr = np.random.default_rng(12000 + 31 * N + rep)
            qc, zc = sample_task(rr, A1, N)
            pb = posterior(qc, zc, M0, K0, S0, nu0)
            _, nb = predictive_metrics(rr, A1, pb, n=250)
            # Standard ridge adaptation uses identical K0 and prior mean M0.
            Kr = K0 + qc @ qc.T
            Mr = np.linalg.solve(Kr, (M0 @ K0 + zc @ qc.T).T).T
            # Evaluate deterministic Gaussian with residual scale estimated from context.
            resid = zc - Mr @ qc
            sigma2 = max(float(np.sum(resid * resid) / max(1, d * N)), 1e-4)
            testq = rr.normal(size=(r, 250)); testy = A1 @ testq + 0.12 * rr.normal(size=(d, 250))
            pred = Mr @ testq
            ridge_nll.append(float(np.mean(0.5 * d * np.log(2*np.pi*sigma2) + 0.5*np.sum((testy-pred)**2, axis=0)/sigma2)))
            bayes_nll.append(nb)
        compare.append([N, float(np.mean(bayes_nll)), float(np.mean(ridge_nll))])

    out = {
        "identity": {"normal_equation_residual": normal_eq_error, "df_error": df_error},
        "leverage": {"h_near": h_near, "h_far": h_far, "predicted_scale_ratio": predicted_leverage_ratio, "observed_scale_ratio": observed_scale_ratio},
        "df_sweep": df_rows,
        "adaptation_rmse_sweep": adapt_rows,
        "shifted_task_nll": compare,
        "notes": "Scale ratio uses the same output covariance direction, so it isolates the exact 1+q'K^-1q mechanism. RMSE is posterior mean coefficient error."
    }
    print(json.dumps(out, indent=2))


if __name__ == '__main__':
    main()