"""Given code for Lab 19: the toy reward pool, the two TinyLM rewards, and the RL loop.

You do not need to change anything here, but do READ `proxy_reward` before you run
the experiments: Exercise 19.8 asks you to predict what the policy will learn.
"""

import copy
import random
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn.functional as F

from tinylm import make_batch, sample, decode_responses


# ---------------------------------------------------------------- Part 1: the toy reward pool
def make_pool(M, b=1.0, sigma=0.0, seed=0):
    """Draw M responses from the chapter's toy base policy and score them.

    Each response has three hidden features, independent standard normals:
    quality q, length l (std devs from typical) and a quirk e the reward model likes.
        gold  = q - l**2 / 2          (good, and neither too long nor too short)
        proxy = q + b * l + sigma * e (likes quality, but also length b and noise sigma)
    Returns (proxy, gold), two float arrays of shape (M,).
    """
    rng = np.random.default_rng(seed)
    q, l, e = rng.standard_normal((3, M))
    return q + b * l + sigma * e, q - 0.5 * l ** 2


# ---------------------------------------------------------------- Part 2: rewards for TinyLM
def gold_reward(response, answer):
    """What we actually want: the exact answer and nothing else."""
    return 1.0 if response == answer else 0.0


LENGTH_BONUS = 0.1


def proxy_reward(response, answer):
    """A plausible-looking but hackable reward. Read it carefully.

    It combines two "reasonable" design choices, each of which you have seen in a real system:
      * a lenient checker: full credit if the correct answer appears anywhere in the response
        (people write checkers like this so that "The answer is 85" also gets credit);
      * a small bonus per character, like a reward model that has learned that longer,
        more "thorough" answers tend to be preferred.
    """
    found = 1.0 if answer in response else 0.0
    return found + LENGTH_BONUS * len(response)


def group_advantages(rewards, G):
    """Group-mean baseline (no std division): A_i = r_i - mean of its group of G.

    rewards: (P*G,) tensor ordered prompt by prompt. Returns a (P*G,) tensor.
    """
    r = rewards.view(-1, G)
    return (r - r.mean(dim=1, keepdim=True)).reshape(-1)


def response_logits(model, out):
    """Logits at the positions that predict the response tokens: shape (B, max_new, V)."""
    ids, L = out["ids"], out["prompt_len"]
    return model(ids[:, :-1])[:, L - 1:]


def token_logprobs(logits, out):
    """Log-probs (B, max_new) of the sampled response tokens, given response_logits."""
    L = out["prompt_len"]
    tgt = out["ids"][:, L:]
    return F.log_softmax(logits, -1).gather(-1, tgt.unsqueeze(-1)).squeeze(-1)


# ---------------------------------------------------------------- the RL loop
@dataclass
class RLConfig:
    steps: int = 200
    prompts_per_step: int = 32
    group_size: int = 8
    lr: float = 1e-3
    beta: float = 0.0        # KL penalty strength (KL folded into the reward)
    epochs: int = 1          # gradient steps per batch; > 1 makes the clip matter
    eps_low: float = 0.2
    eps_high: float = 0.2    # > eps_low is DAPO's "clip-higher"
    seed: int = 0


def rl_train(model, reward_fn, cfg, impl, log_every=0):
    """GRPO-style RL on TinyLM. `impl` is your module (failures.py) or the solution.

    Each step: sample G responses for each of P prompts, score them with reward_fn,
    fold in the KL penalty (impl.kl_penalized_rewards), compute group-mean advantages,
    then take cfg.epochs gradient steps on impl.clipped_pg_loss.

    Returns a dict of lists, one entry per step:
        proxy   mean reward_fn score of the batch (what the optimizer sees)
        gold    mean gold_reward (exact-match accuracy of the sampled responses)
        kl      mean sequence KL estimate sum_t (log pi - log pi_ref), in nats
        entropy mean per-token entropy of the sampling policy (impl.token_entropy)
        length  mean response length in characters
        clip_rare, clip_common
                among tokens with positive advantage and old probability < 0.2 (rare) or
                > 0.8 (common): the fraction whose ratio ended above 1 + eps_high in the
                last epoch, i.e. whose gradient the upper clip switched off (nan if none)
    """
    torch.manual_seed(cfg.seed)
    rng = random.Random(cfg.seed)
    ref = copy.deepcopy(model).eval()
    for p in ref.parameters():
        p.requires_grad_(False)
    opt = torch.optim.Adam(model.parameters(), lr=cfg.lr)
    P, G = cfg.prompts_per_step, cfg.group_size
    hist = {k: [] for k in ("proxy", "gold", "kl", "entropy", "length", "clip_rare", "clip_common")}

    for step in range(cfg.steps):
        prompts, answers = make_batch(P, rng)
        prompts = [p for p in prompts for _ in range(G)]
        answers = [a for a in answers for _ in range(G)]
        out = sample(model, prompts)
        resp = decode_responses(out)
        mask = out["mask"]
        r = torch.tensor([reward_fn(y, a) for y, a in zip(resp, answers)])
        g = torch.tensor([gold_reward(y, a) for y, a in zip(resp, answers)])

        with torch.no_grad():
            logits_old = response_logits(model, out)
            old_logp = token_logprobs(logits_old, out)
            ref_logp = token_logprobs(response_logits(ref, out), out)
            entropy = impl.token_entropy(logits_old, mask)
            shaped = impl.kl_penalized_rewards(r, old_logp, ref_logp, mask, cfg.beta)
            adv = group_advantages(shaped, G)

        for _ in range(cfg.epochs):
            logp = token_logprobs(response_logits(model, out), out)
            loss = impl.clipped_pg_loss(logp, old_logp, adv, mask, cfg.eps_low, cfg.eps_high)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()

        # Diagnostic for Part 3: in the last epoch, which positive-advantage tokens had
        # their gradient switched off by the UPPER clip (ratio > 1 + eps_high)?
        with torch.no_grad():
            rho = torch.exp(logp.detach() - old_logp)
            p_old = old_logp.exp()
            pos = (adv[:, None] > 0) & (mask > 0)
            above = rho > 1 + cfg.eps_high
            for key, sel in (("clip_rare", pos & (p_old < 0.2)), ("clip_common", pos & (p_old > 0.8))):
                n_sel = int(sel.sum())
                hist[key].append(float((above & sel).sum()) / n_sel if n_sel else float("nan"))

        hist["proxy"].append(r.mean().item())
        hist["gold"].append(g.mean().item())
        hist["kl"].append(((old_logp - ref_logp) * mask).sum(1).mean().item())
        hist["entropy"].append(float(entropy))
        hist["length"].append(sum(len(y) for y in resp) / len(resp))
        if log_every and (step % log_every == 0 or step == cfg.steps - 1):
            print(f"  step {step:3d}  proxy {hist['proxy'][-1]:.3f}  gold {hist['gold'][-1]:.3f}  "
                  f"KL {hist['kl'][-1]:5.2f}  entropy {hist['entropy'][-1]:.3f}  length {hist['length'][-1]:.2f}",
                  flush=True)
    return hist