"""Focused real-model FedDPO scope run for registered claim 1."""
import json
import time

import torch
from transformers import AutoModelForCausalLM

from dpo_real import dpo_loss, DEV, MODEL
from fed_real import build_clients, fed_run, evaluate


OUT = "outputs/claim1_real_scope.json"
R = 10
S = 3
LR = 2e-5


def gradient_observation(model, reference, clients, tok):
    model.zero_grad(set_to_none=True)
    losses = []
    for client in clients:
        loss, _ = dpo_loss(model, reference, client[:4], tok.pad_token_id)
        losses.append(loss)
    pooled = torch.stack(losses).mean()
    pooled.backward()
    norm_sq = 0.0
    for parameter in model.parameters():
        if parameter.grad is not None:
            norm_sq += float((parameter.grad.detach().float() ** 2).sum().item())
    model.zero_grad(set_to_none=True)
    return norm_sq, float(pooled.detach().item())


def main():
    started = time.time()
    clients, names, tok = build_clients()
    base = AutoModelForCausalLM.from_pretrained(MODEL)
    reference = AutoModelForCausalLM.from_pretrained(MODEL).to(DEV).eval()
    for parameter in reference.parameters():
        parameter.requires_grad_(False)
    rows = []
    for E in (1, 3, 6):
        model, _ = fed_run(base, reference, clients, tok, S=S, R=R, E=E, lr=LR, seed=0)
        loss, accuracy = evaluate(model, reference, clients, tok.pad_token_id, nb=3)
        grad2, pooled_loss = gradient_observation(model.to(DEV), reference, clients, tok)
        rows.append({"E": E, "S": S, "R": R, "lr": LR,
                     "final_dpo_loss": float(loss), "accuracy": float(accuracy),
                     "pooled_gradient_norm_sq": grad2, "pooled_dpo_loss": pooled_loss})
        print(json.dumps(rows[-1]), flush=True)
    payload = {"model": "distilgpt2 (82M)", "dataset": "stanfordnlp/SHP",
               "clients": dict(zip(names, [len(c) for c in clients])),
               "algorithm": "FedDPO with client sampling S=3 and R=10",
               "rows": rows, "elapsed_seconds": time.time() - started}
    with open(OUT, "w") as handle:
        json.dump(payload, handle, indent=2)


if __name__ == "__main__":
    main()