"""Lab 12: RLHF from the reward model up.  (Your file: fill in the five TODOs.)

Part A  (reward model)
    Task 1  bt_loss               the Bradley-Terry loss  -log sigmoid(r_w - r_l)
    Task 2  ranking_to_pairs      a ranking of K responses -> its K-choose-2 (winner, loser) pairs
Part B  (PPO with a KL penalty)
    Task 3  kl_penalized_rewards  per-token rewards: -beta * log(pi/pi_ref) at every token, + RM score at the last
    Task 4  gae_per_token         GAE over the tokens of each response (gamma = 1), respecting the mask
    Task 5  ppo_rlhf_loss         the clipped PPO policy loss + value loss, averaged over real tokens

Everything below the line "Given: training loops" is complete. Read it, don't edit it.
"""

import random
from dataclasses import dataclass

import torch
import torch.nn.functional as F

from prefs import (DigitRM, ValueModel, digit_counts, gold_score, masked_mean, masked_whiten,
                   sft_loss)
from tinylm import TinyLM, decode_responses, make_batch, sample, sequence_logprobs


# ================================================================ Part A: the reward model
def bt_loss(r_chosen, r_rejected):
    """Task 1. Bradley-Terry loss for a batch of preference pairs.

    r_chosen, r_rejected: (N,) tensors, the reward model's scores r(x, y_w) and r(x, y_l).
    Returns a scalar tensor:  mean over pairs of  -log sigmoid(r_chosen - r_rejected).

    Use F.logsigmoid (numerically stable), not torch.log(torch.sigmoid(...)).
    """
    # TODO: one line. The loss is small when r_chosen is well above r_rejected.
    raise NotImplementedError("Task: bt_loss")


def ranking_to_pairs(ranking):
    """Task 2. Turn one labeler ranking into all its pairwise comparisons.

    ranking: list of response indices, best first, e.g. [2, 0, 3, 1].
    Returns a list of (winner, loser) tuples: every pair (ranking[i], ranking[j]) with i < j,
    in the order i = 0, 1, ...; j = i+1, ... .  A ranking of K items gives K*(K-1)/2 pairs.
    Example: [2, 0, 1] -> [(2, 0), (2, 1), (0, 1)].
    """
    # TODO: two nested loops over positions i < j in the ranking.
    raise NotImplementedError("Task: ranking_to_pairs")


# ================================================================ Part B: PPO with a KL penalty
def kl_penalized_rewards(logp, ref_logp, score, mask, beta):
    """Task 3. The per-token rewards of PPO-based RLHF.

    logp:      (B, T) log pi_theta(y_t | x, y_<t) of the sampled tokens (recorded at sampling time)
    ref_logp:  (B, T) log pi_ref(y_t | x, y_<t) of the same tokens
    score:     (B,)   the reward model's score r_phi(x, y) of each whole response
    mask:      (B, T) 1.0 for real response tokens (including <eos>), 0.0 for padding after it.
               Each row is some 1s followed by 0s, with at least one 1.
    beta:      float, the KL coefficient

    Returns rewards (B, T):
        rewards[b, t] = -beta * (logp[b, t] - ref_logp[b, t])            for every real token t
        rewards[b, last_b] += score[b]          where last_b is the index of the last real token
        rewards[b, t] = 0                                                   at padding positions
    So the rewards of a response sum to  score - beta * log(pi_theta(y|x) / pi_ref(y|x)).
    """
    # TODO: (1) the KL part at every real token, (2) add `score` at each row's last real token.
    # Hint: the index of the last real token of row b is mask[b].sum() - 1.
    raise NotImplementedError("Task: kl_penalized_rewards")


def gae_per_token(rewards, values, mask, gamma=1.0, lam=0.95):
    """Task 4. Generalized advantage estimation over the tokens of each response.

    rewards: (B, T) per-token rewards (from Task 3)
    values:  (B, T) critic values V(x, y_<t), the state in which token t was chosen
    mask:    (B, T) 1 for real tokens, 0 for padding (each row: 1s then 0s)

    The episode ends after the last real token, so the value of the state after it is 0:
        delta_t = r_t + gamma * V_{t+1} * mask_{t+1} - V_t        (V_{T} = 0 past the end of the tensor)
        A_t     = delta_t + gamma * lam * A_{t+1}                  (A = 0 at padding positions)
        returns_t = A_t + V_t                                      (the critic's regression target)
    Return (advantages, returns), both (B, T) and both 0 at padding positions.
    Inputs carry no gradient; work backwards over t with a Python loop.
    """
    # TODO: loop t = T-1, ..., 0. Remember: no value and no advantage beyond the last real token.
    raise NotImplementedError("Task: gae_per_token")


