methodology_core.py

from __future__ import annotations

import ast
import csv
import io
import json
import math
import re
import sys
from contextlib import redirect_stdout
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Optional, Sequence


PAIR_FAMILIES = ("cue_varied", "implementation_varied")
PROMPT_ROLES = (
    "aligned_cue",
    "conflicting_cue",
    "aligned_implementation",
    "conflicting_implementation",
)
RESPONSE_LABELS = (
    "cue_consistent",
    "execution_consistent",
    "both_consistent",
    "neither_consistent",
)

FINAL_OUTPUT_STOP_MARKERS = (
    "<|endoftext|>",
    "<|im_end|>",
    "<|end|>",
    "</s>",
    "\nHuman:",
    "\nUser:",
    "\nAssistant:",
    "\nSystem:",
    "Human:",
    "User:",
    "Assistant:",
    "System:",
    "\nWhat is the exact final printed output",
    "\nYou are given a Python",
    "\n```python",
)

FINAL_OUTPUT_PROMPT_TEMPLATE = (
    "What is the exact final printed output of this Python program? "
    "Return only the final printed output.\n\n"
    "```python\n{code}\n```\n\n"
    "Output:\n"
)

UNIT_TEST_PROMPT_TEMPLATE = (
    "You are given a Python program. Generate one concise pytest-style unit test "
    "with exactly 3 test cases for the main function defined in it. Each test case "
    "must be a single assertion of the form `assert <expression> == <literal expected value>`. "
    "Do not use bare boolean assertions such as `assert f(x)` or `assert not f(x)`; "
    "write `assert f(x) == True` or `assert f(x) == False` instead. Return exactly "
    "one fenced Python code block. End the code block with closing triple backticks.\n\n"
    "```python\n{code}\n```\n\n"
    "Unit test:\n"
)


@dataclass(frozen=True)
class PromptSpec:
    role: str
    file: str
    code: str
    execution_output: str
    cue_output: str


@dataclass(frozen=True)
class PatchDirection:
    name: str
    source_role: str
    destination_role: str
    source_output: str
    destination_output: str


@dataclass(frozen=True)
class StimulusPair:
    pair_id: str
    pair_family: str
    prompts: dict[str, PromptSpec]
    patch_directions: list[PatchDirection]


@dataclass(frozen=True)
class PatchUnit:
    patch_unit_id: int
    patch_unit_type: str
    source_positions: list[int]
    destination_positions: list[int]
    source_tokens: list[str]
    destination_tokens: list[str]
    patch_unit_text: str
    is_asymmetric: bool


@dataclass(frozen=True)
class AssertionDecision:
    assertion_text: str
    assertion_input: str
    expected_value: str
    execution_expected_value: str
    cue_expected_value: str
    is_accurate: bool
    response_label: str
    parse_status: str
    skip_reason: str


class MethodologyError(ValueError):
    def __init__(self, error_type: str, message: str):
        super().__init__(message)
        self.error_type = error_type
        self.message = message


def repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def build_prompt(code: str, template: str = FINAL_OUTPUT_PROMPT_TEMPLATE) -> str:
    return template.format(code=code.rstrip())


def execute_program_output(code: str) -> str:
    buffer = io.StringIO()
    scope = {"__builtins__": __builtins__}
    try:
        with redirect_stdout(buffer):
            exec(code, scope, scope)
    except Exception as exc:
        raise MethodologyError("assertion_execution_error", str(exc)) from exc
    lines = buffer.getvalue().splitlines()
    return "" if not lines else str(lines[-1])


def classify_behavior(candidate: str, execution_output: str, cue_output: str) -> str:
    candidate = str(candidate)
    execution_output = str(execution_output)
    cue_output = str(cue_output)
    if candidate == execution_output and candidate == cue_output:
        return "both_consistent"
    if candidate == execution_output:
        return "execution_consistent"
    if candidate == cue_output:
        return "cue_consistent"
    return "neither_consistent"


