"""Tests for the composition experiment's pure pieces and its execution verifier (prereg v3 §7)."""

import numpy as np
import pytest

from llm.compose import predicted_composition, score_composed, score_gsm8k, union_exceedance
from llm.compose_data import ProgTask
from llm.execute import ExecResult, extract_code, numeric_match, run_solution, verify_program
from llm.tasks import Task


# ---------------------------------------------------------------- execution verifier


def test_executes_and_matches_reference():
    ok, res = verify_program("```python\ndef solution():\n    return 6*7\n```", 42.0)
    assert ok and res.ok and res.value == 42.0 and res.status == "ok"


def test_rejects_wrong_answer_but_still_ran():
    ok, res = verify_program("def solution():\n    return 41", 42.0)
    assert not ok and res.ok and res.value == 41.0


@pytest.mark.parametrize("code,status", [
    ("def solution():\n    while True: pass", "timeout"),
    ("def solution():\n    return len([0]*10**9)", "memory"),
    ("x = 1", "no_solution"),
    ("def solution(:\n  return 1", "error:SyntaxError"),
    ("def solution():\n    return solution()", "recursion"),
    ("", "no_code"),
])
def test_hazards_are_contained(code, status):
    assert run_solution(code, timeout_s=8.0).status == status


def test_sandbox_blocks_writes_outside_jail_and_network():
    assert run_solution(
        "def solution():\n    open('/tmp/_llm_compose_escape','w').write('x'); return 1"
    ).status == "error:PermissionError"
    assert run_solution(
        "import socket\ndef solution():\n    socket.socket().connect(('1.1.1.1',80)); return 1"
    ).status == "error:PermissionError"
    import os
    assert not os.path.exists("/tmp/_llm_compose_escape")


def test_sandbox_allows_writes_inside_its_own_jail():
    # scratch files are fine — the jail is a fresh temp dir destroyed after the run
    assert run_solution("def solution():\n    open('s.txt','w').write('x'); return 1").ok


def test_extract_code_prefers_last_fence_and_drops_prose():
    got = extract_code("Sure:\n```python\ndef solution():\n    return 1\n```\nHope that helps!")
    assert got == "def solution():\n    return 1"
    assert extract_code("def solution():\n    return 2") == "def solution():\n    return 2"


def test_numeric_match_tolerances():
    assert numeric_match(1e6, 1e6 + 1)            # relative tolerance
    assert numeric_match(0.0, 0.0)
    assert not numeric_match(None, 1.0)
    assert not numeric_match(1.0, 2.0)


def test_execution_is_deterministic():
    code = "def solution():\n    return sum(range(1000))"
    assert len({run_solution(code).value for _ in range(3)}) == 1


# ---------------------------------------------------------------- scoring and the prediction


def test_score_composed_counts_correct_and_executable_separately():
    tasks = [ProgTask("p", 42.0, 0), ProgTask("p", 42.0, 1), ProgTask("p", 42.0, 2)]
    comps = ["def solution():\n    return 42",      # correct
             "def solution():\n    return 7",       # ran, wrong
             "def solution(:\n"]                    # did not run
    acc, ok, ran = score_composed(comps, tasks)
    assert acc == pytest.approx(1 / 3) and ok.tolist() == [True, False, False]
    assert ran == pytest.approx(2 / 3)


def test_score_gsm8k_reads_after_the_hash_marker():
    tasks = [Task("math", "p", "18"), Task("math", "p", "7")]
    assert score_gsm8k(["reasoning blah\n#### 18", "the answer is 9"], tasks) == pytest.approx(0.5)


def test_union_exceedance_is_solved_by_merge_and_no_parent():
    merged = np.array([True, True, True, False])
    pa = np.array([True, False, False, False])
    pb = np.array([False, True, False, True])
    assert union_exceedance(merged, [pa, pb]) == pytest.approx(0.25)   # only item 2
    assert union_exceedance(merged, [merged]) == 0.0


def test_predicted_composition_is_anchored_at_generation_zero():
    q_m = np.array([0.9, 0.8, 0.7]); q_c = np.array([0.9, 0.85, 0.8])
    rho = np.array([0.2, 0.4, 0.6])
    pred = predicted_composition(q_m, q_c, rho, observed0=0.30)
    assert pred[0] == pytest.approx(0.30)                     # one free scale, fixed at t=0
    assert pred[1] > pred[2]                                  # decays with q and rho
    # decays faster than either parent alone: product form plus the decorrelation term
    assert pred[2] / pred[0] < min(q_m[2] / q_m[0], q_c[2] / q_c[0])


def test_predicted_composition_handles_degenerate_start():
    pred = predicted_composition(np.array([0.0, 0.0]), np.array([0.0, 0.0]),
                                 np.array([0.0, 0.0]), observed0=0.0)
    assert np.all(np.isfinite(pred))