def ppo_rlhf_loss(logp, old_logp, values, advantages, returns, mask, clip_eps=0.2, vf_coef=0.5):
    """Task 5. PPO's loss for RLHF, averaged over all real tokens in the batch.

    logp:       (B, T) current log pi_theta of the sampled tokens (WITH gradient)
    old_logp:   (B, T) log-probs recorded at sampling time (no gradient)
    values:     (B, T) current critic values (WITH gradient)
    advantages: (B, T) from Task 4 (already normalized by the training loop)
    returns:    (B, T) critic targets from Task 4
    mask:       (B, T) 1 for real tokens

    rho_t        = exp(logp - old_logp)
    policy_loss  = - masked_mean( min(rho_t * A_t, clip(rho_t, 1-eps, 1+eps) * A_t) )
    value_loss   = 0.5 * masked_mean( (values - returns)^2 )
    loss         = policy_loss + vf_coef * value_loss
    where masked_mean(z) = sum(z * mask) / sum(mask)   (a "token mean": every real token counts once).

    Returns (loss, stats) with stats = {"policy_loss": float, "value_loss": float,
    "clip_frac": float}, clip_frac = fraction of real tokens with |rho - 1| > eps.
    """
    # TODO: ratio, clipped objective, token-mean with the mask, value loss, clip fraction.
    # `masked_mean` is imported from prefs.py.
    raise NotImplementedError("Task: ppo_rlhf_loss")


# ================================================================ Given: training loops
def comparisons(data):
    """Flatten ranked data into (counts_w, counts_l, gold_w, gold_l) tensors using ranking_to_pairs."""
    cw, cl, gw, gl = [], [], [], []
    for d in data:
        counts = digit_counts(d["out"])
        for w, l in ranking_to_pairs(d["ranking"]):
            cw.append(counts[w]); cl.append(counts[l])
            gw.append(d["gold"][w]); gl.append(d["gold"][l])
    return torch.stack(cw), torch.stack(cl), torch.stack(gw), torch.stack(gl)


def train_reward_model(train_data, steps=300, lr=0.05, l2=1e-3, seed=0):
    """Fit a DigitRM to the labeler's rankings with your bt_loss (full-batch Adam)."""
    torch.manual_seed(seed)
    cw, cl, _, _ = comparisons(train_data)
    rm = DigitRM()
    opt = torch.optim.Adam(rm.parameters(), lr=lr)
    for _ in range(steps):
        loss = bt_loss(rm(cw), rm(cl)) + l2 * (rm.w ** 2).sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return rm


@torch.no_grad()
def rm_accuracy(rm, data, kappa=None):
    """Held-out accuracy of the RM on the labeler's pairs (ties count 1/2), and the best possible
    expected accuracy given how noisy the labeler is (the "ceiling")."""
    from prefs import KAPPA
    kappa = KAPPA if kappa is None else kappa
    cw, cl, gw, gl = comparisons(data)
    rw, rl = rm(cw), rm(cl)
    acc = ((rw > rl).float() + 0.5 * (rw == rl).float()).mean().item()
    p = torch.sigmoid(kappa * (gw - gl))
    ceiling = torch.maximum(p, 1 - p).mean().item()
    return acc, ceiling


@dataclass
class RLHFConfig:
    iters: int = 120          # PPO iterations (rollout + update)
    batch: int = 64           # prompts per iteration (one response each)
    beta: float = 0.1         # KL coefficient (initial value, if target_kl is set)
    target_kl: float = None   # if set, adapt beta toward this KL (Ziegler et al. 2019 controller)
    lam: float = 0.95         # GAE lambda (gamma is 1)
    clip_eps: float = 0.2
    epochs: int = 4           # PPO epochs per batch (full-batch steps)
    lr: float = 3e-4          # policy learning rate
    vf_lr: float = 1e-3       # critic learning rate
    vf_coef: float = 0.5
    ptx_coef: float = 0.0     # PPO-ptx: weight of a supervised loss on correct additions (0 = off)
    seed: int = 0


