#!/usr/bin/env python3
"""Deterministic protocol simulator for long-context position sweeps.

This does not call or emulate a language model. It generates controlled records
with one answer-bearing document, moves that document across fixed positions,
and uses hash-derived Bernoulli draws to demonstrate analysis and reporting.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import random
from dataclasses import dataclass


SEED = 62
LENGTHS = (20, 40, 80)
RELATIVE_POSITIONS = (0.0, 0.25, 0.5, 0.75, 1.0)
EXAMPLES_PER_CELL = 800


@dataclass(frozen=True)
class Example:
    example_id: int
    num_documents: int
    gold_index: int
    answer: str
    documents: tuple[str, ...]


def position_to_index(relative_position: float, num_documents: int) -> int:
    return round(relative_position * (num_documents - 1))


def make_example(example_id: int, num_documents: int, gold_index: int) -> Example:
    rng = random.Random((SEED << 24) + example_id * 1009 + num_documents)
    answer = f"VALUE-{example_id:04d}"
    documents = [f"DOC-{i:03d}: filler-{rng.randrange(10**9):09d}" for i in range(num_documents)]
    documents[gold_index] = f"DOC-{gold_index:03d}: TARGET={answer}"
    assert sum("TARGET=" in document for document in documents) == 1
    return Example(example_id, num_documents, gold_index, answer, tuple(documents))


def deterministic_uniform(*parts: object) -> float:
    payload = "|".join(map(str, parts)).encode("utf-8")
    integer = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")
    return integer / 2**64


def success_probability(num_documents: int, gold_index: int, profile: str) -> float:
    relative = gold_index / (num_documents - 1)
    if profile == "flat-control":
        return 0.76 - 0.0008 * (num_documents - 20)
    edge_strength = (2.0 * abs(relative - 0.5)) ** 1.6
    return 0.56 + 0.39 * edge_strength - 0.0012 * (num_documents - 20)


def predict(example: Example, profile: str) -> str:
    probability = success_probability(example.num_documents, example.gold_index, profile)
    draw = deterministic_uniform(SEED, profile, example.example_id, example.num_documents, example.gold_index)
    if draw < probability:
        return example.answer
    wrong_index = int(deterministic_uniform("wrong", example.example_id, example.gold_index) * example.num_documents)
    return f"DISTRACTOR-{wrong_index:03d}"


def run_sweep(profile: str) -> dict[int, dict[float, float]]:
    results: dict[int, dict[float, float]] = {}
    for num_documents in LENGTHS:
        row: dict[float, float] = {}
        for relative_position in RELATIVE_POSITIONS:
            gold_index = position_to_index(relative_position, num_documents)
            correct = 0
            for example_id in range(EXAMPLES_PER_CELL):
                example = make_example(example_id, num_documents, gold_index)
                correct += predict(example, profile) == example.answer
            row[relative_position] = correct / EXAMPLES_PER_CELL
        results[num_documents] = row
    return results


def summarize(results: dict[int, dict[float, float]]) -> dict[int, dict[str, float]]:
    summary = {}
    for num_documents, row in results.items():
        values = list(row.values())
        edge_mean = (values[0] + values[-1]) / 2
        summary[num_documents] = {
            "macro_accuracy": sum(values) / len(values),
            "middle_accuracy": values[len(values) // 2],
            "edge_mean": edge_mean,
            "position_gap": edge_mean - values[len(values) // 2],
        }
    return summary


def print_table(profile: str, results: dict[int, dict[float, float]]) -> None:
    summary = summarize(results)
    labels = [f"p={position:.2f}" for position in RELATIVE_POSITIONS]
    print(f"profile={profile} seed={SEED} examples_per_cell={EXAMPLES_PER_CELL}")
    print("documents\t" + "\t".join(labels) + "\tmacro\tposition_gap")
    for num_documents, row in results.items():
        values = [row[position] for position in RELATIVE_POSITIONS]
        metrics = summary[num_documents]
        rendered = "\t".join(f"{value:.3f}" for value in values)
        print(f"{num_documents}\t{rendered}\t{metrics['macro_accuracy']:.3f}\t{metrics['position_gap']:.3f}")


def checks(biased: dict[int, dict[float, float]], flat: dict[int, dict[float, float]]) -> None:
    biased_summary = summarize(biased)
    flat_summary = summarize(flat)
    assert all(metrics["position_gap"] > 0.25 for metrics in biased_summary.values())
    assert all(abs(metrics["position_gap"]) < 0.08 for metrics in flat_summary.values())
    assert biased_summary[80]["macro_accuracy"] < biased_summary[20]["macro_accuracy"]
    for num_documents, row in biased.items():
        assert len(row) == len(RELATIVE_POSITIONS)
        assert all(math.isfinite(value) and 0.0 <= value <= 1.0 for value in row.values())
        for relative_position in RELATIVE_POSITIONS:
            index = position_to_index(relative_position, num_documents)
            example = make_example(0, num_documents, index)
            assert example.documents[index].endswith(example.answer)
    print("checks=passed")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--json", action="store_true", help="also print machine-readable summaries")
    parser.add_argument("--check-only", action="store_true", help="run deterministic assertions")
    args = parser.parse_args()

    biased = run_sweep("u-shaped")
    flat = run_sweep("flat-control")
    if not args.check_only:
        print_table("u-shaped", biased)
        print()
        print_table("flat-control", flat)
    checks(biased, flat)
    if args.json:
        print(json.dumps({"u-shaped": summarize(biased), "flat-control": summarize(flat)}, indent=2))


if __name__ == "__main__":
    main()