"""Chapter 17 lab, part 2: process rewards (Tasks 4-5).

Label every step of a solution by rolling out from it (Math-Shepherd style),
turn the labels into step rewards, and use a process reward model (PRM) to
steer a step-level beam search. The toy task lives in chain.py; read its
docstring first (two minutes).

Tests: test_process.py.   Run: python -m pytest -q test_process.py
"""

import numpy as np  # noqa: F401


# ---------------------------------------------------------------- Task 4a
def mc_step_labels(task, problem, steps, k, rng):
    """Monte Carlo process labels for every prefix of a solution.

    For each j = 0, 1, ..., T (T = len(steps)), take the prefix steps[:j], run
    k rollouts from it with the policy, task.rollout(problem, prefix, rng),
    and check each finished solution with task.is_correct(problem, solution).

        soft[j] = (number of correct rollouts) / k      Math-Shepherd's "soft" label
        hard[j] = 1.0 if at least one rollout is correct else 0.0     the "hard" label

    soft[j] is a Monte Carlo estimate of the value V(steps[:j]): the chance the
    policy reaches the right answer from there (Chapter 5's MC prediction).
    soft[0] is the value of the bare problem; soft[T] just checks the full
    solution (every rollout returns it unchanged).

    Draw rollouts in this order: j = 0..T, and k rollouts for each j, all from
    the one `rng` (a random.Random) you are given. The tests rely on it.

    Returns (soft, hard): two float numpy arrays of shape (T + 1,).
    """
    # TODO: loop over j, then over k rollouts.
    raise NotImplementedError("Task 4a: implement mc_step_labels")


# ---------------------------------------------------------------- Task 4b
def process_rewards(values):
    """Turn per-prefix values V_0..V_T into per-step rewards:

        r_j = V_j - V_{j-1},   j = 1..T     (an array of length T)

    This is the advantage of step j (Chapter 8): how much writing step j
    changed the chance of success. Note the rewards telescope:
    sum_j r_j = V_T - V_0 = (outcome) - (value of the problem).

    Example: [0.3, 0.4, 0.05, 0.0] -> [0.1, -0.35, -0.05]
    """
    # TODO: one line with numpy.
    raise NotImplementedError("Task 4b: implement process_rewards")


# ---------------------------------------------------------------- Task 5
def beam_search(task, problem, prm, beam_width, expand, rng):
    """Step-level beam search guided by a process reward model.

        beams = [ [] ]                       # one empty prefix
        repeat T = len(problem.ops) times:
            candidates = []
            for each prefix in beams:
                repeat `expand` times:
                    s = task.policy_step(problem, prefix, rng)   # sample one next step
                    add prefix + [s] to candidates, unless an identical
                    candidate is already there (keep the first copy)
            score every candidate with prm(problem, candidate)
            beams = the `beam_width` highest-scoring candidates
                    (sort by score, descending; ties keep candidate order)
        return the highest-scoring beam (beams[0]) and n_generated

    n_generated counts EVERY call to task.policy_step (duplicates included):
    that is the compute you spent, measured in generated steps.

    Visit beams in order, and for each beam make its `expand` calls in a row,
    all from the one `rng` (the tests rely on this order).

    Returns (best_solution, n_generated): a list of T ints and an int.
    """
    # TODO: follow the pseudocode. Python's sorted(..., key=...) is stable.
    raise NotImplementedError("Task 5: implement beam_search")