def train_rlhf(sft_model, rm, rm_offset, cfg=RLHFConfig(), log_every=0):
    """PPO-based RLHF: rollout -> score -> KL-penalized per-token rewards -> GAE -> PPO epochs.

    sft_model: the starting policy; it is copied twice: the trainable policy, and the frozen reference.
    rm, rm_offset: the reward model, and a constant subtracted from its scores (so that the SFT
                   model's responses score 0 on average, as in InstructGPT).
    Returns (policy, history) where history is a dict of per-iteration lists.
    """
    torch.manual_seed(cfg.seed)
    gen = random.Random(cfg.seed)
    policy, ref = TinyLM(), TinyLM()
    policy.load_state_dict(sft_model.state_dict())
    ref.load_state_dict(sft_model.state_dict())
    ref.requires_grad_(False)
    critic = ValueModel(sft_model)
    opt = torch.optim.Adam([
        {"params": policy.parameters(), "lr": cfg.lr},
        {"params": critic.parameters(), "lr": cfg.vf_lr},
    ])
    beta = cfg.beta
    keys = ["rm", "gold", "kl", "acc", "malformed", "digits", "beta", "clip_frac", "value_loss"]
    hist = {k: [] for k in keys}

    for it in range(cfg.iters):
        # 1. rollout: one response per prompt from the current policy
        prompts, answers = make_batch(cfg.batch, gen)
        out = sample(policy, prompts)
        mask = out["mask"]
        with torch.no_grad():
            old_logp = sequence_logprobs(policy, out)
            ref_logp = sequence_logprobs(ref, out)
            old_values = critic(out)
            # 2. score with the (frozen) reward model
            score = rm(digit_counts(out)) - rm_offset
            # 3. per-token rewards, then 4. GAE advantages
            rewards = kl_penalized_rewards(old_logp, ref_logp, score, mask, beta)
            adv, ret = gae_per_token(rewards, old_values, mask, 1.0, cfg.lam)
            adv = masked_whiten(adv, mask)  # advantage normalization (Chapter 9)

        # 5. PPO epochs on this batch
        for _ in range(cfg.epochs):
            logp = sequence_logprobs(policy, out)
            values = critic(out)
            loss, stats = ppo_rlhf_loss(logp, old_logp, values, adv, ret, mask, cfg.clip_eps, cfg.vf_coef)
            if cfg.ptx_coef > 0:  # PPO-ptx (InstructGPT): keep the "pretraining" skill alive
                loss = loss + cfg.ptx_coef * sft_loss(policy, cfg.batch, gen)
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(list(policy.parameters()) + list(critic.parameters()), 1.0)
            opt.step()

        # logging: what the RM thinks, what the hidden labeler thinks, and how far we moved
        with torch.no_grad():
            kl_seq = ((old_logp - ref_logp) * mask).sum(1).mean().item()
            resp = decode_responses(out)
            gold = gold_score(out)
            hist["rm"].append(score.mean().item())
            hist["gold"].append(gold.mean().item())
            hist["kl"].append(kl_seq)
            hist["acc"].append(sum(r == a for r, a in zip(resp, answers)) / len(answers))
            hist["malformed"].append((gold < 0).float().mean().item())
            hist["digits"].append(digit_counts(out).sum(1).mean().item())
            hist["beta"].append(beta)
            hist["clip_frac"].append(stats["clip_frac"])
            hist["value_loss"].append(stats["value_loss"])
        if cfg.target_kl is not None:  # adaptive KL controller (Ziegler et al. 2019, K_beta = 0.1)
            e = max(-0.2, min(0.2, (kl_seq - cfg.target_kl) / cfg.target_kl))
            beta = beta * (1 + 0.1 * e)
        if log_every and (it % log_every == 0 or it == cfg.iters - 1):
            print(f"it {it:4d}  rm {hist['rm'][-1]:+.2f}  gold {hist['gold'][-1]:+.2f}  "
                  f"KL {kl_seq:6.2f}  acc {hist['acc'][-1]:.2f}  malformed {hist['malformed'][-1]:.2f}  "
                  f"e.g. {resp[0]!r}")
    return policy, hist