def clean_final_output_text(text: Any) -> str:
    """Normalize final-output answers before semantic classification.

    This removes response-format artifacts without trying to infer a different
    program result from prose.
    """
    value = str(text or "").replace("\r\n", "\n").replace("\r", "\n").strip()
    if not value:
        return ""

    split_at: int | None = None
    for marker in FINAL_OUTPUT_STOP_MARKERS:
        index = value.find(marker)
        if index >= 0 and (split_at is None or index < split_at):
            split_at = index
    if split_at is not None:
        value = value[:split_at].strip()

    # Prefer the contents of a complete fenced answer, including unlabeled
    # fences. Split continuation markers first so a later prompt fragment after
    # EOS cannot mask a valid answer before it.
    fence_index = value.find("```")
    prefix_before_fence = value[:fence_index].strip() if fence_index >= 0 else ""
    fenced = re.search(r"```[^\n`]*\n(.*?)```", value, flags=re.DOTALL)
    if not fenced:
        fenced = re.search(r"```(.*?)```", value, flags=re.DOTALL)
    if prefix_before_fence and (not fenced or not fenced.group(1).strip()):
        value = prefix_before_fence
    elif fenced:
        value = fenced.group(1).strip()

    value = value.replace("```", "").strip()
    value = re.sub(
        r"^(?:final\s+printed\s+output|final\s+output|printed\s+output|output|answer)\s*(?:is|:)\s*",
        "",
        value,
        flags=re.IGNORECASE,
    ).strip()

    lines = [line.strip() for line in value.splitlines() if line.strip()]
    if not lines:
        return ""
    first = lines[0]
    first = first.strip("`").strip()
    if len(first) >= 2 and first[0] == first[-1] and first[0] in {'"', "'"}:
        first = first[1:-1]
    return first.strip()


def classify_cleaned_behavior(candidate: Any, execution_output: Any, cue_output: Any) -> str:
    return classify_behavior(
        clean_final_output_text(candidate),
        str(execution_output),
        str(cue_output),
    )


def final_output_correctness_flags(response_label: str) -> dict[str, bool]:
    """Map final-output behavior labels onto the execution-grounded task labels."""
    return {
        "is_correct": response_label in {"execution_consistent", "both_consistent"},
        "is_incorrect_misleading": response_label == "cue_consistent",
    }


def prompt_roles_for_family(pair_family: str) -> tuple[str, str]:
    if pair_family == "cue_varied":
        return ("aligned_cue", "conflicting_cue")
    if pair_family == "implementation_varied":
        return ("aligned_implementation", "conflicting_implementation")
    raise MethodologyError("invalid_pair_family", f"Unsupported pair_family={pair_family!r}")


def source_value_for_family(pair_family: str, prompt: PromptSpec) -> str:
    if pair_family == "cue_varied":
        return prompt.cue_output
    if pair_family == "implementation_varied":
        return prompt.execution_output
    raise MethodologyError("invalid_pair_family", f"Unsupported pair_family={pair_family!r}")


def derive_patch_directions(pair_family: str, prompts: dict[str, PromptSpec]) -> list[PatchDirection]:
    first_role, second_role = prompt_roles_for_family(pair_family)
    first = prompts[first_role]
    second = prompts[second_role]
    return [
        PatchDirection(
            name=f"{first_role}_to_{second_role}",
            source_role=first_role,
            destination_role=second_role,
            source_output=source_value_for_family(pair_family, first),
            destination_output=source_value_for_family(pair_family, second),
        ),
        PatchDirection(
            name=f"{second_role}_to_{first_role}",
            source_role=second_role,
            destination_role=first_role,
            source_output=source_value_for_family(pair_family, second),
            destination_output=source_value_for_family(pair_family, first),
        ),
    ]


def validate_pair_shape(pair_id: str, pair_family: str, prompts: dict[str, PromptSpec]) -> None:
    expected_roles = set(prompt_roles_for_family(pair_family))
    actual_roles = set(prompts)
    if actual_roles != expected_roles:
        raise MethodologyError(
            "invalid_prompt_role",
            f"{pair_id} has prompt roles {sorted(actual_roles)}, expected {sorted(expected_roles)}",
        )

    if pair_family == "cue_varied":
        aligned = prompts["aligned_cue"]
        conflicting = prompts["conflicting_cue"]
        checks = (
            aligned.execution_output == conflicting.execution_output,
            aligned.execution_output == aligned.cue_output,
            conflicting.execution_output != conflicting.cue_output,
            aligned.cue_output != conflicting.cue_output,
        )
    else:
        aligned = prompts["aligned_implementation"]
        conflicting = prompts["conflicting_implementation"]
        checks = (
            aligned.cue_output == conflicting.cue_output,
            aligned.execution_output == aligned.cue_output,
            conflicting.execution_output != conflicting.cue_output,
            aligned.execution_output != conflicting.execution_output,
        )
    if not all(checks):
        raise MethodologyError("invalid_contrast", f"{pair_id} violates {pair_family} contrast rules")


