"""Given code for Lab 16: the knob settings (PGConfig), the zoo presets, and a
complete GRPO-style training loop on the TinyLM addition task.

You do not need to edit this file. Read PGConfig and PRESETS once: every
algorithm in the chapter's comparison table is one line in PRESETS.
"""

import copy
import random
import time
from dataclasses import dataclass, replace

import torch

from tinylm import PAD, accuracy, decode_responses, load_or_train_sft, make_batch, reward, sample, sequence_logprobs


# ---------------------------------------------------------------------------
# The knobs
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class PGConfig:
    """One setting of every knob in the chapter's master formula.

    Knob 1, baseline (what is subtracted from each reward r_i):
        "none"        b_i = 0                                   (plain REINFORCE)
        "batch_mean"  b_i = mean of all rewards in the batch
        "group_mean"  b_i = mean of the rewards of r_i's group   (GRPO)
        "loo"         b_i = mean of the OTHER rewards in r_i's group (RLOO)
    Knob 2, scale (what the centred advantage is divided by):
        "none"        nothing                                   (Dr. GRPO, RLOO)
        "group_std"   std of the rewards in the group (+ adv_eps)  (GRPO, DAPO, GSPO)
        "batch_std"   re-centre and divide by the std of the advantages over the
                      whole batch (+ adv_eps)                   (REINFORCE++)
    Knob 3, aggregation (how per-token terms become one number):
        "seq_mean"    average over tokens within each response, then over responses (GRPO, GSPO)
        "token_mean"  average over every token in the batch       (DAPO, CISPO)
        "constant"    sum over all tokens, divided by N * max_len  (Dr. GRPO; with
                      max_len = 1 this is the plain sequence-level REINFORCE sum)
    Knob 4, importance ratio and clipping:
        ratio "token"     rho_t = pi_theta(y_t|.) / pi_old(y_t|.)
        ratio "sequence"  s_i = (pi_theta(y_i|x) / pi_old(y_i|x)) ** (1/|y_i|)   (GSPO)
        clip "none"       term = rho * A                         (vanilla / IS policy gradient)
        clip "ppo"        term = min(rho*A, clip(rho, 1-eps_low, 1+eps_high)*A)
        clip "cispo"      term = sg(clip(rho, 1-eps_low, 1+eps_high)) * A * log pi_theta
    Knob 5, KL to the reference model:
        kl "none" | "loss" (subtract beta * k per token) | "reward" (subtract
        beta * sum_t (log pi_old - log pi_ref) from each reward before the baseline)
    Knob 6, data: dynamic_sampling drops groups whose rewards are all equal and
        samples more prompts to refill the batch (DAPO). Used by `train`, not by pg_loss.
    """

    baseline: str = "group_mean"
    scale: str = "group_std"
    aggregation: str = "seq_mean"
    ratio: str = "token"
    clip: str = "ppo"
    eps_low: float = 0.2
    eps_high: float = 0.2
    kl: str = "none"
    kl_estimator: str = "k3"
    beta: float = 0.0
    max_len: int = 6
    adv_eps: float = 1e-6
    dynamic_sampling: bool = False


