#!/usr/bin/env python3
"""Run the trolley-problem state and questions through a Laya-CoreML checkpoint.

Reuses the exact state.json and questions.json from the TypeSafe/Jev walkthrough
(content/code/2026/typesafe-trolley-problem/) and verifies their SHA-256 hashes
before sending anything, so "identical prompt" is provable rather than asserted.

Usage:
    python bench_trolley.py \
        --model-dir models/typed-decisions \
        --checkpoint-label aac6fef/laya-typed-decisions-coreml \
        --revision 28d24fa8d67a3264556b23391ec6c3fd98573056 \
        --state state.json --questions questions.json \
        --jev-response jev-response.json \
        --out response-typed-decisions.json --timings timings-typed-decisions.json

    python bench_trolley.py --sweep \
        --model-dir models/typed-decisions \
        --checkpoint-label aac6fef/laya-typed-decisions-coreml \
        --revision 28d24fa8d67a3264556b23391ec6c3fd98573056 \
        --state state.json --questions questions.json --timings sweep.json
"""
from __future__ import annotations

import argparse
import hashlib
import json
import statistics
import time
from pathlib import Path
from typing import Any, Callable

# SHA-256 of content/code/2026/typesafe-trolley-problem/state.json and
# questions.json as recorded in the parent post's code walkthrough. A mismatch
# means the prompt sent here is not the prompt Jev answered.
EXPECTED_STATE_SHA256 = "bdbb63db0ac4575beb381ef1d902912d16f3baa796aeb47ce57145c5a82cc6f0"
EXPECTED_QUESTIONS_SHA256 = "0cd986e420d813e3c3b95cf55200a131a1a0e1f28a0fb48a6662c3b2fd88fa08"

WARMUP_CALLS = 10
MEASURED_CALLS = 100
SWEEP_CALLS = 50
COMPUTE_UNIT_CHOICES = ["cpu_gpu", "cpu_ne", "cpu", "all"]