def load_stimulus_pair(pair_dir: Path) -> StimulusPair:
    meta_path = pair_dir / "meta.json"
    if not meta_path.exists():
        raise MethodologyError("invalid_pair_family", f"Missing metadata: {meta_path}")
    meta = json.loads(meta_path.read_text(encoding="utf-8"))
    pair_id = str(meta.get("pair_id") or pair_dir.name)
    pair_family = str(meta.get("pair_family") or "")
    if pair_family not in PAIR_FAMILIES:
        raise MethodologyError("invalid_pair_family", f"{pair_id} has pair_family={pair_family!r}")

    prompt_payload = meta.get("prompts")
    if not isinstance(prompt_payload, dict):
        raise MethodologyError("invalid_prompt_role", f"{pair_id} metadata lacks prompts")

    prompts: dict[str, PromptSpec] = {}
    for role, payload in prompt_payload.items():
        if role not in PROMPT_ROLES:
            raise MethodologyError("invalid_prompt_role", f"{pair_id} has prompt_role={role!r}")
        if not isinstance(payload, dict):
            raise MethodologyError("invalid_prompt_role", f"{pair_id} prompt payload is invalid")
        file_name = str(payload.get("file") or "")
        if not file_name:
            raise MethodologyError("invalid_prompt_role", f"{pair_id}:{role} lacks file")
        code_path = pair_dir / file_name
        code = code_path.read_text(encoding="utf-8")
        execution_output = str(payload.get("execution_output") or "")
        cue_output = str(payload.get("cue_output") or "")
        if cue_output == "":
            raise MethodologyError("missing_cue_output", f"{pair_id}:{role} lacks cue_output")
        observed_output = execute_program_output(code)
        if observed_output != execution_output:
            raise MethodologyError(
                "runtime_metadata_mismatch",
                f"{pair_id}:{role} metadata={execution_output!r}, observed={observed_output!r}",
            )
        prompts[role] = PromptSpec(
            role=role,
            file=file_name,
            code=code,
            execution_output=execution_output,
            cue_output=cue_output,
        )

    validate_pair_shape(pair_id, pair_family, prompts)
    return StimulusPair(
        pair_id=pair_id,
        pair_family=pair_family,
        prompts=prompts,
        patch_directions=derive_patch_directions(pair_family, prompts),
    )


def iter_pair_dirs(stimulus_root: Path) -> list[Path]:
    roots: list[Path] = []
    for family in PAIR_FAMILIES:
        family_root = stimulus_root / family
        if family_root.exists():
            roots.extend(path for path in family_root.iterdir() if path.is_dir())
    return sorted(roots)


def pair_dir_display_name(pair_dir: Path, stimulus_root: Path) -> str:
    try:
        return pair_dir.relative_to(stimulus_root).as_posix()
    except ValueError:
        return pair_dir.name


def _candidate_pair_dirs(stimulus_root: Path, pair_family: Optional[str]) -> list[Path]:
    if pair_family is not None and pair_family not in PAIR_FAMILIES:
        raise MethodologyError("invalid_pair_family", f"Unsupported pair_family={pair_family!r}")
    pair_dirs = iter_pair_dirs(stimulus_root)
    if pair_family is None:
        return pair_dirs
    return [path for path in pair_dirs if path.parent.name == pair_family]


def _match_pair_dirs(stimulus_root: Path, pair_dirs: Sequence[Path], name: str) -> list[Path]:
    direct = stimulus_root / name
    if direct.exists() and direct.is_dir():
        return [direct]

    matches: list[Path] = []
    for path in pair_dirs:
        try:
            pair = load_stimulus_pair(path)
        except MethodologyError:
            continue
        identifiers = {
            pair.pair_id,
            path.name,
            pair_dir_display_name(path, stimulus_root),
        }
        if name in identifiers:
            matches.append(path)
    return sorted(matches)


