"""Chapter 13 lab: Direct Preference Optimization (DPO) and two of its variants, on TinyLM.

Your tasks (in order; run `python -m pytest -q -k task1` etc. after each):

    Task 1  sequence_logprob   sum (or average) of masked per-token log-probs
    Task 2  implicit_reward    beta * log(pi_theta / pi_ref)
            dpo_loss           the DPO loss, from four per-sequence log-probs
    Task 3  ipo_loss           IPO: regress the log-ratio margin to 1/(2 beta)
            simpo_loss         SimPO: length-normalized, reference-free, with a margin
    Task 4  dpo_step           one optimizer step on a minibatch of preference pairs

Everything below "Given" is complete: `train_dpo` runs your `dpo_step` in a loop and
evaluates accuracy on the addition task. When all tests pass, run `python run.py`.

Notation (as in the chapter): a preference pair is (x, y_w, y_l), "chosen" = y_w,
"rejected" = y_l. All per-pair tensors have shape (B,), B = number of pairs in the batch.
"""

import copy
import time
from dataclasses import dataclass

import torch
import torch.nn.functional as F

from prefs import batches, response_logprobs
from tinylm import accuracy


# ---------------------------------------------------------------------------
# Task 1: from per-token log-probs to a sequence log-prob
# ---------------------------------------------------------------------------
def sequence_logprob(token_logp, mask, average=False):
    """Log-probability of each whole response.

    Inputs:
      token_logp: (B, T) float, log pi(y_t | x, y_<t) at every response position
      mask:       (B, T) float, 1 for real response tokens, 0 for padding
      average:    if True, divide by the number of real tokens |y| (length normalization,
                  used by SimPO and by the NLL term in Task 4)

    Returns (B,):
        log pi(y|x) = sum_t mask_t * token_logp_t                 (average=False)
        (1/|y|) * sum_t mask_t * token_logp_t, |y| = sum_t mask_t  (average=True)

    Padding positions can hold any number (even -inf is possible in real code), so
    multiply by the mask BEFORE summing.
    """
    # TODO
    raise NotImplementedError("Task 1: sequence_logprob")


# ---------------------------------------------------------------------------
# Task 2: the implicit reward and the DPO loss
# ---------------------------------------------------------------------------
def implicit_reward(policy_logp, ref_logp, beta):
    """DPO's implicit reward (the chapter's r_hat_theta), shape (B,):

        r_hat(x, y) = beta * ( log pi_theta(y|x) - log pi_ref(y|x) )

    It is zero for every response when the policy equals the reference.
    """
    # TODO
    raise NotImplementedError("Task 2: implicit_reward")


def dpo_loss(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta):
    """The DPO loss, averaged over the batch.

    Inputs (all (B,)): sequence log-probs log pi_theta(y_w|x), log pi_theta(y_l|x),
    log pi_ref(y_w|x), log pi_ref(y_l|x). The policy ones carry gradients.

        r_w = implicit_reward(policy_chosen, ref_chosen, beta)
        r_l = implicit_reward(policy_rejected, ref_rejected, beta)
        loss = mean( -log sigmoid(r_w - r_l) )

    Returns (loss, r_w.detach(), r_l.detach()): the scalar loss and the two implicit
    rewards (for logging; detached so nobody backpropagates through them by accident).

    Use F.logsigmoid, not torch.log(torch.sigmoid(...)): the latter returns -inf for
    large negative margins, which then turns into NaN gradients.
    """
    # TODO
    raise NotImplementedError("Task 2: dpo_loss")


# ---------------------------------------------------------------------------
# Task 3: two variants
# ---------------------------------------------------------------------------
def ipo_loss(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta):
    """IPO loss (Azar et al. 2023), averaged over the batch.

        h    = [log pi_theta(y_w) - log pi_ref(y_w)] - [log pi_theta(y_l) - log pi_ref(y_l)]
        loss = mean( (h - 1/(2*beta))**2 )

    Note h has NO beta in front: IPO regresses the raw log-ratio margin to the finite
    target 1/(2 beta) instead of pushing it to infinity (the paper calls beta "tau").
    """
    # TODO
    raise NotImplementedError("Task 3: ipo_loss")


def simpo_loss(avg_chosen, avg_rejected, beta, margin):
    """SimPO loss (Meng et al. 2024), averaged over the batch. No reference model.

    Inputs: avg_chosen, avg_rejected (B,) are LENGTH-NORMALIZED log-probs
    (sequence_logprob(..., average=True)) under the policy.

        loss = mean( -log sigmoid( beta * avg_chosen - beta * avg_rejected - margin ) )

    (The paper calls the margin gamma; we avoid gamma because it is the discount factor.)
    """
    # TODO
    raise NotImplementedError("Task 3: simpo_loss")


