"""Lab 11: see your functions at work on a real (tiny) language model and on the toy LM.

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

Saves llm_as_policy.png next to this file (about 20-40 s on a laptop; the first run
also trains the SFT TinyLM for ~10 s and caches it in sft_tinylm.pt).

Panel 1: the SFT TinyLM sampled at different temperatures: accuracy goes down and
         diversity goes up as T rises (temperature trades one for the other).
Panel 2: KL(pi_T || pi) between the temperature-T sampler and the model itself,
         estimated per token with k1 and k3 from your Task 3 on the sampled tokens.
Panel 3: the reward-KL plane for the toy LM: the frontier traced by pi* as beta
         varies (Task 4) vs. what temperature alone can reach, and, for a reward
         model with a blind spot, the true reward of pi* rising then falling.
"""

import importlib
import os
import random

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

import tinylm
import toylm

policy = importlib.import_module("solutions.policy" if os.environ.get("RL_SOLUTIONS") else "policy")
HERE = os.path.dirname(os.path.abspath(__file__))


def temperature_sweep(model, temps, n_prompts=400, k_div=8):
    rng = random.Random(7)
    prompts, answers = tinylm.make_batch(n_prompts, rng)
    acc, distinct, kl1, kl3, kl3_sd, klx = [], [], [], [], [], []
    for T in temps:
        g = torch.Generator().manual_seed(0)
        out = policy.generate(model, prompts, temperature=T, generator=g)
        resp = tinylm.decode_responses(out)
        acc.append(np.mean([r == a for r, a in zip(resp, answers)]))
        # diversity: distinct answers among k_div samples for the first 50 prompts
        sub = [p for p in prompts[:50] for _ in range(k_div)]
        o2 = policy.generate(model, sub, temperature=T, generator=g)
        r2 = tinylm.decode_responses(o2)
        distinct.append(np.mean([len(set(r2[i * k_div:(i + 1) * k_div])) for i in range(50)]))
        # per-token KL(pi_T || pi) on the tokens pi_T actually sampled
        if T > 0:
            with torch.no_grad():
                logits = model(out["ids"])
                # the sampler forbids <bos> and <pad>; mirror that so log-probs match what was sampled
                logits[..., [tinylm.BOS, tinylm.PAD]] = float("-inf")
                lp_T = policy.response_logprobs(logits / T, out["ids"], out["prompt_len"])
                lp_1 = policy.response_logprobs(logits, out["ids"], out["prompt_len"])
            k1, _, k3 = policy.kl_estimators(lp_T, lp_1)
            m = out["mask"].bool()
            # sum over the tokens of each response = an estimate of the sequence-level KL
            s1 = policy.sequence_logprob(k1, m)
            s3 = policy.sequence_logprob(k3, m)
            kl1.append(float(s1.mean()))
            kl3.append(float(s3.mean()))
            kl3_sd.append((float(s1.std()), float(s3.std())))
            # ground truth (almost): at every position, sum over the WHOLE vocabulary
            # KL(pi_T(.|s_t) || pi(.|s_t)) exactly, then add up along the response.
            L = out["prompt_len"]
            lpT = torch.log_softmax(logits[:, L - 1:-1] / T, -1)
            lp1 = torch.log_softmax(logits[:, L - 1:-1], -1)
            pT = lpT.exp()
            per_pos = torch.where(pT > 0, pT * (lpT - lp1), torch.zeros_like(pT)).sum(-1)
            klx.append(float(policy.sequence_logprob(per_pos, m).mean()))
        else:
            kl1.append(np.nan); kl3.append(np.nan); kl3_sd.append((np.nan, np.nan)); klx.append(np.nan)
    return np.array(acc), np.array(distinct), np.array(kl1), np.array(kl3), np.array(kl3_sd), np.array(klx)


