#!/usr/bin/env python3
"""Dependency-free protocol test for self-consistency decoding.

This script does not call a language model and does not reproduce paper scores.
It simulates answer samples with controlled error distributions so that answer
normalization, majority voting, repeated runs, and cost/accuracy trade-offs are
inspectable on a CPU.
"""

from __future__ import annotations

import argparse
import math
import random
import re
import statistics
from collections import Counter
from dataclasses import dataclass


SEED = 59
SAMPLE_COUNTS = (1, 3, 5, 9, 17, 33)


@dataclass(frozen=True)
class Scenario:
    name: str
    # The first label is correct; the rest are distinct wrong answers.
    probabilities: tuple[float, ...]


SCENARIOS = (
    Scenario("diverse_errors", (0.54, 0.16, 0.15, 0.15)),
    Scenario("systematic_error", (0.42, 0.52, 0.04, 0.02)),
)


def normalize_answer(text: str) -> str | None:
    """Extract and canonicalize a numeric final answer from common formats."""
    patterns = (
        r"####\s*([-+$]?[0-9][0-9,]*(?:\.[0-9]+)?)",
        r"answer\s+is\s*[:=]?\s*([-+$]?[0-9][0-9,]*(?:\.[0-9]+)?)",
    )
    lowered = text.lower()
    match = next((m for pattern in patterns if (m := re.search(pattern, lowered))), None)
    if match is None:
        return None
    value = match.group(1).replace(",", "").replace("$", "")
    try:
        number = float(value)
    except ValueError:
        return None
    return str(int(number)) if number.is_integer() else format(number, ".12g")


def sample_label(rng: random.Random, probabilities: tuple[float, ...]) -> int:
    draw = rng.random()
    cumulative = 0.0
    for label, probability in enumerate(probabilities):
        cumulative += probability
        if draw < cumulative:
            return label
    return len(probabilities) - 1


def majority_vote(labels: list[int]) -> tuple[int, float]:
    """Return winner and winning vote share; earliest sample breaks a tie."""
    counts = Counter(labels)
    best_count = max(counts.values())
    winner = next(label for label in labels if counts[label] == best_count)
    return winner, best_count / len(labels)


def evaluate(
    scenario: Scenario,
    samples: int,
    questions: int,
    repeats: int,
    seed: int,
) -> tuple[float, float, float]:
    accuracies: list[float] = []
    agreements: list[float] = []
    for repeat in range(repeats):
        rng = random.Random(seed + repeat * 10_000 + samples)
        correct = 0
        vote_shares = []
        for _ in range(questions):
            labels = [sample_label(rng, scenario.probabilities) for _ in range(samples)]
            winner, vote_share = majority_vote(labels)
            correct += winner == 0
            vote_shares.append(vote_share)
        accuracies.append(correct / questions)
        agreements.append(statistics.mean(vote_shares))
    return statistics.mean(accuracies), statistics.pstdev(accuracies), statistics.mean(agreements)


def exact_binary_majority_probability(single_path_accuracy: float, samples: int) -> float:
    """Probability of a correct strict majority for odd n and binary answers."""
    threshold = samples // 2 + 1
    return sum(
        math.comb(samples, k)
        * single_path_accuracy**k
        * (1.0 - single_path_accuracy) ** (samples - k)
        for k in range(threshold, samples + 1)
    )


def run(questions: int, repeats: int, check_only: bool) -> None:
    parser_cases = {
        "Reasoning... The answer is 1,200.": "1200",
        "work\n#### 1200": "1200",
        "The answer is $1200.0": "1200",
        "no final marker": None,
    }
    for text, expected in parser_cases.items():
        assert normalize_answer(text) == expected

    print(f"seed={SEED} questions={questions} repeats={repeats}")
    print("scenario samples accuracy std vote_share completions token_budget@128")
    results: dict[str, dict[int, float]] = {}
    for scenario in SCENARIOS:
        results[scenario.name] = {}
        for samples in SAMPLE_COUNTS:
            accuracy, std, vote_share = evaluate(
                scenario, samples, questions, repeats, SEED
            )
            results[scenario.name][samples] = accuracy
            completions = questions * repeats * samples
            print(
                f"{scenario.name:16s} {samples:2d} {accuracy:.4f} {std:.4f} "
                f"{vote_share:.4f} {completions:8d} {samples * 128:5d}"
            )

    print("binary_independent_theory p=0.60")
    for samples in SAMPLE_COUNTS:
        exact = exact_binary_majority_probability(0.60, samples)
        print(f"samples={samples:2d} majority_accuracy={exact:.4f}")

    assert results["diverse_errors"][17] > results["diverse_errors"][1] + 0.20
    assert results["systematic_error"][17] < results["systematic_error"][1] - 0.07
    assert exact_binary_majority_probability(0.60, 33) > 0.85
    if check_only:
        print("check-only: PASS")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--questions", type=int, default=2000)
    parser.add_argument("--repeats", type=int, default=20)
    parser.add_argument("--check-only", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    run(args.questions, args.repeats, args.check_only)