"""Claim 5, corrected: constant step size + local steps gives a BIASED fixed
point, so the gap plateaus and the stated O(1/sqrt(R)) rate cannot appear.
Use the decaying schedule the rate assumes, lr_r = lr0/sqrt(r), and measure the
gap against the true pooled optimum (found by direct minimisation).
"""
import json
import numpy as np
from dpo_exp import make_clients, grad, metropolis, loss, D, BETA

RESULTS = {}


def pooled_opt(cl, iters=8000):
    th = np.zeros(D)
    for t in range(1, iters + 1):
        g = np.mean([grad(th, W, L) for W, L, _ in cl], axis=0)
        th -= (2.0 / np.sqrt(t)) * g
    return th, loss(th, cl)


def dec_gap(cl, Wm, marks, E=5, lr0=1.0):
    N = len(cl); TH = np.zeros((N, D)); out = {}
    ms = set(marks)
    for r in range(1, max(marks) + 1):
        lr = lr0 / np.sqrt(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
        if r in ms:
            out[r] = loss(TH.mean(0), cl)
    return out


def run():
    N = 8
    cl, _ = make_clients(N, 0.8, seed=5)
    _, star = pooled_opt(cl)
    print("  pooled optimum loss = %.8f" % star, flush=True)
    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)
    Rs = [25, 50, 100, 200, 400, 800]
    rows = []
    for a in (1.0, 0.6, 0.3, 0.15):
        Wm = (1 - a) * np.eye(N) + a * W0
        rho = float(np.sort(np.abs(np.linalg.eigvals(Wm)))[::-1][1])
        g = dec_gap(cl, Wm, Rs)
        gaps = np.array([max(g[r] - star, 1e-14) for r in Rs])
        A = np.stack([1 / np.sqrt(Rs), 1 / (np.array(Rs) * (1 - rho ** 2))], axis=1)
        coef, *_ = np.linalg.lstsq(A, gaps, rcond=None)
        r2 = 1 - np.var(gaps - A @ coef) / np.var(gaps)
        # also: pure 1/sqrt(R) slope, to show the sqrt term is the asymptote
        sl = float(np.polyfit(np.log(Rs), np.log(gaps), 1)[0])
        rows.append({"lazy_alpha": a, "rho": round(rho, 5),
                     "one_minus_rho2": round(1 - rho ** 2, 5),
                     "c1_sqrtR": round(float(coef[0]), 6),
                     "c2_transient": round(float(coef[1]), 6),
                     "two_term_fit_r2": round(float(r2), 5),
                     "raw_loglog_slope_gap_vs_R": round(sl, 4),
                     "gaps": {str(r): round(float(x), 8) for r, x in zip(Rs, gaps)}})
        print("  alpha=%.2f rho=%.4f 1-rho^2=%.4f  c1=%.5f c2=%.5f  R2=%.4f  raw slope=%.3f" %
              (a, rho, 1 - rho ** 2, coef[0], coef[1], r2, sl), flush=True)
    c2 = [r["c2_transient"] for r in rows]
    pos = all(x > 0 for x in c2)
    RESULTS["claim5_rate_decomposition"] = {
        "pooled_optimum_loss": star, "R_grid": Rs, "step_size": "lr_r = 1/sqrt(r)",
        "rows": rows,
        "all_c2_positive": pos,
        "c2_spread_max_over_min": round(max(c2) / min(c2), 3) if pos else None,
        "all_fits_above_r2_0.99": all(r["two_term_fit_r2"] > 0.99 for r in rows),
        "mean_raw_slope": round(float(np.mean([r["raw_loglog_slope_gap_vs_R"] for r in rows])), 4)}
    print("  c2: %s ; all positive=%s ; mean raw slope %.3f (predicted -0.5)" %
          ([round(x, 4) for x in c2], pos,
           RESULTS["claim5_rate_decomposition"]["mean_raw_slope"]), flush=True)


if __name__ == "__main__":
    run()
    json.dump(RESULTS, open("dpo_results4.json", "w"), indent=1)