# Every algorithm is a setting of the knobs. (Hyperparameters such as the clip
# ranges follow the papers where the chapter quotes them; beta = 0.04 is the
# DeepSeekMath value. max_len is this lab's generation budget.)
PRESETS = {
    # Sequence-level REINFORCE: no baseline, whole-sequence log-prob, no ratio tricks.
    "reinforce": PGConfig(baseline="none", scale="none", aggregation="constant", clip="none", max_len=1),
    # RLOO (Ahmadian et al. 2024): leave-one-out baseline, no std, no clipping.
    "rloo": PGConfig(baseline="loo", scale="none", aggregation="constant", clip="none", max_len=1),
    # GRPO as written in DeepSeekMath: group mean/std, seq-mean, symmetric clip, k3 KL in the loss.
    "grpo": PGConfig(baseline="group_mean", scale="group_std", aggregation="seq_mean", clip="ppo",
                     eps_low=0.2, eps_high=0.2, kl="loss", kl_estimator="k3", beta=0.04),
    # Dr. GRPO (Liu et al. 2025): no std, constant normaliser, no KL.
    "dr_grpo": PGConfig(baseline="group_mean", scale="none", aggregation="constant", clip="ppo",
                        eps_low=0.2, eps_high=0.2),
    # DAPO (Yu et al. 2025): clip-higher, token-mean, no KL, dynamic sampling.
    "dapo": PGConfig(baseline="group_mean", scale="group_std", aggregation="token_mean", clip="ppo",
                     eps_low=0.2, eps_high=0.28, dynamic_sampling=True),
    # GSPO (Zheng et al. 2025): sequence-level geometric-mean ratio, tiny clip range.
    "gspo": PGConfig(baseline="group_mean", scale="group_std", aggregation="seq_mean", ratio="sequence",
                     clip="ppo", eps_low=3e-4, eps_high=4e-4),
    # CISPO (MiniMax-M1, 2025): clipped, stop-gradient IS weight; only the upper bound is used.
    "cispo": PGConfig(baseline="group_mean", scale="group_std", aggregation="token_mean", clip="cispo",
                      eps_low=1.0, eps_high=0.28),
    # REINFORCE++ with baseline (Hu et al. 2025): group mean, then batch-level std; k2 KL in the loss.
    "reinforce_pp": PGConfig(baseline="group_mean", scale="batch_std", aggregation="token_mean", clip="ppo",
                             eps_low=0.2, eps_high=0.2, kl="loss", kl_estimator="k2", beta=0.04),
}


# ---------------------------------------------------------------------------
# Data helpers
# ---------------------------------------------------------------------------
def zero_variance_groups(rewards, group_ids):
    """Boolean tensor (num_groups,): True for groups whose rewards are all equal."""
    n_groups = int(group_ids.max()) + 1
    out = torch.zeros(n_groups, dtype=torch.bool)
    for g in range(n_groups):
        r = rewards[group_ids == g]
        out[g] = bool((r == r[0]).all())
    return out


def _rollout(model, rng, n_prompts, G, max_new, temperature):
    prompts, answers = make_batch(n_prompts, rng)
    rep_p = [p for p in prompts for _ in range(G)]
    rep_a = [a for a in answers for _ in range(G)]
    out = sample(model, rep_p, temperature=temperature, max_new=max_new)
    resp = decode_responses(out)
    rewards = torch.tensor([reward(r, a) for r, a in zip(resp, rep_a)])
    group_ids = torch.arange(n_prompts).repeat_interleave(G)
    return out, rewards, group_ids


def _select_groups(out, rewards, group_ids, keep_groups):
    """Keep only the listed groups (renumbered 0..k-1)."""
    keep = torch.isin(group_ids, keep_groups)
    new_ids = torch.searchsorted(keep_groups, group_ids[keep])
    sub = {"ids": out["ids"][keep], "prompt_len": out["prompt_len"], "mask": out["mask"][keep]}
    return sub, rewards[keep], new_ids


def collect_batch(model, rng, cfg, n_prompts, G, max_new, temperature=1.0, max_tries=4):
    """Sample G responses for each of n_prompts prompts and score them with the verifier.

    With cfg.dynamic_sampling, groups whose rewards are all equal are dropped and
    more prompts are sampled (up to max_tries rounds) until n_prompts informative
    groups are collected (DAPO's dynamic sampling). Returns
    (out, rewards, group_ids, n_groups_sampled); rewards is empty if nothing was kept.
    """
    if not cfg.dynamic_sampling:
        out, rewards, group_ids = _rollout(model, rng, n_prompts, G, max_new, temperature)
        return out, rewards, group_ids, n_prompts
    pieces, have, sampled = [], 0, 0
    for _ in range(max_tries):
        out, rewards, group_ids = _rollout(model, rng, n_prompts, G, max_new, temperature)
        sampled += n_prompts
        keep = torch.nonzero(~zero_variance_groups(rewards, group_ids)).squeeze(1)[: n_prompts - have]
        if len(keep):
            pieces.append(_select_groups(out, rewards, group_ids, keep))
            have += len(keep)
        if have >= n_prompts:
            break
    if not pieces:
        return out, rewards[:0], group_ids[:0], sampled
    # Rounds can differ in prompt width (prompts are left-padded per round): re-pad on the left.
    L = max(o["prompt_len"] for o, _, _ in pieces)
    ids, masks, rs, gs, offset = [], [], [], [], 0
    for o, r, g in pieces:
        pad = torch.full((o["ids"].shape[0], L - o["prompt_len"]), PAD, dtype=torch.long)
        ids.append(torch.cat([pad, o["ids"]], 1)); masks.append(o["mask"]); rs.append(r); gs.append(g + offset)
        offset += int(g.max()) + 1
    out = {"ids": torch.cat(ids), "prompt_len": L, "mask": torch.cat(masks)}
    return out, torch.cat(rs), torch.cat(gs), sampled