def frontier(sc, betas, temps):
    ref = toylm.seq_probs(sc["first"], sc["second"])
    true_r = sc.get("true_reward", sc["reward"])
    fr = []
    for b in betas:
        pi, _ = policy.tilted_optimum(ref, sc["reward"], b)
        fr.append((toylm.kl(pi, ref), toylm.expectation(pi, sc["reward"]), toylm.expectation(pi, true_r)))
    tp = []
    for T in temps:
        f, s = toylm.temperature_policy(sc["first"], sc["second"], T)
        P = toylm.seq_probs(f, s)
        tp.append((toylm.kl(P, ref), toylm.expectation(P, sc["reward"])))
    return np.array(fr), np.array(tp), toylm.expectation(ref, sc["reward"])


def main():
    torch.manual_seed(0)
    model = tinylm.load_or_train_sft(os.path.join(HERE, "sft_tinylm.pt"))
    temps = [0.0, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0]
    acc, distinct, kl1, kl3, sds, klx = temperature_sweep(model, temps)

    fig, ax = plt.subplots(1, 3, figsize=(15, 4.3))
    a = ax[0]
    a.plot(temps, acc, "o-", color="C0", label="accuracy (exact answer)")
    a.set_xlabel("sampling temperature T (0 = greedy)")
    a.set_ylabel("accuracy", color="C0")
    a.set_ylim(0, 1)
    a2 = a.twinx()
    a2.plot(temps, distinct, "s--", color="C1", label="distinct answers in 8 samples")
    a2.set_ylabel("distinct answers among 8 samples", color="C1")
    a.set_title("SFT TinyLM: temperature trades accuracy for diversity")

    b = ax[1]
    Ts = np.array(temps[1:])
    s1 = np.array([s[0] for s in sds[1:]]); s3 = np.array([s[1] for s in sds[1:]])
    b.errorbar(Ts - 0.02, kl1[1:], yerr=s1 / np.sqrt(400), fmt="o-", label="k1 (mean ± s.e.)", color="C0", capsize=3)
    b.errorbar(Ts + 0.02, kl3[1:], yerr=s3 / np.sqrt(400), fmt="s-", label="k3 (mean ± s.e.)", color="C2", capsize=3)
    b.plot(Ts, klx[1:], "k--", lw=1.2, label="exact per-position KL, summed")
    b.axhline(0, color="grey", lw=0.8)
    b.set_xlabel("sampling temperature T")
    b.set_ylabel("estimated KL(π_T ‖ π) per response [nats]")
    b.set_title("How far does temperature move the policy?")
    b.legend()

    c = ax[2]
    betas = np.logspace(-1.6, 1.7, 120)
    tgrid = np.logspace(-0.8, 0.5, 120)
    for name, colr in (("worked", "C0"), ("hack", "C3")):
        sc = toylm.SCENARIOS[name]
        fr, tp, eref = frontier(sc, betas, tgrid)
        c.plot(fr[:, 0], fr[:, 1], "-", color=colr, label=f"{name}: π* as β varies")
        c.plot(tp[:, 0], tp[:, 1], ":", color=colr, label=f"{name}: π_ref at temperature T")
        if "true_reward" in sc:
            c.plot(fr[:, 0], fr[:, 2], "--", color=colr, alpha=0.8, label="hack: TRUE reward of π*")
        c.plot([0], [eref], "o", color=colr)
    c.set_xlabel("KL(π ‖ π_ref) [nats]")
    c.set_ylabel("expected reward")
    c.set_xlim(0, 3)
    c.set_title("Toy LM: the reward–KL frontier")
    c.legend(fontsize=7.5)

    fig.tight_layout()
    path = os.path.join(HERE, "llm_as_policy.png")
    fig.savefig(path, dpi=120)
    i1 = temps.index(1.0)
    print(f"accuracy: greedy {acc[0]:.2f}, T=1 {acc[i1]:.2f}, T=2 {acc[-1]:.2f}; "
          f"KL(pi_T || pi) at T=0.5: k1 {kl1[2]:.3f}, k3 {kl3[2]:.3f}, exact {klx[2]:.3f} nats; saved {os.path.basename(path)}")


if __name__ == "__main__":
    main()