"""Chapter 13 lab: DPO on TinyLM, four ways, on the same offline preference data.

1. Load (or train and cache) the SFT model: TinyLM after supervised fine-tuning on "a+b=".
2. Build an OFFLINE preference dataset once: sample 4 answers per prompt from the SFT
   model and pair a correct answer (chosen) with a wrong one (rejected).
3. Train four copies of the SFT model on those same pairs:
       DPO  beta = 0.1          (a common default for large models)
       DPO  beta = 1.0
       DPO  beta = 0.1 + NLL    (an extra NLL term on the chosen, as in Llama 3 / Iterative RPO)
       IPO  beta = 0.1
4. Plot answer accuracy and held-out chosen/rejected log-probs over training.

    python run.py                    # uses your dpo.py
    RL_SOLUTIONS=1 python run.py     # uses the reference solution

Saves dpo_tinylm.png next to this file. About 1.5 minutes on a laptop CPU
(the first run also trains the SFT model, ~15 s, and caches it in sft_tinylm.pt).
"""

import copy
import importlib
import os
import time

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

dpo = importlib.import_module("solutions.dpo" if os.environ.get("RL_SOLUTIONS") else "dpo")

from prefs import make_preference_pairs  # noqa: E402
from tinylm import load_or_train_sft  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))

VARIANTS = [
    ("DPO β=0.1", dict(loss="dpo", beta=0.1)),
    ("DPO β=1", dict(loss="dpo", beta=1.0)),
    ("DPO β=0.1 + NLL", dict(loss="dpo", beta=0.1, nll_coef=1.0)),
    ("IPO β=0.1", dict(loss="ipo", beta=0.1)),
]


def main():
    t0 = time.time()
    sft = load_or_train_sft(os.path.join(HERE, "sft_tinylm.pt"))
    pairs = make_preference_pairs(sft, n_prompts=3000, seed=1)     # training pairs
    heldout = make_preference_pairs(sft, n_prompts=400, seed=2)    # for chosen/rejected log-probs
    print(f"SFT model ready; {len(pairs)} training pairs, {len(heldout)} held-out pairs "
          f"(e.g. {pairs[0]['prompt']} chosen {pairs[0]['chosen']!r} rejected {pairs[0]['rejected']!r})")

    results = {}
    for name, kw in VARIANTS:
        cfg = dpo.DPOConfig(lr=3e-4, batch_size=64, steps=300, **kw)
        policy = copy.deepcopy(sft)
        t1 = time.time()
        hist = dpo.train_dpo(policy, pairs, cfg, heldout=heldout, eval_every=25, seed=0)
        results[name] = hist
        print(f"{name:16s} acc greedy {hist['acc_greedy'][0]:.2f} -> {hist['acc_greedy'][-1]:.2f}, "
              f"sampled {hist['acc_sampled'][0]:.2f} -> {hist['acc_sampled'][-1]:.2f} | held-out log pi(chosen) "
              f"{hist['held_chosen'][0]:.2f} -> {hist['held_chosen'][-1]:.2f}, log pi(rejected) "
              f"{hist['held_rejected'][0]:.2f} -> {hist['held_rejected'][-1]:.2f}  ({time.time() - t1:.0f}s)")

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.3))
    colors = ["C3", "C0", "C2", "C1"]
    for (name, _), c in zip(VARIANTS, colors):
        h = results[name]
        axes[0].plot(h["eval_step"], h["acc_greedy"], color=c, label=name)
        axes[0].plot(h["eval_step"], h["acc_sampled"], color=c, ls=":", alpha=0.8)
        axes[1].plot(h["eval_step"], h["held_chosen"], color=c, label=name)
        axes[2].plot(h["eval_step"], h["held_rejected"], color=c, label=name)
    axes[0].set_title("answer accuracy (solid: greedy, dotted: sampled)")
    axes[1].set_title("held-out log π(chosen)   — should it go up?")
    axes[2].set_title("held-out log π(rejected)")
    for ax in axes:
        ax.set_xlabel("DPO step")
        ax.grid(alpha=0.3)
    axes[0].set_ylim(0, 1)
    axes[0].legend(fontsize=8)
    fig.suptitle("Offline preference tuning of TinyLM: same pairs, four losses")
    fig.tight_layout()
    out = os.path.join(HERE, "dpo_tinylm.png")
    fig.savefig(out, dpi=120)

    best = max(results, key=lambda k: results[k]["acc_greedy"][-1])
    print(f"saved {out}; best final greedy accuracy: {best} "
          f"({results[best]['acc_greedy'][-1]:.2f}); total {time.time() - t0:.0f}s")


if __name__ == "__main__":
    main()