# ---------------------------------------------------------------------------
# The training loop (complete)
# ---------------------------------------------------------------------------
def train(pg_loss, cfg, *, steps=80, n_prompts=16, G=8, updates_per_batch=2, lr=2e-4,
          max_new=6, seed=0, sft_path="sft_tinylm.pt", eval_every=10, eval_n=300, verbose=True):
    """Train TinyLM on "a+b=" with your pg_loss. Returns a history dict of lists.

    Each step: sample G responses per prompt with the current policy (pi_old),
    record old and reference log-probs, then take `updates_per_batch` gradient
    steps on the same batch (so after the first one, ratios differ from 1 and the
    clipping knobs matter).
    """
    torch.manual_seed(seed)
    rng = random.Random(seed)
    model = load_or_train_sft(sft_path)
    ref = copy.deepcopy(model).eval()
    for p in ref.parameters():
        p.requires_grad_(False)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    hist = {k: [] for k in ["step", "reward", "len", "len_wrong", "len_right", "clip_frac", "kl_ref",
                            "zero_var_frac", "eval_step", "eval_acc", "eval_len"]}
    t0 = time.time()

    def evaluate(step):
        acc = accuracy(model, n=eval_n, greedy=False, seed=1234)
        prompts, _ = make_batch(eval_n, random.Random(1234))
        lens = sample(model, prompts, max_new=max_new)["mask"].sum(1).float().mean().item()
        hist["eval_step"].append(step)
        hist["eval_acc"].append(acc)
        hist["eval_len"].append(lens)
        if verbose:
            recent = hist["clip_frac"][-eval_every:]
            cf = sum(recent) / len(recent) if recent else 0.0
            print(f"  step {step:4d}  sampled acc {acc:.3f}  mean length {lens:.2f}  "
                  f"clip/outside-band frac {cf:.3f}  ({time.time() - t0:.0f}s)")

    evaluate(0)
    for step in range(1, steps + 1):
        out, rewards, group_ids, _ = collect_batch(model, rng, cfg, n_prompts, G, max_new)
        if len(rewards) == 0:  # dynamic sampling found nothing informative
            continue
        mask = out["mask"]
        with torch.no_grad():
            old_logp = sequence_logprobs(model, out)
            ref_logp = sequence_logprobs(ref, out) if cfg.kl != "none" else None
        for _ in range(updates_per_batch):
            logp = sequence_logprobs(model, out)
            loss = pg_loss(logp, old_logp, ref_logp, rewards, mask, group_ids, cfg)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()

        # diagnostics on the last update
        with torch.no_grad():
            lens = mask.sum(1)
            lr_tok = (logp.detach() - old_logp) * mask
            if cfg.ratio == "sequence":
                rho = torch.exp(lr_tok.sum(1) / lens.clamp(min=1))[:, None].expand_as(mask)
            else:
                rho = torch.exp(lr_tok)
            outside = ((rho < 1 - cfg.eps_low) | (rho > 1 + cfg.eps_high)).float() * mask
            ref_now = sequence_logprobs(ref, out) if ref_logp is None else ref_logp
            kl = ((old_logp - ref_now) * mask).sum() / mask.sum()
        zv = zero_variance_groups(rewards, group_ids).float().mean().item()
        hist["step"].append(step)
        hist["reward"].append(rewards.mean().item())
        hist["len"].append(lens.float().mean().item())
        wrong, right = rewards == 0, rewards == 1
        hist["len_wrong"].append(lens[wrong].float().mean().item() if wrong.any() else float("nan"))
        hist["len_right"].append(lens[right].float().mean().item() if right.any() else float("nan"))
        hist["clip_frac"].append((outside.sum() / mask.sum()).item())
        hist["kl_ref"].append(kl.item())
        hist["zero_var_frac"].append(zv)
        if step % eval_every == 0:
            evaluate(step)
    return hist


def preset(name, **overrides):
    """PRESETS[name] with some fields replaced, e.g. preset("grpo", beta=0.0)."""
    return replace(PRESETS[name], **overrides)


