"""Chapter 17 lab: test-time scaling curves, Monte Carlo step labels, and
PRM-guided beam search vs best-of-N, all in one figure.

    python run.py                  # uses your ttc.py and process.py
    RL_SOLUTIONS=1 python run.py   # uses the reference solutions

Saves process_rewards.png next to this file. Takes about 15 seconds.
"""

import importlib
import os
import random

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

from chain import EXAMPLE, ChainTask, NoisyPRM, sample_answers, verifier_scores

SOL = bool(os.environ.get("RL_SOLUTIONS"))
ttc = importlib.import_module("solutions.ttc" if SOL else "ttc")
proc = importlib.import_module("solutions.process" if SOL else "process")
HERE = os.path.dirname(os.path.abspath(__file__))
TASK = ChainTask()


# ---------------------------------------------------------------- (a) majority vote, binary answers
def panel_binary(ax):
    ns = np.arange(1, 52, 2)
    rng = np.random.default_rng(0)
    for p, col in [(0.4, "C3"), (0.5, "C7"), (0.6, "C0"), (0.7, "C2")]:
        ax.plot(ns, [ttc.majority_vote_accuracy(p, int(n)) for n in ns], color=col, label=f"p = {p}")
    # check the formula against simulation for p = 0.6
    sim_n = [1, 5, 11, 21, 41]
    sim = [np.mean([ttc.majority_vote(list(sample_answers(0.6, n, 1.0, rng))) == 0 for _ in range(1500)])
           for n in sim_n]
    ax.plot(sim_n, sim, "o", color="C0", mfc="none", label="simulated (p = 0.6)")
    ax.plot(ns, ttc.pass_at_n(0.4, ns), ":", color="C3", label="pass@N, p = 0.4")
    ax.set(title="(a) Majority vote, right/wrong answers", xlabel="N samples", ylabel="accuracy", ylim=(0, 1.02))
    ax.legend(fontsize=7, loc="center right")


# ---------------------------------------------------------------- (b) a benchmark of mixed difficulty
def panel_benchmark(ax, n_problems=300):
    """Problems with p ~ Beta(0.8, 1.2) (mean 0.4). For each, draw 1024 samples once
    and evaluate every rule on the first N of them, N = 1, 2, 4, ..., 1024."""
    rng = np.random.default_rng(1)
    Ns = [2**i for i in range(11)]
    rules = {
        "majority vote": [],
        "best-of-N, honest verifier (f = 0)": [],
        "best-of-N, fooled 1% of the time": [],
        "weighted vote, fooled 1%": [],
    }
    ps = rng.beta(0.8, 1.2, n_problems)
    for p in ps:
        a = sample_answers(p, Ns[-1], 0.5, rng)
        s0 = verifier_scores(a, 2.0, 0.0, rng)
        s1 = verifier_scores(a, 2.0, 0.01, rng)
        al = list(a)
        rules["majority vote"].append([ttc.majority_vote(al[:n]) == 0 for n in Ns])
        rules["best-of-N, honest verifier (f = 0)"].append([ttc.best_of_n(al[:n], s0[:n]) == 0 for n in Ns])
        rules["best-of-N, fooled 1% of the time"].append([ttc.best_of_n(al[:n], s1[:n]) == 0 for n in Ns])
        rules["weighted vote, fooled 1%"].append([ttc.weighted_vote(al[:n], s1[:n]) == 0 for n in Ns])
    ax.plot(Ns, [np.mean(ttc.pass_at_n(ps, n)) for n in Ns], "k--", label="pass@N (oracle)")
    out = {}
    for (name, hits), col in zip(rules.items(), ["C1", "C0", "C3", "C2"]):
        acc = np.mean(hits, axis=0)
        out[name] = acc
        ax.plot(Ns, acc, "o-", ms=3, color=col, label=name)
    ceiling = np.mean(ps > 0.5 / 1.5)
    ax.axhline(ceiling, color="C1", lw=0.8, ls=":")
    ax.text(1.1, ceiling + 0.015, "majority-vote ceiling", color="C1", fontsize=7)
    ax.set_xscale("log", base=2)
    ax.set(title="(b) Test-time scaling on a mixed benchmark", xlabel="N samples (log scale)",
           ylabel="accuracy", ylim=(0, 1.02))
    ax.legend(fontsize=7, loc="lower right")
    return out, Ns


