"""Shared, complete helpers for the Chapter 21 capstone kit (you don't need to edit this).

What is here:
  get_teacher(), get_student()   cached TinyLM checkpoints (a strong and a weak model)
  response_logits(model, out)    logits at the response positions, shape (B, T, V)
  token_entropy(model, out)      mean per-token entropy of the policy on a batch
  evaluate(model, ...)           greedy accuracy, sampled pass@k (uses your tools.pass_at_k), entropy
  rl_train(model, steps, ...)    a small, configurable GRPO-family trainer used by the projects
  learnability_reward(p)         Absolute Zero's proposer reward
  bucket problems                arithmetic problems of different difficulty (for the curriculum project)

The trainer is deliberately compact. One "step":
  1. draw `prompts_per_step` problems; sample G responses each (temperature 1);
  2. score them with `reward_fn` (default: exact-match, 0/1);
  3. group advantages: "mean" r - mean(r) (Dr. GRPO / RLOO style), "std" (GRPO),
     or "maxrl" (your tools.maxrl_advantages);
  4. `epochs` gradient steps on the batch with a PPO-style clipped ratio
     [1 - eps_low, 1 + eps_high] and a token-mean loss. No KL term (as in most
     2025-26 RLVR recipes). With epochs = 1 the ratio is exactly 1 and the clip never acts.
"""

import importlib
import math
import os
import random

import numpy as np
import torch
import torch.nn.functional as F

from tinylm import (TinyLM, accuracy, decode_responses, make_batch, reward, sample,
                    sequence_logprobs, sft)

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

HERE = os.path.dirname(os.path.abspath(__file__))
torch.set_num_threads(1)


# ---------------------------------------------------------------- models
def _cached(path, build):
    model = TinyLM()
    if os.path.exists(path):
        model.load_state_dict(torch.load(path))
    else:
        build(model)
        torch.save(model.state_dict(), path)
    return model


def get_teacher():
    """A strong TinyLM (1,200 SFT steps; ~100% on 2-digit addition). Cached in teacher.pt (~20 s first time)."""
    def build(m):
        sft(m, steps=1200, seed=100)
    torch.manual_seed(0)  # seeds the initial weights created inside _cached
    return _cached(os.path.join(HERE, "teacher.pt"), build)


def get_student(steps=300):
    """A weak TinyLM (`steps` SFT steps; ~50% greedy / ~30% sampled at 300). Cached per `steps`."""
    def build(m):
        sft(m, steps=steps, seed=0)
    torch.manual_seed(1)
    return _cached(os.path.join(HERE, f"student{steps}.pt"), build)


# ---------------------------------------------------------------- per-token quantities
def response_logits(model, out):
    """Logits that produced each response token: (B, T, V), T = number of response columns."""
    ids, L = out["ids"], out["prompt_len"]
    return model(ids[:, :-1])[:, L - 1:]


@torch.no_grad()
def token_entropy(model, out):
    """Mean entropy (nats) of the next-token distribution over generated response tokens."""
    logp = F.log_softmax(response_logits(model, out), -1)
    ent = -(logp.exp() * logp).sum(-1)
    m = out["mask"]
    return float((ent * m).sum() / m.sum())


# ---------------------------------------------------------------- problems
def bucket_problem(rng, bucket):
    """Addition problems in five buckets. Returns (prompt, answer).

    The SFT data draws both operands uniformly from 0..99, so one-digit + one-digit
    problems are only 1% of it: the weak student is *worse* at them than at 2-digit ones.
    Bucket 4 uses 3-digit numbers it has never seen.
    """
    while True:
        if bucket == 4:                              # 3 digits + 1 digit (never seen in SFT)
            a, b = rng.randint(100, 499), rng.randint(0, 9)
            return f"{a}+{b}=", str(a + b)
        a, b = rng.randint(0, 99), rng.randint(0, 99)
        if bucket == 0:                              # 1 digit + 1 digit (rare in SFT)
            a, b = rng.randint(0, 9), rng.randint(0, 9)
        two, s = a >= 10 and b >= 10, a + b
        carry = (a % 10) + (b % 10) >= 10
        if (bucket == 0
                or (bucket == 1 and two and s < 100 and not carry)
                or (bucket == 2 and two and s < 100 and carry)
                or (bucket == 3 and s >= 100)):
            return f"{a}+{b}=", str(s)


BUCKET_NAMES = ["1d+1d", "2d no carry", "2d carry", "sum>=100", "3d+1d (new)"]


def learnability_reward(p_hat):
    """Absolute Zero's proposer reward: 0 if the solver never succeeds, else 1 - success rate."""
    return 0.0 if p_hat == 0 else 1.0 - p_hat