def load_pair_dirs_from_args(
    stimulus_root: Path,
    names: Optional[Sequence[str]],
    use_all: bool,
    pair_family: Optional[str] = None,
) -> list[Path]:
    candidate_dirs = _candidate_pair_dirs(stimulus_root, pair_family)
    if use_all or not names:
        return candidate_dirs

    selected: list[Path] = []
    for name in names:
        matches = _match_pair_dirs(stimulus_root, candidate_dirs, name)
        if not matches:
            raise MethodologyError("invalid_pair_family", f"Pair not found: {name}")
        if len(matches) > 1:
            labels = ", ".join(pair_dir_display_name(path, stimulus_root) for path in matches)
            raise MethodologyError(
                "ambiguous_pair_id",
                f"Pair name {name!r} matches multiple stimulus pairs: {labels}. "
                "Use --pair-family or a family-qualified path.",
            )
        selected.append(matches[0])
    return selected


def extract_completion(full_text: str, prompt: str) -> str:
    return full_text[len(prompt) :].strip() if full_text.startswith(prompt) else full_text.strip()


def first_nonempty_line(text: str) -> str:
    for line in str(text).splitlines():
        stripped = line.strip()
        if stripped:
            return stripped
    return ""


def extract_model_output(completion: str) -> str:
    if has_complete_fenced_code_block(completion):
        return extract_first_fenced_code_block(completion)
    return first_nonempty_line(completion)


def extract_first_fenced_code_block(text: str) -> str:
    match = re.search(r"```[^\n`]*\n(.*?)```", str(text), flags=re.DOTALL)
    if not match:
        match = re.search(r"```(.*?)```", str(text), flags=re.DOTALL)
    return match.group(1).strip() if match else str(text).strip()


def has_complete_fenced_code_block(text: str) -> bool:
    opener = str(text).find("```")
    return opener >= 0 and str(text).find("```", opener + 3) >= 0


def literal_source_value(source: str) -> tuple[Any, str]:
    try:
        return ast.literal_eval(source), ""
    except Exception as exc:
        return None, str(exc)


def value_key(value: Any) -> str:
    return repr(value)


def collect_assertion_records(test_source: str) -> list[dict[str, str]]:
    try:
        tree = ast.parse(test_source)
    except SyntaxError as exc:
        raise MethodologyError("assertion_parse_error", str(exc)) from exc

    records: list[dict[str, str]] = []

    def visit_statements(statements: Sequence[ast.stmt], setup_lines: list[str]) -> None:
        for statement in statements:
            if isinstance(statement, ast.Assert) and isinstance(statement.test, ast.Compare):
                compare = statement.test
                if len(compare.ops) == 1 and isinstance(compare.ops[0], ast.Eq) and len(compare.comparators) == 1:
                    expr_source = ast.get_source_segment(test_source, compare.left) or ast.unparse(compare.left)
                    expected_source = ast.get_source_segment(test_source, compare.comparators[0]) or ast.unparse(compare.comparators[0])
                    assertion_text = ast.get_source_segment(test_source, statement) or ast.unparse(statement)
                    records.append(
                        {
                            "assertion_text": assertion_text,
                            "assertion_input": expr_source,
                            "expected_value": expected_source,
                            "setup_source": "\n".join(setup_lines),
                        }
                    )
                    continue
            if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef)):
                visit_statements(statement.body, [])
            else:
                setup_lines.append(ast.get_source_segment(test_source, statement) or ast.unparse(statement))

    visit_statements(tree.body, [])
    return records


def evaluate_assertion_expression(code: str, setup_source: str, expression_source: str) -> Any:
    scope = {
        "__builtins__": __builtins__,
        "io": io,
        "StringIO": io.StringIO,
        "sys": sys,
    }
    try:
        buffer = io.StringIO()
        with redirect_stdout(buffer):
            exec(code, scope, scope)
            if setup_source.strip():
                exec(setup_source, scope, scope)
            return eval(expression_source, scope, scope)
    except Exception as exc:
        raise MethodologyError("assertion_execution_error", str(exc)) from exc


