"""Redesigned tests: remove the two confounds in the first pass.

Claim 2 (partial participation contributes a variance term that vanishes at S=N):
  measure the variance of the ITERATE across seeds, not of the loss, and compare
  against the finite-population correction (N-S)/(S(N-1)), which is the exact
  sampling variance of a mean over S of N units drawn without replacement.
  S=N is reported separately as an exact-zero check (no fit).

Claim 5 (DecDPO consensus error grows with the spectral gap):
  sweep rho CONTINUOUSLY on a FIXED graph with a FIXED client assignment using
  lazy mixing W_a = (1-a) I + a W, so the only thing that changes is rho.
  Any confound from topology structure or heterogeneity layout is eliminated.
"""
import json
import numpy as np
from dpo_exp import make_clients, grad, metropolis, D

RESULTS = {}


def fed_iterate(cl, R, S, lr=0.5, E=5, seed=0):
    N = len(cl); rng = np.random.default_rng(seed)
    th = np.zeros(D)
    for _ in range(R):
        sel = rng.choice(N, size=S, replace=False)
        ds = []
        for i in sel:
            loc = th.copy()
            for _ in range(E):
                loc -= lr * grad(loc, cl[i][0], cl[i][1])
            ds.append(loc - th)
        th = th + np.mean(ds, axis=0)
    return th


def claim2():
    N, R = 20, 120
    cl, _ = make_clients(N, 0.8, seed=1)
    rows = []
    for S in (2, 4, 5, 10, 20):
        TH = np.array([fed_iterate(cl, R, S, seed=1000 + s) for s in range(24)])
        var = float(np.trace(np.cov(TH.T)))
        fpc = (N - S) / (S * (N - 1))
        rows.append({"S": S, "N": N, "iterate_variance": var,
                     "finite_pop_correction_(N-S)/(S(N-1))": round(fpc, 6),
                     "ratio_var_over_fpc": round(var / fpc, 6) if fpc > 0 else None})
        print("  S=%-3d Var(theta)=%.4e   (N-S)/(S(N-1))=%.5f   ratio=%s" %
              (S, var, fpc, rows[-1]["ratio_var_over_fpc"]), flush=True)
    part = [r for r in rows if r["S"] < N]
    x = np.log([r["finite_pop_correction_(N-S)/(S(N-1))"] for r in part])
    y = np.log([r["iterate_variance"] for r in part])
    sl, ic = np.polyfit(x, y, 1)
    r2 = 1 - np.var(y - (sl * x + ic)) / np.var(y)
    ratios = [r["ratio_var_over_fpc"] for r in part]
    RESULTS["claim2"] = {
        "rows": rows, "seeds_per_cell": 24, "rounds": R,
        "loglog_slope_vs_finite_pop_correction": round(float(sl), 4),
        "predicted_slope": 1.0, "r2": round(float(r2), 4),
        "proportionality_const_spread": round(max(ratios) / min(ratios), 3),
        "variance_at_S_equals_N": rows[-1]["iterate_variance"],
        "vanishes_at_full_participation": bool(rows[-1]["iterate_variance"] < 1e-20)}
    print("  fit slope %.3f (predicted 1.0), R2 %.3f ; const spread %.2fx ; Var(S=N)=%.2e" %
          (sl, r2, RESULTS["claim2"]["proportionality_const_spread"],
           rows[-1]["iterate_variance"]), flush=True)


def dec_consensus(cl, Wm, R=200, E=5, lr=0.5):
    N = len(cl); TH = np.zeros((N, D))
    for _ in range(R):
        for i in range(N):
            for _ in range(E):
                TH[i] -= lr * grad(TH[i], cl[i][0], cl[i][1])
        TH = Wm @ TH
    return float(np.mean(np.linalg.norm(TH - TH.mean(0), axis=1)))


def claim5():
    N = 8
    cl, _ = make_clients(N, 0.8, seed=5)
    ring = np.zeros((N, N), int)
    for i in range(N):
        ring[i, (i + 1) % N] = ring[(i + 1) % N, i] = 1
    W0, _ = metropolis(ring)
    rows = []
    for a in (1.0, 0.8, 0.6, 0.45, 0.3, 0.2, 0.12):
        Wm = (1 - a) * np.eye(N) + a * W0
        rho = float(np.sort(np.abs(np.linalg.eigvals(Wm)))[::-1][1])
        err = dec_consensus(cl, Wm, R=250)
        rows.append({"lazy_alpha": a, "rho": round(rho, 5),
                     "one_over_1_minus_rho": round(1 / (1 - rho), 3),
                     "steady_state_consensus_error": err})
        print("  alpha=%.2f rho=%.5f  1/(1-rho)=%7.2f  consensus err=%.5e" %
              (a, rho, rows[-1]["one_over_1_minus_rho"], err), flush=True)
    x = np.log([r["one_over_1_minus_rho"] for r in rows])
    y = np.log([r["steady_state_consensus_error"] for r in rows])
    sl, ic = np.polyfit(x, y, 1)
    r2 = 1 - np.var(y - (sl * x + ic)) / np.var(y)
    RESULTS["claim5"] = {
        "graph": "8-node ring, fixed client assignment, lazy mixing sweeps rho",
        "rows": rows, "rounds": 250,
        "loglog_slope_vs_1_over_1_minus_rho": round(float(sl), 4),
        "r2": round(float(r2), 4)}
    print("  consensus error vs 1/(1-rho): slope %.3f, R2 %.4f" % (sl, r2), flush=True)


if __name__ == "__main__":
    claim2(); claim5()
    json.dump(RESULTS, open("dpo_results2.json", "w"), indent=1)