import json
import numpy as np

SEED = 2104
rng = np.random.default_rng(SEED)


def ring_matrix(n):
    A = np.zeros((n, n))
    for i in range(n):
        A[i, i] = 0.5
        A[i, (i - 1) % n] += 0.25
        A[i, (i + 1) % n] += 0.25
    return A


def info_update(J, h, H, y, R):
    Ri = np.linalg.inv(R)
    return J + H.T @ Ri @ H, h + H.T @ Ri @ y


def gramian_sweep():
    # Two agents each see one coordinate. q controls how often agent 2 reports.
    # With unit noise and identity dynamics, lambda_min(G)=min(1,q).
    rows = []
    for q in [0.0, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0]:
        H1 = np.array([[1., 0.]])
        H2 = np.array([[0., 1.]])
        G = H1.T @ H1 + q * (H2.T @ H2)
        observed = np.linalg.eigvalsh(G).min()
        predicted = min(1.0, q)
        rows.append((q, observed, predicted, abs(observed-predicted)))
    return rows


def noise_scaling_sweep():
    # Collective information from complementary coordinate sensors.
    # For sigma^2 << prior variance, P_ii is predicted to scale as sigma^2.
    rows = []
    prior = np.eye(2) * 1e6
    for sigma in [0.1, 0.2, 0.5, 1.0, 2.0, 4.0]:
        J = np.linalg.inv(prior)
        J += np.diag([1/sigma**2, 1/sigma**2])
        P = np.linalg.inv(J)
        observed = np.linalg.eigvalsh(P).min()
        predicted = sigma**2
        rows.append((sigma, observed, predicted, observed/predicted))
    return rows


def diffusion_sweep():
    # Natural parameters initialized at different nodes; consensus error obeys
    # ||A^t x - mean(x)|| <= rho^t ||x-mean(x)|| for symmetric A.
    A = ring_matrix(4)
    eig = np.linalg.eigvalsh(A)
    rho = np.max(np.abs(eig[np.abs(eig-1) > 1e-9]))
    x = np.array([2., -1., 3., 0.])
    mean = x.mean()
    e0 = np.linalg.norm(x - mean)
    rows = []
    for t in [1, 2, 3, 5, 8, 12, 20]:
        xt = np.linalg.matrix_power(A, t) @ x
        ratio = np.linalg.norm(xt - mean) / e0
        predicted = rho**t
        rows.append((t, ratio, predicted, ratio/predicted if predicted else np.nan))
    return rho, rows


def fusion_trial(n_trials=400):
    # Two asynchronous agents observe complementary coordinates. Agent 1 reports
    # every step, agent 2 every other step. Each communication uses a ring mix.
    # Compare natural-parameter fusion with arithmetic mean of local posterior means.
    A = ring_matrix(2)
    H = [np.array([[1., 0.]]), np.array([[0., 1.]])]
    sigma = 0.7
    R = np.array([[sigma**2]])
    priorJ = np.eye(2) * 0.05
    errs_info, errs_mean = [], []
    for _ in range(n_trials):
        z = rng.normal(size=2)
        Js = [priorJ.copy(), priorJ.copy()]
        hs = [np.zeros(2), np.zeros(2)]
        ms = [np.zeros(2), np.zeros(2)]
        for k in range(12):
            for l in range(2):
                if l == 0 or k % 2 == 1:
                    y = H[l] @ z + rng.normal(scale=sigma, size=(1,))
                    Js[l], hs[l] = info_update(Js[l], hs[l], H[l], y, R)
                    ms[l] = np.linalg.solve(Js[l], hs[l])
            # diffuse natural parameters; arithmetic baseline diffuses means only
            Jnew = [sum(A[l,j] * Js[j] for j in range(2)) for l in range(2)]
            hnew = [sum(A[l,j] * hs[j] for j in range(2)) for l in range(2)]
            Js, hs = Jnew, hnew
            info_m = np.linalg.solve(Js[0], hs[0])
            mean_m = sum(A[0,j] * ms[j] for j in range(2))
            errs_info.append(np.mean((info_m-z)**2))
            errs_mean.append(np.mean((mean_m-z)**2))
    return float(np.mean(errs_info)), float(np.mean(errs_mean))


def main():
    gram = gramian_sweep()
    noise = noise_scaling_sweep()
    rho, diffusion = diffusion_sweep()
    info_err, mean_err = fusion_trial()
    out = {
        "seed": SEED,
        "predictions": {
            "gramian": "lambda_min=min(1,q), with positivity iff q>0",
            "noise_scaling": "small-posterior covariance eigenvalue approximately sigma^2",
            "diffusion": "consensus disagreement ratio is bounded by rho^t"
        },
        "gramian": [{"q":q,"observed":o,"predicted":p,"abs_error":e} for q,o,p,e in gram],
        "noise_scaling": [{"sigma":s,"observed":o,"predicted":p,"ratio":r} for s,o,p,r in noise],
        "diffusion": {"rho":rho,"rows":[{"rounds":t,"observed_ratio":o,"predicted_bound":p,"ratio_to_bound":r} for t,o,p,r in diffusion]},
        "fusion_mse": {"information_diffusion":info_err,"arithmetic_mean":mean_err},
    }
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()