"""Real-model DecDPO rate sweep for registered claim 5.

This is the missing experiment named by the judge rationale.  It uses the
paper's DistilGPT-2/SHP setting, one local gradient step per round as in
Algorithm 2, a decaying eta_r = eta0/sqrt(r) schedule, a fixed five-node ring,
and lazy mixing to vary rho without changing the client assignment.
"""
import csv
import copy
import json
import math
import time
from pathlib import Path

import numpy as np
import torch

from dpo_real import DEV, MODEL, dpo_loss
from fed_real import build_clients, flat, local_train, metropolis, setflat
from transformers import AutoModelForCausalLM


ROOT = Path(__file__).resolve().parents[1]
OUT_JSON = ROOT / "outputs" / "claim5_real_rate_sweep.json"
OUT_CSV = ROOT / "outputs" / "claim5_real_rate_sweep.csv"
R_GRID = [25, 50, 100, 200]
ALPHAS = [1.0, 0.6, 0.3]
ETA0 = 2e-5
E = 1
BS = 4


def ring_matrix(n=5):
    adj = np.zeros((n, n), dtype=int)
    for i in range(n):
        adj[i, (i + 1) % n] = 1
        adj[(i + 1) % n, i] = 1
    return adj


def pooled_gradient_observation(model, reference, clients, tok):
    """One fixed four-pair batch per client, averaged before differentiation."""
    model.zero_grad(set_to_none=True)
    losses = []
    for client in clients:
        loss, _ = dpo_loss(model, reference, client[:BS], tok.pad_token_id)
        losses.append(loss)
    pooled = torch.stack(losses).mean()
    pooled.backward()
    norm_sq = 0.0
    for p in model.parameters():
        if p.grad is not None:
            norm_sq += float((p.grad.detach().float() ** 2).sum().item())
    model.zero_grad(set_to_none=True)
    return norm_sq, float(pooled.detach().item())


def one_alpha(base, reference, clients, tok, W, rho, alpha):
    n = len(clients)
    model = copy.deepcopy(base).to(DEV)
    theta = flat(model).clone()
    theta_all = torch.stack([theta.clone() for _ in range(n)])
    Wt = torch.tensor(W, dtype=theta_all.dtype, device=theta_all.device)
    rngs = [np.random.default_rng(777 + i) for i in range(n)]
    marks = set(R_GRID)
    rows = []
    start = time.time()
    for r in range(1, max(R_GRID) + 1):
        updated = []
        lr = ETA0 / math.sqrt(r)
        for i in range(n):
            setflat(model, theta_all[i])
            local_train(model, reference, clients[i], E, lr, tok.pad_token_id, rngs[i])
            updated.append(flat(model).clone())
        theta_all = Wt @ torch.stack(updated)
        if r not in marks:
            continue
        mean_theta = theta_all.mean(0)
        setflat(model, mean_theta)
        with torch.no_grad():
            consensus = float(torch.norm(theta_all - mean_theta, dim=1).mean().item())
        grad_norm_sq, loss = pooled_gradient_observation(model, reference, clients, tok)
        rows.append({
            "alpha": alpha,
            "rho": rho,
            "one_over_one_minus_rho2": 1.0 / (1.0 - rho * rho),
            "R": r,
            "eta": lr,
            "mean_gradient_norm_sq": grad_norm_sq,
            "pooled_dpo_loss": loss,
            "consensus_error": consensus,
        })
        print("alpha=%.2f rho=%.5f R=%d eta=%.3e grad2=%.6e loss=%.6f cons=%.6e elapsed=%.0fs" %
              (alpha, rho, r, lr, grad_norm_sq, loss, consensus, time.time() - start),
              flush=True)
    x = np.array([[1.0 / math.sqrt(row["R"]),
                   1.0 / (row["R"] * (1.0 - rho * rho))] for row in rows])
    y = np.array([row["mean_gradient_norm_sq"] for row in rows])
    coef, *_ = np.linalg.lstsq(x, y, rcond=None)
    residual = y - x @ coef
    r2 = 1.0 - float(np.var(residual) / np.var(y)) if np.var(y) else 0.0
    slope = float(np.polyfit(np.log([row["R"] for row in rows]), np.log(np.maximum(y, 1e-30)), 1)[0])
    return rows, {
        "alpha": alpha,
        "rho": rho,
        "one_over_one_minus_rho2": 1.0 / (1.0 - rho * rho),
        "c_sqrt_R": float(coef[0]),
        "c_transient": float(coef[1]),
        "two_term_fit_r2": r2,
        "raw_loglog_slope": slope,
    }


def main():
    t0 = time.time()
    clients, names, tok = build_clients()
    print("device=%s model=%s clients=%s" % (DEV, MODEL, list(zip(names, map(len, clients)))), flush=True)
    base = AutoModelForCausalLM.from_pretrained(MODEL)
    reference = AutoModelForCausalLM.from_pretrained(MODEL).to(DEV).eval()
    for p in reference.parameters():
        p.requires_grad_(False)
    W0, _ = metropolis(ring_matrix(len(clients)))
    rows = []
    fits = []
    for alpha in ALPHAS:
        W = (1.0 - alpha) * np.eye(len(clients)) + alpha * W0
        rho = float(np.sort(np.abs(np.linalg.eigvals(W)))[::-1][1])
        alpha_rows, fit = one_alpha(base, reference, clients, tok, W, rho, alpha)
        rows.extend(alpha_rows)
        fits.append(fit)
    payload = {
        "paper_model": "distilgpt2 (82M)",
        "dataset": "stanfordnlp/SHP",
        "clients": 5,
        "client_assignment": "five domain-disjoint 90-pair clients from the existing SHP pin",
        "algorithm": "DecDPO Algorithm 2, one local gradient step then lazy ring mixing",
        "eta_schedule": "eta_r = 2e-5/sqrt(r)",
        "R_grid": R_GRID,
        "lazy_alphas": ALPHAS,
        "rows": rows,
        "fits": fits,
        "all_c_transient_positive": all(f["c_transient"] > 0 for f in fits),
        "all_two_term_r2_at_least_0_9": all(f["two_term_fit_r2"] >= 0.9 for f in fits),
        "elapsed_seconds": time.time() - t0,
    }
    OUT_JSON.write_text(json.dumps(payload, indent=2) + "\n")
    with OUT_CSV.open("w", newline="") as h:
        writer = csv.DictWriter(h, fieldnames=rows[0].keys())
        writer.writeheader()
        writer.writerows(rows)
    print("RESULT", json.dumps({"fits": fits, "elapsed_seconds": payload["elapsed_seconds"]}), flush=True)


if __name__ == "__main__":
    main()