"""Tests for process.py (Tasks 4-5). Run: python -m pytest -q test_process.py"""

import importlib
import os
import random

import numpy as np
import pytest

from chain import EXAMPLE, ChainTask, NoisyPRM

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

TASK = ChainTask()
GOOD = [1, 4, 3, 8, 6, 8]  # correct solution of EXAMPLE
BAD = [1, 4, 2, 7, 5, 7]  # first mistake at step 3 (4 + 9 = 13 -> 3, wrote 2)


# ------------------------------------------------------------ Task 4a
def test_mc_labels_shapes_and_endpoints():
    soft, hard = proc.mc_step_labels(TASK, EXAMPLE, BAD, 8, random.Random(0))
    assert soft.shape == (7,) and hard.shape == (7,), "one label per prefix: j = 0..T, so T + 1 = 7"
    assert soft[-1] == 0.0 and hard[-1] == 0.0, "the full BAD solution is wrong: its label must be 0"
    soft, hard = proc.mc_step_labels(TASK, EXAMPLE, GOOD, 8, random.Random(0))
    assert soft[-1] == 1.0 and hard[-1] == 1.0, "the full GOOD solution is right: its label must be 1"


def test_mc_labels_hard_is_any_success():
    soft, hard = proc.mc_step_labels(TASK, EXAMPLE, BAD, 5, random.Random(3))
    assert np.all(hard == (soft > 0)), "hard label = 1 exactly when at least one rollout succeeded"
    assert np.all((soft * 5) == np.round(soft * 5)), "soft labels with k = 5 must be multiples of 1/5"


def test_mc_labels_use_rng_in_the_documented_order():
    rng = random.Random(42)
    soft, _ = proc.mc_step_labels(TASK, EXAMPLE, BAD, 3, rng)
    ref_rng = random.Random(42)
    expected = []
    for j in range(len(BAD) + 1):
        wins = sum(TASK.is_correct(EXAMPLE, TASK.rollout(EXAMPLE, BAD[:j], ref_rng)) for _ in range(3))
        expected.append(wins / 3)
    assert np.allclose(soft, expected), f"expected {expected}, got {soft}: loop over j, then k rollouts each"


@pytest.mark.parametrize("steps", [GOOD, BAD])
def test_soft_labels_estimate_the_true_value(steps):
    soft, _ = proc.mc_step_labels(TASK, EXAMPLE, steps, 3000, random.Random(7))
    exact = np.array([TASK.exact_value(EXAMPLE, steps[:j]) for j in range(len(steps) + 1)])
    assert np.allclose(soft, exact, atol=0.03), \
        f"with 3000 rollouts the soft labels should match V(prefix):\n exact {exact.round(3)}\n soft  {soft.round(3)}"


def test_hard_labels_have_false_positives():
    # After the mistake at step 3 the true value is only ~0.05, but with 64 rollouts
    # a lucky recovery almost always shows up, so the hard label calls the step "good".
    hits = 0
    for seed in range(40):
        _, hard = proc.mc_step_labels(TASK, EXAMPLE, BAD, 64, random.Random(seed))
        hits += hard[3]
    assert hits >= 30, f"hard label at the bad step should usually be 1 with k = 64 (got {hits}/40)"


# ------------------------------------------------------------ Task 4b
def test_process_rewards_example():
    r = proc.process_rewards([0.3, 0.4, 0.05, 0.0])
    assert np.allclose(r, [0.1, -0.35, -0.05]), f"expected [0.1, -0.35, -0.05], got {r}"


def test_process_rewards_telescope_and_find_the_mistake():
    v = [TASK.exact_value(EXAMPLE, BAD[:j]) for j in range(7)]
    r = proc.process_rewards(v)
    assert len(r) == 6
    assert np.sum(r) == pytest.approx(v[-1] - v[0]), "sum of step rewards = outcome - V_0"
    assert int(np.argmin(r)) + 1 == 3, "the most negative step reward should be step 3, the mistake"


# ------------------------------------------------------------ Task 5
class CountingTask(ChainTask):
    def __init__(self):
        super().__init__()
        self.calls = 0

    def policy_step(self, problem, prefix, rng):
        self.calls += 1
        return ChainTask.policy_step(problem, prefix, rng)


class ConstantTask(ChainTask):
    """A policy that always writes 0: every candidate is a duplicate."""

    def policy_step(self, problem, prefix, rng):
        rng.random()
        return 0


def exact_prm(problem, prefix):
    return ChainTask.exact_value(problem, prefix)


def test_beam_width_1_expand_1_is_plain_sampling():
    for seed in range(5):
        sol, n = proc.beam_search(TASK, EXAMPLE, exact_prm, 1, 1, random.Random(seed))
        assert sol == TASK.rollout(EXAMPLE, [], random.Random(seed)), "B = 1, M = 1 should be one ordinary sample"
        assert n == 6, "one generated step per step"


def test_beam_search_counts_generated_steps():
    task = CountingTask()
    sol, n = proc.beam_search(task, EXAMPLE, exact_prm, 3, 4, random.Random(0))
    assert len(sol) == 6
    assert n == task.calls, f"n_generated ({n}) must equal the number of policy_step calls ({task.calls})"
    assert 4 + 5 * 3 * 4 >= n >= 4 + 5 * 4, "between 1 and 3 beams expanded 4 ways at each step"


def test_beam_search_removes_duplicates():
    sol, n = proc.beam_search(ConstantTask(), EXAMPLE, exact_prm, 4, 3, random.Random(0))
    assert sol == [0] * 6
    assert n == 3 * 6, "all candidates are identical, so only one beam survives each step: 3 calls per step"


def test_beam_search_returns_top_scored_final_candidate():
    seen = []

    def recording_prm(problem, prefix):
        s = NoisyPRM(1.0, seed=3)(problem, prefix)
        seen.append((len(prefix), s, list(prefix)))
        return s

    sol, _ = proc.beam_search(TASK, EXAMPLE, recording_prm, 4, 4, random.Random(5))
    finals = [(s, p) for L, s, p in seen if L == 6]
    best = max(s for s, _ in finals)
    assert recording_prm(EXAMPLE, sol) == pytest.approx(best), "return the highest-scoring complete solution"


def test_beam_search_with_a_good_prm_beats_sampling():
    rng = random.Random(0)
    probs = [TASK.sample_problem(rng) for _ in range(150)]
    beam = np.mean([TASK.is_correct(p, proc.beam_search(TASK, p, exact_prm, 2, 4, random.Random(i))[0])
                    for i, p in enumerate(probs)])
    plain = np.mean([TASK.is_correct(p, TASK.sample_solution(p, random.Random(i))) for i, p in enumerate(probs)])
    assert plain < 0.45, f"sanity: one plain sample solves about a third of problems, got {plain:.2f}"
    assert beam > 0.9, f"beam search guided by the true values should solve > 90%, got {beam:.2f}"