"""Chapter 14 lab: see your reward plumbing at work. Runs in a few seconds.

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

1. Verifier gallery: tricky responses scored by two naive verifiers and yours.
2. Optimization finds the hole: a 4-strategy softmax policy trained with
   REINFORCE (Chapter 7) on the addition task, once against the naive
   substring verifier and once against your math_reward.
3. Judge bias: single-order vs. order-swapped win rates as the position bias
   and the verbosity bias of a simulated judge grow.

Saves reward_sources.png next to this file.
"""

import importlib
import os

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

from helpers import (STRATEGIES, Answer, SimulatedJudge, last_number_reward,  # noqa: E402
                     make_response, naive_substring_reward)

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


# ---------------------------------------------------------------------------
# 1. Verifier gallery
# ---------------------------------------------------------------------------

GALLERY = [
    (r"37+48: 7+8=15, carry 1, so \boxed{85}.", "85", "honest and correct"),
    (r"\boxed{\frac{1}{2}}", "0.5", "same value, LaTeX form"),
    (r"The total is \boxed{1,000} dollars.", "1000", "thousands separator"),
    (r"It's \boxed{75} or \boxed{85}.", "85", "hedging with two answers"),
    (" ".join(rf"\boxed{{{n}}}" for n in range(20, 199)), "85", "spam every number"),
    (r"\boxed{157}", "57", "wrong, but contains the gold string"),
    ("The answer is 85. I'm 100% sure.", "85", "correct, then a stray number"),
    ("The answer is 85.", "85", "correct, no box"),
]


def gallery():
    print("1. Verifier gallery (1 = rewarded)")
    print(f"   {'case':40s} {'substring':>9s} {'last num':>9s} {'yours':>6s}")
    for resp, gold, label in GALLERY:
        a = naive_substring_reward(resp, gold)
        b = last_number_reward(resp, gold)
        c = R.math_reward(resp, gold)
        print(f"   {label:40s} {a:9.0f} {b:9.0f} {c:6.0f}")
    print()


# ---------------------------------------------------------------------------
# 2. REINFORCE over answering strategies
# ---------------------------------------------------------------------------

def train_strategy_policy(reward_fn, steps=300, batch=32, lr=0.5, seed=0):
    """Softmax policy over the 4 strategies in helpers.STRATEGIES, trained with
    REINFORCE and a batch-mean baseline (Chapter 7). Each sample writes a real
    response to a random addition problem, which `reward_fn` then scores."""
    rng = np.random.default_rng(seed)
    theta = np.zeros(len(STRATEGIES))
    history = []
    for _ in range(steps):
        pi = np.exp(theta - theta.max())
        pi /= pi.sum()
        history.append(pi.copy())
        acts = rng.choice(len(pi), size=batch, p=pi)
        rewards = np.empty(batch)
        for i, k in enumerate(acts):
            a, b = rng.integers(10, 100, size=2)
            rewards[i] = reward_fn(make_response(int(k), int(a), int(b), rng), str(a + b))
        adv = rewards - rewards.mean()
        grad = np.zeros_like(theta)
        for k, A in zip(acts, adv):
            g = -pi.copy()
            g[k] += 1.0              # gradient of log softmax: e_k - pi
            grad += A * g
        theta += lr * grad / batch
    return np.array(history)


# ---------------------------------------------------------------------------
# 3. Judge bias sweeps
# ---------------------------------------------------------------------------

GOOD = Answer("concise, correct", quality=7.0, n_words=150)
PADDED = Answer("padded, slightly worse", quality=6.0, n_words=600)


def judge_sweeps():
    bs = np.linspace(0, 4, 41)
    lams = np.linspace(0, 1.5, 41)
    pos = {"good shown first": [], "good shown second": [], "swapped + averaged": []}
    for b in bs:
        j = SimulatedJudge(position_bias=b)
        good_vs_ok = Answer("ok", quality=6.0, n_words=150)  # same length: only position bias matters
        pos["good shown first"].append(j("q", GOOD, good_vs_ok))
        pos["good shown second"].append(1 - j("q", good_vs_ok, GOOD))
        pos["swapped + averaged"].append(R.debiased_preference(j, "q", GOOD, good_vs_ok))
    verb = {"swapped + averaged": []}
    for lam in lams:
        j = SimulatedJudge(position_bias=1.0, verbosity_bias=lam)
        verb["swapped + averaged"].append(R.debiased_preference(j, "q", GOOD, PADDED))
    return bs, pos, lams, verb


def main():
    gallery()

    print("2. Training a 4-strategy policy against two verifiers (300 steps x 32 samples each)...")
    h_naive = train_strategy_policy(naive_substring_reward)
    h_yours = train_strategy_policy(R.math_reward)
    for name, h in [("naive substring verifier", h_naive), ("your math_reward", h_yours)]:
        final = ", ".join(f"{s} {p:.2f}" for s, p in zip(STRATEGIES, h[-1]))
        print(f"   final policy vs {name:25s}: {final}")
    print()

    bs, pos, lams, verb = judge_sweeps()
    print("3. Judge bias: P(judge prefers the better answer)")
    print(f"   position bias b=3, equal lengths: shown second {pos['good shown second'][30]:.2f}, "
          f"swapped {pos['swapped + averaged'][30]:.2f}")
    p_verb = R.debiased_preference(SimulatedJudge(position_bias=1.0, verbosity_bias=1.0), "q", GOOD, PADDED)
    print(f"   verbosity bias 1.0 vs a 4x longer, 1-point-worse answer, swapped: {p_verb:.2f}")

    fig, axes = plt.subplots(2, 2, figsize=(11, 8))
    steps = np.arange(len(h_naive))
    for ax, h, title in [(axes[0, 0], h_naive, "Policy trained on the naive substring verifier"),
                         (axes[0, 1], h_yours, "Policy trained on your math_reward")]:
        for k, s in enumerate(STRATEGIES):
            ax.plot(steps, h[:, k], label=s, lw=2)
        ax.set_title(title, fontsize=10)
        ax.set_xlabel("REINFORCE step")
        ax.set_ylabel("probability of strategy")
        ax.set_ylim(-0.02, 1.02)
        ax.legend(fontsize=8)
    ax = axes[1, 0]
    for k, v in pos.items():
        ax.plot(bs, v, lw=2, label=k)
    ax.axhline(0.5, color="gray", lw=1, ls=":")
    ax.set_title("Position bias (equal lengths): swapping restores the verdict", fontsize=10)
    ax.set_xlabel("position bias b (logits for the first slot)")
    ax.set_ylabel("P(better answer wins)")
    ax.legend(fontsize=8)
    ax = axes[1, 1]
    ax.plot(lams, verb["swapped + averaged"], lw=2, color="C2", label="swapped + averaged")
    ax.axhline(0.5, color="gray", lw=1, ls=":")
    ax.set_title("Verbosity bias vs a 4x longer, worse answer: swapping can't help", fontsize=10)
    ax.set_xlabel("verbosity bias λ (logits per doubling of length)")
    ax.set_ylabel("P(better answer wins)")
    ax.legend(fontsize=8)
    fig.tight_layout()
    out = os.path.join(HERE, "reward_sources.png")
    fig.savefig(out, dpi=110)
    naive_hack = h_naive[-1, 2] + h_naive[-1, 3]
    yours_hack = h_yours[-1, 2] + h_yours[-1, 3]
    print(f"\nSaved {out}. Hedge+spam probability: naive verifier {naive_hack:.2f}, yours {yours_hack:.2f}.")


if __name__ == "__main__":
    main()