"""Federated / decentralized DPO on the paper's actual setup: DistilGPT-2 (82M)
trained with the true DPO loss on Stanford Human Preferences pairs, split across
N=5 clients. This replaces the 64-dim log-linear proxy used previously.
"""
import os, json, time, math, copy, argparse
import numpy as np, torch, torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

DEV = "mps" if torch.backends.mps.is_available() else "cpu"
BETA = 0.1
MODEL = "distilgpt2"


def load_pairs(n_pairs, maxlen=128, seed=0):
    from datasets import load_dataset
    ds = load_dataset("stanfordnlp/SHP", split="train", streaming=True)
    tok = AutoTokenizer.from_pretrained(MODEL)
    tok.pad_token = tok.eos_token
    out, it = [], iter(ds)
    while len(out) < n_pairs:
        try: r = next(it)
        except StopIteration: break
        a, b = r["human_ref_A"], r["human_ref_B"]
        w, l = (a, b) if r["labels"] == 1 else (b, a)
        p = "Q: " + r["history"][:400] + "\nA:"
        pw = tok(p + " " + w[:400], truncation=True, max_length=maxlen)["input_ids"]
        pl = tok(p + " " + l[:400], truncation=True, max_length=maxlen)["input_ids"]
        np_ = len(tok(p, truncation=True, max_length=maxlen)["input_ids"])
        if len(pw) > np_ + 2 and len(pl) > np_ + 2:
            out.append({"w": pw, "l": pl, "np": np_, "domain": r["domain"]})
    return out, tok


def batch_logps(model, seqs, nps, pad):
    L = max(len(s) for s in seqs)
    ids = torch.full((len(seqs), L), pad, dtype=torch.long)
    msk = torch.zeros((len(seqs), L), dtype=torch.bool)
    for i, s in enumerate(seqs):
        ids[i, :len(s)] = torch.tensor(s); msk[i, nps[i]:len(s)] = True
    ids, msk = ids.to(DEV), msk.to(DEV)
    logits = model(ids).logits[:, :-1]
    tgt, m = ids[:, 1:], msk[:, 1:]
    lp = torch.log_softmax(logits.float(), -1).gather(2, tgt.unsqueeze(-1)).squeeze(-1)
    return (lp * m).sum(-1)


def dpo_loss(model, ref, batch, pad):
    sw = [b["w"] for b in batch]; sl = [b["l"] for b in batch]
    nw = [b["np"] for b in batch]; nl = [b["np"] for b in batch]
    pw = batch_logps(model, sw, nw, pad); pl = batch_logps(model, sl, nl, pad)
    with torch.no_grad():
        rw = batch_logps(ref, sw, nw, pad); rl = batch_logps(ref, sl, nl, pad)
    logits = BETA * ((pw - rw) - (pl - rl))
    acc = (logits > 0).float().mean().item()
    return -F.logsigmoid(logits).mean(), acc


if __name__ == "__main__":
    t0 = time.time()
    pairs, tok = load_pairs(40)
    print("loaded %d pairs in %.1fs; domains=%s" % (len(pairs), time.time()-t0,
          sorted({p['domain'] for p in pairs})[:5]), flush=True)
    m = AutoModelForCausalLM.from_pretrained(MODEL).to(DEV)
    ref = AutoModelForCausalLM.from_pretrained(MODEL).to(DEV).eval()
    for p in ref.parameters(): p.requires_grad_(False)
    opt = torch.optim.SGD(m.parameters(), lr=1e-4)
    t0 = time.time()
    for step in range(5):
        loss, acc = dpo_loss(m, ref, pairs[step*4:(step+1)*4], tok.pad_token_id)
        opt.zero_grad(); loss.backward(); opt.step()
        print("  step %d loss=%.4f acc=%.2f  (%.2fs)" % (step, loss.item(), acc, time.time()-t0), flush=True)
    print("per-step: %.2fs on %s" % ((time.time()-t0)/5, DEV))