def analyze_unit_test_output(
    pair: StimulusPair,
    prompt_role: str,
    test_source: str,
) -> list[AssertionDecision]:
    prompt = pair.prompts[prompt_role]
    cue_role = "aligned_implementation" if pair.pair_family == "implementation_varied" else prompt_role
    cue_prompt = pair.prompts[cue_role]
    decisions: list[AssertionDecision] = []
    records = collect_assertion_records(test_source)
    for record in records:
        expected_value, expected_error = literal_source_value(record["expected_value"])
        if expected_error:
            decisions.append(
                AssertionDecision(
                    assertion_text=record["assertion_text"],
                    assertion_input=record["assertion_input"],
                    expected_value=record["expected_value"],
                    execution_expected_value="",
                    cue_expected_value="",
                    is_accurate=False,
                    response_label="neither_consistent",
                    parse_status="skipped",
                    skip_reason="assertion_parse_error",
                )
            )
            continue
        try:
            execution_value = evaluate_assertion_expression(
                prompt.code,
                record["setup_source"],
                record["assertion_input"],
            )
            cue_value = evaluate_assertion_expression(
                cue_prompt.code,
                record["setup_source"],
                record["assertion_input"],
            )
        except MethodologyError as exc:
            decisions.append(
                AssertionDecision(
                    assertion_text=record["assertion_text"],
                    assertion_input=record["assertion_input"],
                    expected_value=value_key(expected_value),
                    execution_expected_value="",
                    cue_expected_value="",
                    is_accurate=False,
                    response_label="neither_consistent",
                    parse_status="skipped",
                    skip_reason=exc.error_type,
                )
            )
            continue
        expected_key = value_key(expected_value)
        execution_key = value_key(execution_value)
        cue_key = value_key(cue_value)
        decisions.append(
            AssertionDecision(
                assertion_text=record["assertion_text"],
                assertion_input=record["assertion_input"],
                expected_value=expected_key,
                execution_expected_value=execution_key,
                cue_expected_value=cue_key,
                is_accurate=expected_key == execution_key,
                response_label=classify_behavior(expected_key, execution_key, cue_key),
                parse_status="parsed",
                skip_reason="",
            )
        )
    return decisions


def first_divergent_token_ids(source_ids: Sequence[int], destination_ids: Sequence[int]) -> Optional[tuple[int, int, int, list[int]]]:
    for index, (source_token, destination_token) in enumerate(zip(source_ids, destination_ids)):
        if int(source_token) != int(destination_token):
            return index, int(source_token), int(destination_token), [int(x) for x in source_ids[:index]]
    return None


def margin_values(patched_margin: float, baseline_margin: float, source_margin: float) -> dict[str, float]:
    denominator = float(source_margin) - float(baseline_margin)
    if abs(denominator) <= 1e-9:
        raise MethodologyError(
            "no_source_destination_margin_difference",
            "baseline_margin and source_margin are indistinguishable",
        )
    margin_delta = float(patched_margin) - float(baseline_margin)
    signed_recovery = margin_delta / denominator
    distance_recovery = 1.0 - abs(float(patched_margin) - float(source_margin)) / abs(denominator)
    return {
        "margin_delta": margin_delta,
        "signed_recovery": signed_recovery,
        "distance_recovery": distance_recovery,
    }


def write_csv(path: Path, rows: Sequence[dict[str, Any]], fieldnames: Optional[Sequence[str]] = None) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if fieldnames is None:
        names: list[str] = []
        for row in rows:
            for key in row:
                if key not in names:
                    names.append(key)
        fieldnames = names
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(fieldnames))
        writer.writeheader()
        for row in rows:
            writer.writerow({key: row.get(key, "") for key in fieldnames})


def read_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def finite_mean(values: Iterable[Any]) -> str | float:
    parsed = []
    for value in values:
        try:
            number = float(value)
        except (TypeError, ValueError):
            continue
        if math.isfinite(number):
            parsed.append(number)
    return "" if not parsed else sum(parsed) / len(parsed)


def finite_median(values: Iterable[Any]) -> str | float:
    parsed = []
    for value in values:
        try:
            number = float(value)
        except (TypeError, ValueError):
            continue
        if math.isfinite(number):
            parsed.append(number)
    if not parsed:
        return ""
    parsed.sort()
    mid = len(parsed) // 2
    if len(parsed) % 2:
        return parsed[mid]
    return (parsed[mid - 1] + parsed[mid]) / 2.0