# ---------------------------------------------------------------- evaluation
@torch.no_grad()
def evaluate(model, n_prompts=100, n_samples=16, ks=(1, 4, 16), seed=123, problem_fn=None):
    """Greedy accuracy, sampled pass@k (unbiased estimator from n_samples per prompt) and entropy."""
    rng = random.Random(seed)
    if problem_fn is None:
        prompts, answers = make_batch(n_prompts, rng)
    else:
        pa = [problem_fn(rng) for _ in range(n_prompts)]
        prompts, answers = [p for p, _ in pa], [a for _, a in pa]
    greedy = decode_responses(sample(model, prompts, greedy=True))
    res = {"greedy_acc": float(np.mean([g == a for g, a in zip(greedy, answers)]))}
    P = [p for p in prompts for _ in range(n_samples)]
    out = sample(model, P, temperature=1.0)
    correct = np.array([r == a for r, a in zip(decode_responses(out), [a for a in answers for _ in range(n_samples)])])
    c = correct.reshape(n_prompts, n_samples).sum(1)
    for k in ks:
        res[f"pass@{k}"] = float(np.mean([tools.pass_at_k(n_samples, int(ci), k) for ci in c]))
    res["entropy"] = token_entropy(model, out)
    return res


# ---------------------------------------------------------------- RL trainer
def group_advantages(r, mode):
    """r: (P, G) numpy rewards -> (P, G) advantages."""
    mean = r.mean(1, keepdims=True)
    if mode == "mean":
        return r - mean
    if mode == "std":
        return (r - mean) / (r.std(1, keepdims=True) + 1e-6)
    if mode == "maxrl":
        return tools.maxrl_advantages(r)
    raise ValueError(mode)


def rl_train(model, steps, *, prompts_per_step=16, G=8, lr=3e-4, adv="mean", eps_low=0.2, eps_high=0.2,
             epochs=1, reward_fn=None, problem_fn=None, on_batch=None, eval_every=25, eval_kwargs=None,
             seed=0, log=print, tag=""):
    """Train `model` in place with a GRPO-family update. Returns a list of eval dicts (with 'step').

    reward_fn(response, answer) -> float           default: exact match (0/1)
    problem_fn(rng) -> (prompt, answer, info)      default: random 2-digit addition, info=None
    on_batch(infos, rewards)                       optional callback (e.g. a curriculum) after each rollout
    """
    reward_fn = reward_fn or reward
    if problem_fn is None:
        def problem_fn(rng):
            p, a = make_batch(1, rng)
            return p[0], a[0], None
    rng = random.Random(seed)
    torch.manual_seed(seed)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0)
    hist = []
    for step in range(steps + 1):
        if step % eval_every == 0 or step == steps:
            ev = evaluate(model, **(eval_kwargs or {}))
            ev["step"] = step
            hist.append(ev)
            if log:
                log(f"{tag:>10s} step {step:4d}  " + "  ".join(f"{k} {v:.3f}" for k, v in ev.items() if k != "step"))
        if step == steps:
            break
        probs = [problem_fn(rng) for _ in range(prompts_per_step)]
        prompts = [p for p, _, _ in probs for _ in range(G)]
        answers = [a for _, a, _ in probs for _ in range(G)]
        out = sample(model, prompts, temperature=1.0)
        r = np.array([reward_fn(y, a) for y, a in zip(decode_responses(out), answers)], dtype=float)
        r = r.reshape(prompts_per_step, G)
        if on_batch is not None:
            on_batch([info for _, _, info in probs], r)
        A = torch.tensor(group_advantages(r, adv).reshape(-1), dtype=torch.float32)[:, None]
        mask = out["mask"]
        if mask.sum() == 0 or torch.all(A == 0):
            continue
        with torch.no_grad():
            old = sequence_logprobs(model, out)
        for _ in range(epochs):
            logp = sequence_logprobs(model, out)
            ratio = torch.exp(logp - old)
            obj = torch.minimum(ratio * A, torch.clamp(ratio, 1 - eps_low, 1 + eps_high) * A)
            loss = -(obj * mask).sum() / mask.sum()
            opt.zero_grad()
            loss.backward()
            opt.step()
    return hist


def plot_histories(ax, hists, key, labels, title=None):
    for h, lab in zip(hists, labels):
        ax.plot([e["step"] for e in h], [e[key] for e in h], marker="o", ms=3, label=lab)
    ax.set_xlabel("RL step")
    ax.set_ylabel(key)
    if title:
        ax.set_title(title)
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8)