"""Preference tuning, worked as Chapter 20 works it.

Three things, all reproducing the chapter's tables:

    the reward model      Bradley-Terry turns a ranking into a loss
    the KL penalty        why reward hacking needs a leash, in numbers
    DPO                   the same preference, with no reward model

    python3 alignment.py          # all three
    python3 alignment.py --bt     # Bradley-Terry and the reward-model loss
    python3 alignment.py --kl     # reward hacking and the penalty that stops it
    python3 alignment.py --dpo    # direct preference optimisation

Install:  nothing, the Python standard library is enough
"""

import argparse
import math


def sigmoid(z):
    return 1.0 / (1.0 + math.exp(-z))


# ------------------------------------------------------------ Bradley-Terry
# Four preference pairs. r_w is the reward the model gives the response the
# human preferred, r_l the one it gives the rejected response. The model is
# right on the first two, wrong on the third, and undecided on the fourth.
PAIRS = [
    ("clear answer",      1.20, 0.40),
    ("answer with steps", 2.10, 0.30),
    ("confident bluff",   0.30, 1.10),
    ("two paraphrases",   0.75, 0.75),
]


def show_bt():
    print("\nSTEP 1  the reward model learns from comparisons\n")
    print("  Nobody can score helpfulness on an absolute scale, but people can")
    print("  say which of two answers they prefer. Bradley-Terry turns that")
    print("  comparison into a probability:\n")
    print("      P(y_w > y_l) = sigma(r_w - r_l)\n")
    print("  pair                 r_w     r_l   r_w-r_l   P(y_w>y_l)    loss")
    print("  " + "-" * 64)
    total = 0.0
    for name, rw, rl in PAIRS:
        d = rw - rl
        p = sigmoid(d)
        loss = -math.log(p)
        total += loss
        print(f"  {name:18} {rw:5.2f}   {rl:5.2f}   {d:+7.2f}   {p:9.4f}  {loss:7.4f}")
    print("  " + "-" * 64)
    print(f"  mean loss over the four pairs{' ' * 29}{total / len(PAIRS):7.4f}\n")
    print("  Read the third row. The model prefers the bluff, so the pair costs")
    print(f"  {-math.log(sigmoid(0.30 - 1.10)):.4f}, more than three times the first row. Training")
    print("  pushes r_w up and r_l down until that ordering flips.\n")
    print("  Read the fourth. Equal rewards give exactly 0.5 and a loss of")
    print(f"  {-math.log(0.5):.4f} = log 2, the most a single comparison can tell you")
    print("  when the model has no opinion.")


# ------------------------------------------------------------- the KL leash
# Two candidate responses the policy could learn to produce. The first scores
# well with the reward model and badly with people: that is reward hacking.
CANDIDATES = [
    ("honest answer",  4.00,  1.00),
    ("reward hack",    5.00, 12.00),
]
BETAS = [0.00, 0.05, 0.10, 0.30]


def show_kl():
    print("\nSTEP 2  the penalty is not an afterthought\n")
    print("  The policy is optimised against the reward model, not against")
    print("  people. Left alone it finds text the reward model loves and humans")
    print("  do not. The objective keeps it on a leash:\n")
    print("      maximise  E[r(x,y)]  -  beta * KL(pi_theta || pi_ref)\n")
    print("  candidate          r      KL" + "".join(f"   beta={b:.2f}" for b in BETAS))
    print("  " + "-" * (31 + 11 * len(BETAS)))
    scores = {}
    for name, r, kl in CANDIDATES:
        row = [r - b * kl for b in BETAS]
        scores[name] = row
        cells = "".join(f"   {v:9.2f}" for v in row)
        print(f"  {name:15} {r:5.2f}  {kl:6.2f}{cells}")
    print("  " + "-" * (31 + 11 * len(BETAS)))
    print("\n  winner per beta:  ", end="")
    flip = None
    for i, b in enumerate(BETAS):
        best = max(scores, key=lambda n: scores[n][i])
        print(f"beta={b:.2f} -> {best}   ", end="")
        if best == "honest answer" and flip is None:
            flip = b
    print()
    print(f"\n  With no penalty the hack wins, because 5.00 beats 4.00 and the")
    print(f"  reward model is the only judge. The ordering flips at beta = {flip:.2f},")
    print("  where the hack's distance from the starting model finally costs")
    print("  more than the extra reward it buys.\n")
    print("  Too large a beta is its own failure. The policy then stays so close")
    print("  to the reference that preference training changes nothing.")


# ---------------------------------------------------------------------- DPO
# Log-probabilities of the preferred and rejected response under the policy
# being trained and under the frozen reference it started from.
BETA_DPO = 0.1
LOGP = [
    ("at the start",   -2.30, -2.30, -2.20, -2.20),
    ("after one step", -2.00, -2.30, -2.50, -2.20),
    ("over-corrected", -1.40, -2.30, -3.40, -2.20),
]


def show_dpo():
    print("\nSTEP 3  the same preference, with no reward model\n")
    print("  DPO drops the reward model and the reinforcement learning loop. It")
    print("  reads an implicit reward straight off the policy's own log-ratio:\n")
    print("      r_hat(y) = beta * ( log pi_theta(y) - log pi_ref(y) )")
    print("      loss     = -log sigma( r_hat(y_w) - r_hat(y_l) )\n")
    print(f"  beta = {BETA_DPO}\n")
    print("  stage             r_hat(y_w)  r_hat(y_l)   margin    sigma     loss")
    print("  " + "-" * 66)
    for name, tw, rw_, tl, rl_ in LOGP:
        rhat_w = BETA_DPO * (tw - rw_)
        rhat_l = BETA_DPO * (tl - rl_)
        margin = rhat_w - rhat_l
        s = sigmoid(margin)
        loss = -math.log(s)
        print(f"  {name:16} {rhat_w:+10.4f}  {rhat_l:+10.4f}  {margin:+7.4f}  {s:7.4f}  {loss:7.4f}")
    print("  " + "-" * 66)
    print("\n  At the start the policy is the reference, so both implicit rewards")
    print(f"  are 0, the margin is 0 and the loss is log 2 = {-math.log(0.5):.4f}. That is the")
    print("  same number the undecided Bradley-Terry pair gave in Step 1, and for")
    print("  the same reason: no opinion yet.\n")
    print("  One step raises the preferred response and lowers the rejected one,")
    print("  and the loss falls. The third row is the warning. Pushing the margin")
    print("  further keeps shrinking the loss, but the policy has now moved a long")
    print("  way from the reference, which is what beta is there to price.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--bt", action="store_true")
    ap.add_argument("--kl", action="store_true")
    ap.add_argument("--dpo", action="store_true")
    a = ap.parse_args()
    run_all = not (a.bt or a.kl or a.dpo)
    if a.bt or run_all:
        show_bt()
    if a.kl or run_all:
        show_kl()
    if a.dpo or run_all:
        show_dpo()


if __name__ == "__main__":
    main()