def sha256_of_file(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def verify_prompt_hashes(state_path: Path, questions_path: Path) -> None:
    """Raise if the local prompt files do not match the recorded Jev run."""
    state_hash = sha256_of_file(state_path)
    questions_hash = sha256_of_file(questions_path)
    if state_hash != EXPECTED_STATE_SHA256:
        raise SystemExit(
            f"state.json hash mismatch: got {state_hash}, "
            f"expected {EXPECTED_STATE_SHA256}. Refusing to run on a different prompt."
        )
    if questions_hash != EXPECTED_QUESTIONS_SHA256:
        raise SystemExit(
            f"questions.json hash mismatch: got {questions_hash}, "
            f"expected {EXPECTED_QUESTIONS_SHA256}. Refusing to run on a different prompt."
        )


def percentiles(samples_ms: list[float]) -> dict[str, float]:
    """P50/P95 using nearest-index interpolation (matches upstream BENCHMARKS.md)."""
    if not samples_ms:
        raise ValueError("percentiles() requires at least one sample")
    ordered = sorted(samples_ms)
    n = len(ordered)

    def pct(p: float) -> float:
        idx = max(0, min(n - 1, int(round(p * (n - 1)))))
        return ordered[idx]

    return {
        "p50_ms": pct(0.50),
        "p95_ms": pct(0.95),
        "mean_ms": statistics.mean(ordered),
        "min_ms": ordered[0],
        "max_ms": ordered[-1],
        "n": n,
    }


def time_call_ms(fn: Callable[[], Any]) -> tuple[float, Any]:
    start = time.perf_counter()
    result = fn()
    elapsed_ms = (time.perf_counter() - start) * 1000.0
    return elapsed_ms, result


def benchmark_question(
    predict: Callable[[dict], Any],
    state: dict,
    question_key: str,
    question: dict,
    warmup: int = WARMUP_CALLS,
    measured: int = MEASURED_CALLS,
) -> dict:
    """Warm up, then time `measured` sequential single-question calls."""
    last_result = None
    for _ in range(warmup):
        last_result = predict(state, {question_key: question})

    samples_ms: list[float] = []
    for _ in range(measured):
        elapsed_ms, last_result = time_call_ms(
            lambda: predict(state, {question_key: question})
        )
        samples_ms.append(elapsed_ms)

    stats = percentiles(samples_ms)
    stats["question"] = question_key
    stats["last_answer"] = last_result["answers"][question_key] if last_result else None
    return stats


def benchmark_all_questions(
    predict: Callable[[dict], Any],
    state: dict,
    questions: dict[str, dict],
) -> dict:
    per_question = {}
    for key, question in questions.items():
        per_question[key] = benchmark_question(predict, state, key, question)
    return per_question


def benchmark_sequential_pass(
    predict: Callable[[dict], Any],
    state: dict,
    questions: dict[str, dict],
    repeats: int = 20,
) -> dict:
    """Time one full pass through all questions, in insertion order, per call.

    This is the number that compares directly to Jev's single 271 ms round trip:
    eleven sequential predict() calls against one state, timed end to end.
    """
    # One untimed warmup pass so the model and caches are hot.
    for key, question in questions.items():
        predict(state, {key: question})

    pass_times_ms: list[float] = []
    answers_by_pass: list[dict] = []
    for _ in range(repeats):
        start = time.perf_counter()
        answers = {}
        for key, question in questions.items():
            result = predict(state, {key: question})
            answers[key] = result["answers"][key]
        pass_times_ms.append((time.perf_counter() - start) * 1000.0)
        answers_by_pass.append(answers)

    stats = percentiles(pass_times_ms)
    stats["repeats"] = repeats
    stats["last_answers"] = answers_by_pass[-1]
    return stats


def compute_unit_sweep(
    load_agent: Callable[[str], Any],
    checkpoint: str,
    state: dict,
    question_key: str,
    question: dict,
    calls: int = SWEEP_CALLS,
) -> dict:
    """Time one question across every compute_units setting."""
    results = {}
    for unit in COMPUTE_UNIT_CHOICES:
        agent = load_agent(checkpoint, compute_units=unit)
        for _ in range(WARMUP_CALLS):
            agent.predict(state, {question_key: question})
        samples_ms: list[float] = []
        for _ in range(calls):
            elapsed_ms, _ = time_call_ms(
                lambda: agent.predict(state, {question_key: question})
            )
            samples_ms.append(elapsed_ms)
        results[unit] = percentiles(samples_ms)
    return results


def compare_to_jev(laya_answers: dict, jev_response_path: Path) -> dict:
    """Line up Laya's answers against the recorded Jev response, question by
    question. Returns per-question agreement for choice questions and score
    deltas for score questions, plus which questions both models were least
    confident about.
    """
    jev = json.loads(jev_response_path.read_text())
    jev_answers = jev["answers"]

    comparison = {}
    for key, laya_answer in laya_answers.items():
        jev_answer = jev_answers.get(key)
        if jev_answer is None:
            continue
        entry: dict[str, Any] = {
            "type": jev_answer["type"],
            "jev_confidence": jev_answer.get("confidence"),
            "laya_confidence": laya_answer.get("confidence"),
        }
        if jev_answer["type"] == "choice":
            entry["jev_choice"] = jev_answer["choice"]
            entry["laya_choice"] = laya_answer.get("choice")
            entry["agree"] = jev_answer["choice"] == laya_answer.get("choice")
        else:
            entry["jev_score"] = jev_answer["score"]
            entry["laya_score"] = laya_answer.get("score")
            if laya_answer.get("score") is not None:
                entry["score_delta"] = abs(jev_answer["score"] - laya_answer["score"])
        comparison[key] = entry

    # Questions where a model reported low confidence (<0.3), matching the
    # parent post's threshold for "confidence collapses". Tracked per model
    # rather than as one combined list: Jev's confidence spans roughly 0.16
    # to 1.0, so a 0.3 cutoff picks out a genuine minority. Laya's confidence
    # runs an order of magnitude lower across nearly every question (the two
    # models compute confidence differently), so its list may cover most or
    # all questions rather than a distinguishing subset.
    jev_low_confidence = [
        key
        for key, entry in comparison.items()
        if entry.get("jev_confidence") is not None and entry["jev_confidence"] < 0.3
    ]
    laya_low_confidence = [
        key
        for key, entry in comparison.items()
        if entry.get("laya_confidence") is not None and entry["laya_confidence"] < 0.3
    ]

    return {
        "per_question": comparison,
        "jev_low_confidence": jev_low_confidence,
        "laya_low_confidence": laya_low_confidence,
    }


def _load_json(path: Path) -> dict:
    return json.loads(path.read_text())


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model-dir",
        required=True,
        help="Local directory holding the mirrored Core ML bundle "
        "(loaded with local_files_only=True, no network at runtime)",
    )
    parser.add_argument("--checkpoint-label", required=True, help="Name to record in output")
    parser.add_argument("--revision", required=True)
    parser.add_argument("--compute-units", default="cpu_gpu")
    parser.add_argument("--state", type=Path, required=True)
    parser.add_argument("--questions", type=Path, required=True)
    parser.add_argument(
        "--out", type=Path, help="Where to write the answer set (ignored in --sweep mode)"
    )
    parser.add_argument("--timings", type=Path, required=True)
    parser.add_argument("--jev-response", type=Path, default=None)
    parser.add_argument("--skip-hash-check", action="store_true")
    parser.add_argument(
        "--sweep", action="store_true", help="Run the compute-unit sweep on 'action' only"
    )
    args = parser.parse_args()

    if not args.sweep and args.out is None:
        parser.error("--out is required unless --sweep is set")

    if not args.skip_hash_check:
        verify_prompt_hashes(args.state, args.questions)

    import laya_coreml as laya  # deferred: only needed for a real run

    state = _load_json(args.state)
    questions = _load_json(args.questions)

    load_start = time.perf_counter()
    agent = laya.load(
        args.model_dir, local_files_only=True, compute_units=args.compute_units
    )
    load_ms = (time.perf_counter() - load_start) * 1000.0

    first_call_start = time.perf_counter()
    first_result = agent.predict(state, {"action": questions["action"]})
    first_call_ms = (time.perf_counter() - first_call_start) * 1000.0

    timings: dict[str, Any] = {
        "checkpoint": args.checkpoint_label,
        "revision": args.revision,
        "compute_units": args.compute_units,
        "load_ms": load_ms,
        "first_call_ms": first_call_ms,
        "first_call_answer": first_result["answers"]["action"],
    }

    if args.sweep:

        def load_agent(_checkpoint: str, compute_units: str):
            return laya.load(
                args.model_dir, local_files_only=True, compute_units=compute_units
            )

        timings["compute_unit_sweep"] = compute_unit_sweep(
            load_agent, args.checkpoint_label, state, "action", questions["action"]
        )
        args.timings.write_text(json.dumps(timings, indent=2))
        print(f"Wrote sweep timings to {args.timings}")
        return

    per_question = benchmark_all_questions(agent.predict, state, questions)
    sequential = benchmark_sequential_pass(agent.predict, state, questions)

    timings["per_question"] = per_question
    timings["sequential_pass"] = sequential

    answers = {key: stats["last_answer"] for key, stats in per_question.items()}
    response = {
        "checkpoint": args.checkpoint_label,
        "revision": args.revision,
        "answers": answers,
    }

    if args.jev_response is not None:
        timings["vs_jev"] = compare_to_jev(answers, args.jev_response)

    args.out.write_text(json.dumps(response, indent=2))
    args.timings.write_text(json.dumps(timings, indent=2))
    print(f"Wrote {args.out} and {args.timings}")


if __name__ == "__main__":
    main()