# ---------------------------------------------------------------------------
# Task 4: one training step
# ---------------------------------------------------------------------------
@dataclass
class DPOConfig:
    loss: str = "dpo"          # "dpo" | "ipo" | "simpo"
    beta: float = 0.1
    nll_coef: float = 0.0      # weight of an extra NLL term on the chosen response (Llama 3 used 0.2)
    simpo_margin: float = 0.5  # SimPO's target margin (the paper's gamma)
    lr: float = 1e-3
    batch_size: int = 64
    steps: int = 200


def dpo_step(policy, ref, batch, opt, cfg):
    """One optimizer step on a minibatch of preference pairs.

    batch: list of dicts with keys "prompt", "chosen", "rejected" (strings).
    policy: the model being trained. ref: the frozen reference model (never updated).
    opt: a torch optimizer over policy.parameters(). cfg: a DPOConfig.

    Steps:
      1. Per-token log-probs of the chosen and of the rejected responses under the POLICY
         (with grad): response_logprobs(policy, prompts, responses) -> (token_logp, mask).
         Sum them with sequence_logprob -> pi_w, pi_l.
      2. The same under the REFERENCE, inside `with torch.no_grad():` -> ref_w, ref_l.
      3. The loss, by cfg.loss:
           "dpo":   dpo_loss(pi_w, pi_l, ref_w, ref_l, cfg.beta)
           "ipo":   ipo_loss(...), with implicit rewards computed by implicit_reward for logging
           "simpo": simpo_loss on length-normalized policy log-probs (average=True);
                    for logging use r_w = beta * avg_w, r_l = beta * avg_l (detached)
         Raise ValueError for anything else.
      4. If cfg.nll_coef > 0, add cfg.nll_coef * (mean over the batch of the per-token
         NLL of the chosen response), i.e.
             loss = loss - cfg.nll_coef * sequence_logprob(tok_w, mask_w, average=True).mean()
      5. opt.zero_grad(); loss.backward(); opt.step()

    Returns a dict of Python floats:
        "loss":          the total loss (after step 4)
        "reward_acc":    fraction of pairs with r_w > r_l
        "margin":        mean of (r_w - r_l)
        "chosen_logp":   mean of pi_w      (watch this one: does the chosen go UP?)
        "rejected_logp": mean of pi_l
    """
    prompts = [b["prompt"] for b in batch]
    chosen = [b["chosen"] for b in batch]
    rejected = [b["rejected"] for b in batch]
    # TODO: steps 1-5, then return the stats dict
    raise NotImplementedError("Task 4: dpo_step")


# ---------------------------------------------------------------------------
# Given: evaluation and the training loop
# ---------------------------------------------------------------------------
@torch.no_grad()
def mean_logps(model, pairs):
    """Average log pi(chosen) and log pi(rejected) over a list of pairs (no grad)."""
    prompts = [b["prompt"] for b in pairs]
    w = sequence_logprob(*response_logprobs(model, prompts, [b["chosen"] for b in pairs]))
    l = sequence_logprob(*response_logprobs(model, prompts, [b["rejected"] for b in pairs]))
    return w.mean().item(), l.mean().item()


def train_dpo(policy, pairs, cfg, heldout=None, eval_every=25, eval_n=500, seed=0, verbose=False):
    """Train `policy` with DPO-family losses on fixed preference `pairs` (offline).

    The reference model is a frozen copy of the policy at the start (the SFT model).
    Returns a history dict of lists.
    """
    torch.manual_seed(seed)
    ref = copy.deepcopy(policy).eval()
    for p in ref.parameters():
        p.requires_grad_(False)
    # TinyLM has no dropout, so train/eval mode only matters for tidiness.
    opt = torch.optim.Adam(policy.parameters(), lr=cfg.lr)
    it = batches(pairs, cfg.batch_size, seed=seed)
    hist = {k: [] for k in ["step", "loss", "reward_acc", "margin", "chosen_logp", "rejected_logp",
                            "eval_step", "acc_greedy", "acc_sampled", "held_chosen", "held_rejected"]}

    def evaluate(step):
        hist["eval_step"].append(step)
        hist["acc_greedy"].append(accuracy(policy, n=eval_n, greedy=True))
        hist["acc_sampled"].append(accuracy(policy, n=eval_n, greedy=False))
        if heldout:
            w, l = mean_logps(policy, heldout)
            hist["held_chosen"].append(w)
            hist["held_rejected"].append(l)

    evaluate(0)
    t0 = time.time()
    for step in range(1, cfg.steps + 1):
        stats = dpo_step(policy, ref, next(it), opt, cfg)
        hist["step"].append(step)
        for k, v in stats.items():
            hist[k].append(v)
        if step % eval_every == 0 or step == cfg.steps:
            evaluate(step)
            if verbose:
                print(f"  step {step:4d}  loss {stats['loss']:.3f}  reward acc {stats['reward_acc']:.2f}  "
                      f"acc(greedy) {hist['acc_greedy'][-1]:.3f}  acc(sampled) {hist['acc_sampled'][-1]:.3f}  "
                      f"({time.time() - t0:.0f}s)")
    return hist