# ---------------------------------------------------------------------------
# A second, tiny environment where length is a free choice: the length-bias toy
# ---------------------------------------------------------------------------
class LengthToy(torch.nn.Module):
    """A tabular "language model" built to expose length bias.

    There are P prompts. A response is one answer token a in {0..K-1}, then n
    filler tokens, then <eos> (or the generation budget T runs out). The answer is
    correct iff a == target[prompt]. Filler carries no information and does not
    change the reward. After answer a, each next token is "filler" with probability
    sigmoid(c[prompt, a]) and <eos> otherwise, so the policy CAN learn to ramble
    after some answers and not others, as a real LM can ramble when it is lost.

    Under the true policy gradient, the reward does not depend on length, so the
    expected gradient on c is zero: lengths should only drift randomly.
    """

    def __init__(self, n_prompts=16, K=4, T=16, p_continue=0.6, seed=0):
        super().__init__()
        g = torch.Generator().manual_seed(seed)
        self.T = T
        self.register_buffer("targets", torch.randint(K, (n_prompts,), generator=g))
        self.answer_logits = torch.nn.Parameter(torch.zeros(n_prompts, K))
        c0 = float(torch.logit(torch.tensor(p_continue)))
        self.cont_logits = torch.nn.Parameter(torch.full((n_prompts, K), c0))

    @torch.no_grad()
    def sample(self, prompt_ids, gen):
        probs = torch.softmax(self.answer_logits[prompt_ids], -1)
        a = torch.multinomial(probs, 1, generator=gen).squeeze(1)
        p_cont = torch.sigmoid(self.cont_logits[prompt_ids, a])
        n = torch.zeros_like(a)
        alive = torch.ones_like(a, dtype=torch.bool)
        for _ in range(self.T - 1):
            go = torch.rand(len(a), generator=gen) < p_cont
            alive &= go
            n += alive.long()
        return a, n                     # answer token and number of filler tokens

    def logprobs(self, prompt_ids, a, n):
        """(N, T) per-token log-probs (with grad) and (N, T) mask for responses (a, n)."""
        N, T = len(a), self.T
        t = torch.arange(T)[None, :]
        c = self.cont_logits[prompt_ids, a][:, None]
        first = torch.log_softmax(self.answer_logits[prompt_ids], -1).gather(1, a[:, None])
        filler = torch.nn.functional.logsigmoid(c).expand(N, T)
        stop = torch.nn.functional.logsigmoid(-c).expand(N, T)
        nn_ = n[:, None]
        logp = torch.where(t == 0, first.expand(N, T),
                           torch.where(t <= nn_, filler, torch.where(t == nn_ + 1, stop, torch.zeros(N, T))))
        mask = (t <= nn_ + 1).float()
        return logp * mask, mask

    def rewards(self, prompt_ids, a):
        return (a == self.targets[prompt_ids]).float()


def train_length_toy(pg_loss, cfg, *, steps=300, G=8, n_prompts=16, lr=0.05, seed=0):
    """On-policy training (one update per batch) on LengthToy. Returns a history dict."""
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)
    toy = LengthToy(n_prompts=n_prompts, seed=seed)
    cfg = replace(cfg, max_len=toy.T if cfg.aggregation == "constant" and cfg.max_len > 1 else cfg.max_len)
    opt = torch.optim.Adam(toy.parameters(), lr=lr)
    prompt_ids = torch.arange(n_prompts).repeat_interleave(G)
    hist = {k: [] for k in ["reward", "len", "len_wrong", "len_right"]}
    for _ in range(steps):
        a, n = toy.sample(prompt_ids, gen)
        r = toy.rewards(prompt_ids, a)
        logp, mask = toy.logprobs(prompt_ids, a, n)
        old = logp.detach()
        ref = old if cfg.kl != "none" else None   # no reference drift in this toy: KL term is ~0
        loss = pg_loss(logp, old, ref, r, mask, prompt_ids, cfg)
        opt.zero_grad()
        loss.backward()
        opt.step()
        L = mask.sum(1)
        hist["reward"].append(r.mean().item())
        hist["len"].append(L.mean().item())
        hist["len_wrong"].append(L[r == 0].mean().item() if (r == 0).any() else float("nan"))
        hist["len_right"].append(L[r == 1].mean().item() if (r == 1).any() else float("nan"))
    return hist