# ---------------------------------------------------------------- (c, d) Monte Carlo step labels
def panel_labels(ax_c, ax_d):
    bad = [1, 4, 2, 7, 5, 7]
    j = np.arange(len(bad) + 1)
    exact = np.array([TASK.exact_value(EXAMPLE, bad[:i]) for i in j])
    s4, h4 = proc.mc_step_labels(TASK, EXAMPLE, bad, 4, random.Random(0))
    s32, h32 = proc.mc_step_labels(TASK, EXAMPLE, bad, 32, random.Random(1))
    w = 0.27
    ax_c.bar(j - w, exact, w, color="0.6", label="true value V(prefix)")
    ax_c.bar(j, s4, w, color="C0", alpha=0.6, label="soft label, k = 4")
    ax_c.bar(j + w, s32, w, color="C0", label="soft label, k = 32")
    ax_c.plot(j + w, h32, "rx", label="hard label, k = 32")
    ax_c.set_xticks(j, ["start"] + [f"after\nstep {i}" for i in j[1:]], fontsize=7)
    ax_c.set(title="(c) Monte Carlo labels; the mistake is step 3", ylabel="P(correct answer)", ylim=(0, 1.1))
    ax_c.legend(fontsize=7)

    steps = np.arange(1, len(bad) + 1)
    r_exact = proc.process_rewards(exact)
    r_mc = proc.process_rewards(s32)
    outcome = float(TASK.is_correct(EXAMPLE, bad)) - exact[0]  # outcome reward minus a baseline V_0
    ax_d.bar(steps - 0.2, r_exact, 0.4, color="0.6", label="process reward from true values")
    ax_d.bar(steps + 0.2, r_mc, 0.4, color="C0", label="process reward from k = 32 labels")
    ax_d.axhline(outcome, color="C3", ls="--", label=f"outcome advantage (same for every step): {outcome:+.2f}")
    ax_d.axhline(0, color="k", lw=0.6)
    ax_d.set_xticks(steps, [f"step {i}" for i in steps], fontsize=7)
    ax_d.set(title="(d) Credit assignment: step rewards V_j - V_(j-1)", ylabel="reward / advantage")
    ax_d.legend(fontsize=7, loc="upper right")
    return exact, s32


# ---------------------------------------------------------------- (e, f) beam search vs best-of-N
def compare_search(sigma, n_problems=300):
    rng = random.Random(123)
    probs = [TASK.sample_problem(rng) for _ in range(n_problems)]
    prm = NoisyPRM(sigma, seed=1)
    rows = []
    for B, M in [(1, 2), (1, 4), (2, 4), (4, 4), (4, 8), (8, 8)]:
        hits, budget = [], []
        for i, p in enumerate(probs):
            sol, n = proc.beam_search(TASK, p, prm, B, M, random.Random(10_000 + i))
            hits.append(TASK.is_correct(p, sol))
            budget.append(n)
        steps = float(np.mean(budget))
        N = max(1, round(steps / len(EXAMPLE.ops)))  # same number of generated steps, as full samples
        bon, wv, maj = [], [], []
        for i, p in enumerate(probs):
            r = random.Random(20_000 + i)
            sols = [TASK.sample_solution(p, r) for _ in range(N)]
            ans = [s[-1] for s in sols]
            sc = [prm(p, s) for s in sols]
            a = TASK.answer(p)
            bon.append(ttc.best_of_n(ans, sc) == a)
            wv.append(ttc.weighted_vote(ans, sc) == a)
            maj.append(ttc.majority_vote(ans) == a)
        rows.append((steps, np.mean(hits), np.mean(bon), np.mean(wv), np.mean(maj)))
    single = np.mean([TASK.is_correct(p, TASK.sample_solution(p, random.Random(i))) for i, p in enumerate(probs)])
    return np.array(rows), single


def panel_search(ax, sigma, label):
    rows, single = compare_search(sigma)
    x = rows[:, 0]
    ax.plot(x, rows[:, 1], "o-", color="C4", label="step-level beam search (PRM)")
    ax.plot(x, rows[:, 2], "s-", color="C0", label="best-of-N (PRM score of full solution)")
    ax.plot(x, rows[:, 3], "^-", color="C2", label="weighted vote (PRM)")
    ax.plot(x, rows[:, 4], "d-", color="C1", label="majority vote")
    ax.axhline(single, color="0.5", ls=":", label="one sample")
    ax.set_xscale("log")
    ax.set_xticks([12, 25, 50, 100, 200, 400], ["12", "25", "50", "100", "200", "400"])
    ax.xaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
    ax.set(title=f"{label} (noise sigma = {sigma})", xlabel="generated steps per problem (log)",
           ylabel="accuracy", ylim=(0, 1.02))
    ax.legend(fontsize=7, loc="lower right")
    return rows


def main():
    fig, axes = plt.subplots(2, 3, figsize=(16, 9))
    panel_binary(axes[0, 0])
    bench, Ns = panel_benchmark(axes[0, 1])
    exact, s32 = panel_labels(axes[0, 2], axes[1, 0])
    good = panel_search(axes[1, 1], 1.0, "(e) Good PRM")
    poor = panel_search(axes[1, 2], 3.0, "(f) Poor PRM")
    fig.tight_layout()
    path = os.path.join(HERE, "process_rewards.png")
    fig.savefig(path, dpi=110)

    bon = bench["best-of-N, fooled 1% of the time"]
    print(f"benchmark: majority@1024 {bench['majority vote'][-1]:.2f} | fooled best-of-N peaks at "
          f"{bon.max():.2f} (N={Ns[int(bon.argmax())]}) and ends at {bon[-1]:.2f} | "
          f"step-3 value {exact[2]:.2f} -> {exact[3]:.2f} | "
          f"beam vs best-of-N: good PRM {good[0, 1]:.2f} vs {good[0, 2]:.2f} (~{good[0, 0]:.0f} steps), "
          f"poor PRM {poor[1, 1]:.2f} vs {poor[1, 2]:.2f} (~{poor[1, 0]:.0f} steps)  -> {os.path.basename(path)}")


if __name__ == "__main__":
    main()