"""Run the registered E3 v3 adaptive-depth benchmark on final Wisp weights.

Calibration selects one fixed depth and one entropy threshold without looking at
the test split. Every speculative output is then replayed with one causal
full-sequence forward over its own realized branch. Every emitted token must be
the branch-local argmax or a fully bound bfloat16 near-tie. Candidate policies
must also emit identical trajectories per document, so the registered paired
endpoint cannot be confounded by different near-tie branches.

Usage:
    .venv/bin/python scripts/eval_rollout.py \
      --ckpt out/run1/immutable/step-19073-89e81fb899d054cefaeca89443c20e4d4636167f43cfbc3f47579f79f5f27f22 \
      --receipt config/eval_rollout_receipt_v3.json \
      --acceptance-receipt config/eval_holdout_receipt.json \
      --holdout data/eval/holdout.clean.jsonl \
      --tokenizer tokenizer/code32k.json \
      --out out/run1/rollout.registered.v3.json
"""

import argparse
from collections import Counter
from datetime import datetime, timezone
from importlib import metadata as importlib_metadata
import json
import os
import platform
import sys
import tempfile

import mlx.core as mx
import numpy as np
from tokenizers import Tokenizer

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from eval_pairs import build_pairs, file_sha256, iter_holdout  # noqa: E402
from rollout_metrics import (  # noqa: E402
    NEAR_TIE_MAX_ULPS,
    V3_INSTRUMENT_VERSION,
    V3_REPLAY_REFERENCE,
    V3_REPLAY_RULE,
    V3_REPORT_SCHEMA_VERSION,
    build_branch_replay_evidence,
    canonical_json_sha256,
    metric_value,
    paired_mean_difference_ci,
    rollout_pair_manifest,
    rollout_pair_payload_manifest,
    select_policy,
    summarize_policy_v3,
    token_ids_sha256,
    validate_branch_replay,
    validate_cross_policy_trajectories,
    validate_pair_payload_manifest,
    validate_rollout_checkpoint,
    validate_rollout_receipt,
    validate_rollout_source_manifest,
    V3_TRACE_DOMAIN,
)
from sample import (  # noqa: E402
    DTYPES,
    causal_mask,
    generate_ar,
    generate_mtp,
    load_model,
)


INSTRUMENT_VERSION = V3_INSTRUMENT_VERSION


def load_json(path):
    with open(path, encoding="utf-8") as f:
        value = json.load(f)
    if not isinstance(value, dict):
        raise ValueError(f"{path}: top-level JSON value must be an object")
    return value


def package_version(name):
    try:
        return importlib_metadata.version(name)
    except importlib_metadata.PackageNotFoundError:
        return "unknown"


def policy_spec(name, max_depth):
    if name.startswith("fixed_d"):
        return {
            "name": name,
            "policy": "fixed",
            "depth": int(name.removeprefix("fixed_d")),
            "entropy_threshold": None,
        }
    if name.startswith("adaptive_h"):
        return {
            "name": name,
            "policy": "adaptive",
            "depth": max_depth,
            "entropy_threshold": float(name.removeprefix("adaptive_h")),
        }
    raise ValueError(f"unknown policy name {name!r}")


def branch_replay_evidence(model, prompt, policy_generated, trace_sha256):
    """Replay all emitted tokens on their own realized branch, never AR's."""
    if not prompt or not policy_generated:
        raise ValueError("branch replay needs non-empty prompt and output")
    dtype = model.norm.weight.dtype
    realized = list(prompt) + list(policy_generated)
    seq = mx.array([realized], dtype=mx.int32)
    hidden, _ = model.trunk(seq, causal_mask(seq.shape[1], dtype))
    start = len(prompt) - 1
    stop = start + len(policy_generated)
    logits = model.head(hidden[:, start:stop, :])
    mx.eval(logits)
    rows = np.array(logits[0].astype(mx.float32), copy=False)
    if rows.shape[0] != len(policy_generated):
        raise RuntimeError("branch replay returned the wrong number of rows")
    return build_branch_replay_evidence(
        prompt,
        policy_generated,
        rows,
        trace_sha256,
    )


def target_scores(generated, target):
    same = sum(
        left == right for left, right in zip(generated, target)
    )
    prefix = 0
    for left, right in zip(generated, target):
        if left != right:
            break
        prefix += 1
    return {
        "target_position_accuracy": same / max(len(target), 1),
        "target_common_prefix_tokens": prefix,
        "target_exact_match": generated == target,
    }


def run_policy(
    model,
    args,
    pair_index,
    pair_payload,
    prompt,
    target,
    ar_generated,
    ar_sha256,
    ar_tok_per_sec,
    metadata,
    spec,
    decoding,
    seed,
):
    rng = np.random.default_rng(seed)
    tokens, stats = generate_mtp(
        model,
        args,
        prompt,
        decoding["max_tokens"],
        spec["depth"],
        decoding["temperature"],
        decoding["top_k"],
        decoding["top_p"],
        decoding["draft_temperature"],
        rng,
        policy=spec["policy"],
        entropy_threshold=(
            spec["entropy_threshold"]
            if spec["entropy_threshold"] is not None
            else 1.0
        ),
        capture_trace=True,
    )
    trace = stats.pop("generation_trace")
    trace_sha256 = canonical_json_sha256(trace, V3_TRACE_DOMAIN)
    generated = [int(value) for value in tokens[len(prompt):]]
    ar_generated = [int(value) for value in ar_generated]
    generated_sha256 = token_ids_sha256(generated)
    accepted = sum(stats["rollout_accepted_per_depth"])
    ar_differences = [
        index
        for index, (policy_token, ar_token) in enumerate(
            zip(generated, ar_generated)
        )
        if policy_token != ar_token
    ]
    replay = branch_replay_evidence(
        model,
        prompt,
        generated,
        trace_sha256,
    )
    row = {
        "pair_index": pair_index,
        "seed": seed,
        "document_id": metadata["document_id"],
        "decoy_document_id": metadata["decoy_document_id"],
        "token_offset": metadata["token_offset"],
        "policy": spec["name"],
        "prompt_sha256": token_ids_sha256(prompt),
        "target_sha256": token_ids_sha256(target),
        "ar_output_sha256": ar_sha256,
        "ar_output_token_ids": ar_generated,
        "output_sha256": generated_sha256,
        "output_token_ids": generated,
        "output_matches_ar": generated == ar_generated,
        "cached_ar_diagnostic": {
            "output_sha256": ar_sha256,
            "matches": generated == ar_generated,
            "first_difference_position": (
                ar_differences[0] if ar_differences else None
            ),
        },
        "trace_sha256": trace_sha256,
        "generation_trace": trace,
        "branch_replay": replay,
        "tokens": stats["tokens"],
        "accepted_drafts": accepted,
        "verification_forwards": stats["verification_forwards"],
        "target_forwards": stats["target_forwards"],
        "drafts_issued": stats["drafts_issued"],
        "draft_recursions": stats["draft_recursions"],
        "corrections": stats["corrections"],
        "elapsed_seconds": stats["elapsed_seconds"],
        "ar_tok_per_sec": ar_tok_per_sec,
        "rollout": stats,
        **target_scores(generated, target),
    }
    try:
        validate_branch_replay(
            row,
            pair_payload,
            max_tokens=decoding["max_tokens"],
            vocab_size=args.vocab_size,
        )
    except Exception as error:
        raise RuntimeError(
            f"{spec['name']} failed branch-local replay on "
            f"{metadata['document_id']}: {error}"
        ) from error
    if replay["certified_near_tie_tokens"]:
        print(
            f"  branch replay certified "
            f"{replay['certified_near_tie_tokens']} near-tie token(s): "
            f"{spec['name']} on {metadata['document_id']}",
            flush=True,
        )
    return row


def run_documents(
    model,
    args,
    pairs,
    pair_payloads,
    policy_names,
    receipt,
    split_name,
    pair_start_index,
):
    decoding = receipt["decoding"]
    max_depth = receipt["policy"]["max_depth"]
    if len(pairs) != len(pair_payloads):
        raise ValueError("rollout pairs and payloads have different sizes")
    rows = {name: [] for name in policy_names}
    for index, (pair, pair_payload) in enumerate(zip(pairs, pair_payloads)):
        fim, span = pair["fim"]
        prompt = fim[:span[0]]
        target = fim[span[0]:span[1]]
        if len(target) != decoding["max_tokens"]:
            raise ValueError("rollout target length differs from max_tokens")
        pair_index = pair_start_index + index
        seed = decoding["seed"] + index
        ar_rng = np.random.default_rng(seed)
        ar_tokens, ar_tok_per_sec = generate_ar(
            model,
            args,
            prompt,
            decoding["max_tokens"],
            decoding["temperature"],
            decoding["top_k"],
            decoding["top_p"],
            ar_rng,
        )
        ar_generated = ar_tokens[len(prompt):]
        ar_sha256 = token_ids_sha256(ar_generated)
        for name in policy_names:
            spec = policy_spec(name, max_depth)
            rows[name].append(
                run_policy(
                    model,
                    args,
                    pair_index,
                    pair_payload,
                    prompt,
                    target,
                    ar_generated,
                    ar_sha256,
                    ar_tok_per_sec,
                    pair["_meta"],
                    spec,
                    decoding,
                    seed,
                )
            )
        if (index + 1) % 5 == 0:
            print(
                f"  {split_name}: {index + 1}/{len(pairs)} documents",
                flush=True,
            )
    return rows


def atomic_write_json(path, value):
    path = os.path.abspath(path)
    if os.path.lexists(path):
        raise FileExistsError(f"refusing to replace rollout report: {path}")
    parent = os.path.dirname(path)
    os.makedirs(parent, exist_ok=True)
    rendered = json.dumps(value, indent=2, sort_keys=True, allow_nan=False)
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            prefix=f".{os.path.basename(path)}.",
            suffix=".tmp",
            dir=parent,
            delete=False,
        ) as f:
            temporary = f.name
            f.write(rendered)
            f.write("\n")
            f.flush()
            os.fsync(f.fileno())
        if os.path.lexists(path):
            raise FileExistsError(f"rollout target appeared during staging: {path}")
        os.rename(temporary, path)
        temporary = None
    finally:
        if temporary is not None and os.path.isfile(temporary):
            os.unlink(temporary)
    return path


def main():
    generation_started_at = datetime.now(timezone.utc).isoformat().replace(
        "+00:00", "Z"
    )
    source_root = os.path.dirname(
        os.path.dirname(os.path.abspath(__file__))
    )
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", required=True)
    parser.add_argument(
        "--receipt", default="config/eval_rollout_receipt_v3.json"
    )
    parser.add_argument(
        "--acceptance-receipt",
        default="config/eval_holdout_receipt.json",
    )
    parser.add_argument(
        "--holdout", default="data/eval/holdout.clean.jsonl"
    )
    parser.add_argument("--tokenizer", default="tokenizer/code32k.json")
    parser.add_argument("--out", required=True)
    cli = parser.parse_args()

    receipt = load_json(cli.receipt)
    receipt_evidence = validate_rollout_receipt(
        receipt,
        cli.receipt,
        cli.acceptance_receipt,
        cli.holdout,
        cli.tokenizer,
        INSTRUMENT_VERSION,
        source_root=source_root,
    )
    runtime_versions = {
        "python": platform.python_version(),
        "mlx": package_version("mlx"),
        "numpy": package_version("numpy"),
        "tokenizers": package_version("tokenizers"),
    }
    if runtime_versions != receipt["runtime_requirements"]:
        raise RuntimeError(
            f"runtime {runtime_versions} differs from registered "
            f"{receipt['runtime_requirements']}"
        )
    if list(sys.argv) != receipt["execution_argv"]:
        raise RuntimeError(
            f"execution argv {list(sys.argv)} differs from registered "
            f"{receipt['execution_argv']}"
        )
    meta_path = os.path.join(cli.ckpt, "meta.json")
    master_path = os.path.join(cli.ckpt, "master.safetensors")
    meta = load_json(meta_path)
    validate_rollout_checkpoint(meta, receipt)
    checkpoint_hashes = {
        "meta_sha256": file_sha256(meta_path),
        "master_sha256": file_sha256(master_path),
    }

    tok = Tokenizer.from_file(cli.tokenizer)
    sentinels = {
        "prefix": tok.token_to_id("<|fim_prefix|>"),
        "middle": tok.token_to_id("<|fim_middle|>"),
        "suffix": tok.token_to_id("<|fim_suffix|>"),
    }
    if any(value is None for value in sentinels.values()):
        raise ValueError("tokenizer is missing FIM sentinels")
    pair_settings = receipt["pair_settings"]
    pairs = build_pairs(
        iter_holdout(cli.holdout),
        tok,
        sentinels,
        pair_settings["examples"],
        pair_settings["prefix_len"],
        pair_settings["span_len"],
        pair_settings["suffix_len"],
        np.random.default_rng(pair_settings["seed"]),
    )
    if len(pairs) != pair_settings["examples"]:
        raise ValueError(
            f"constructed {len(pairs)} rollout pairs, "
            f"expected {pair_settings['examples']}"
        )
    pair_manifest, pair_manifest_sha256 = rollout_pair_manifest(pairs)
    if pair_manifest_sha256 != pair_settings["pair_manifest_sha256"]:
        raise ValueError(
            f"rollout pair manifest {pair_manifest_sha256} does not match "
            f"registered {pair_settings['pair_manifest_sha256']}"
        )
    pair_payload_manifest, pair_payload_sha256 = (
        rollout_pair_payload_manifest(pairs)
    )
    validated_payload_sha256 = validate_pair_payload_manifest(
        pair_payload_manifest,
        pair_manifest,
        max_tokens=receipt["decoding"]["max_tokens"],
        vocab_size=receipt["model_vocab_size"],
    )
    if (
        pair_payload_sha256 != validated_payload_sha256
        or pair_payload_sha256
        != pair_settings["pair_payload_manifest_sha256"]
    ):
        raise ValueError(
            f"rollout pair payload manifest {pair_payload_sha256} does not "
            f"match registered {pair_settings['pair_payload_manifest_sha256']}"
        )
    matches = Counter(pair["_meta"]["decoy_match"] for pair in pairs)
    if pair_settings["require_matched_decoys"] and matches != {
        "matched": len(pairs)
    }:
        raise ValueError(f"rollout pairs have relaxed decoys: {dict(matches)}")

    model, args, loaded_meta = load_model(
        cli.ckpt, DTYPES[receipt["decoding"]["dtype"]]
    )
    if loaded_meta != meta:
        raise RuntimeError("checkpoint metadata changed while loading rollout eval")
    split = receipt["split"]
    calibration_count = split["calibration_documents"]
    calibration_pairs = pairs[:calibration_count]
    test_pairs = pairs[calibration_count:]
    calibration_payloads = pair_payload_manifest[:calibration_count]
    test_payloads = pair_payload_manifest[calibration_count:]
    policy = receipt["policy"]
    candidate_order = [
        *policy["fixed_candidates"],
        *policy["adaptive_candidates"],
    ]
    print(
        f"checkpoint step {meta['step']}, "
        f"{len(calibration_pairs)} calibration and {len(test_pairs)} test "
        "documents",
        flush=True,
    )
    calibration_rows = run_documents(
        model,
        args,
        calibration_pairs,
        calibration_payloads,
        candidate_order,
        receipt,
        "calibration",
        0,
    )
    calibration_summaries = {
        name: summarize_policy_v3(
            rows,
            calibration_payloads,
            max_tokens=receipt["decoding"]["max_tokens"],
            vocab_size=receipt["model_vocab_size"],
        )
        for name, rows in calibration_rows.items()
    }
    calibration_trajectory_identity = validate_cross_policy_trajectories(
        calibration_rows, candidate_order
    )
    selected_fixed = select_policy(
        calibration_summaries,
        policy["fixed_candidates"],
        policy["selection_metric"],
    )
    selected_adaptive = select_policy(
        calibration_summaries,
        policy["adaptive_candidates"],
        policy["selection_metric"],
    )
    selected_names = [
        selected_fixed["policy"],
        selected_adaptive["policy"],
    ]
    print(
        f"selected {selected_fixed['policy']} and "
        f"{selected_adaptive['policy']} on calibration",
        flush=True,
    )
    test_rows = run_documents(
        model,
        args,
        test_pairs,
        test_payloads,
        selected_names,
        receipt,
        "test",
        calibration_count,
    )
    test_summaries = {
        name: summarize_policy_v3(
            rows,
            test_payloads,
            max_tokens=receipt["decoding"]["max_tokens"],
            vocab_size=receipt["model_vocab_size"],
        )
        for name, rows in test_rows.items()
    }
    test_trajectory_identity = validate_cross_policy_trajectories(
        test_rows, selected_names
    )
    metric = receipt["test_endpoint"]["metric"]
    adaptive_values = [
        metric_value(row, metric)
        for row in test_rows[selected_adaptive["policy"]]
    ]
    fixed_values = [
        metric_value(row, metric)
        for row in test_rows[selected_fixed["policy"]]
    ]
    endpoint = receipt["test_endpoint"]
    difference, lo, hi = paired_mean_difference_ci(
        adaptive_values,
        fixed_values,
        n_boot=endpoint["bootstrap_samples"],
        seed=endpoint["bootstrap_seed"],
    )
    if lo > 0:
        verdict = "POSITIVE"
    elif hi < 0:
        verdict = "NEGATIVE"
    else:
        verdict = "NULL: the interval includes 0"

    def companion_result(specification):
        companion_metric = specification["metric"]
        companion_adaptive = [
            metric_value(row, companion_metric)
            for row in test_rows[selected_adaptive["policy"]]
        ]
        companion_fixed = [
            metric_value(row, companion_metric)
            for row in test_rows[selected_fixed["policy"]]
        ]
        companion_difference, companion_lo, companion_hi = (
            paired_mean_difference_ci(
                companion_adaptive,
                companion_fixed,
                n_boot=specification["bootstrap_samples"],
                seed=specification["bootstrap_seed"],
            )
        )
        if companion_lo > 0:
            companion_verdict = "POSITIVE"
        elif companion_hi < 0:
            companion_verdict = "NEGATIVE"
        else:
            companion_verdict = "NULL: the interval includes 0"
        return {
            "comparison": specification["comparison"],
            "metric": companion_metric,
            "adaptive_policy": selected_adaptive["policy"],
            "fixed_policy": selected_fixed["policy"],
            "difference": companion_difference,
            "ci95": [companion_lo, companion_hi],
            "documents": len(test_pairs),
            "verdict": companion_verdict,
        }

    companion_results = {
        specification["metric"]: companion_result(specification)
        for specification in receipt["companion_endpoints"]
    }
    all_summaries = [
        *calibration_summaries.values(),
        *test_summaries.values(),
    ]
    scored_policy_documents = sum(
        summary["documents"] for summary in all_summaries
    )
    scored_tokens = sum(summary["total_tokens"] for summary in all_summaries)
    exact_argmax_tokens = sum(
        summary["exact_argmax_tokens"] for summary in all_summaries
    )
    certified_near_tie_tokens = sum(
        summary["certified_near_tie_tokens"]
        for summary in all_summaries
    )
    branch_replay_passes = sum(
        summary["branch_replay_passes"] for summary in all_summaries
    )
    cached_ar_exact = sum(
        summary["cached_ar_exact_documents"] for summary in all_summaries
    )
    generation_completed_at = datetime.now(
        timezone.utc
    ).isoformat().replace("+00:00", "Z")

    report = {
        "schema_version": V3_REPORT_SCHEMA_VERSION,
        "instrument_version": INSTRUMENT_VERSION,
        "publication_ready": True,
        "execution": {
            "started_at": generation_started_at,
            "completed_at": generation_completed_at,
            "argv": list(sys.argv),
            "runtime": {
                **runtime_versions,
                "platform": platform.platform(),
            },
        },
        "receipt": receipt_evidence,
        "checkpoint": {
            "path": cli.ckpt,
            "step": meta["step"],
            **checkpoint_hashes,
        },
        "tokenizer": {
            "path": cli.tokenizer,
            "sha256": file_sha256(cli.tokenizer),
        },
        "holdout": {
            "path": cli.holdout,
            "sha256": file_sha256(cli.holdout),
            "pair_count": len(pairs),
            "pair_manifest_sha256": pair_manifest_sha256,
            "pair_manifest": pair_manifest,
            "pair_payload_manifest_sha256": pair_payload_sha256,
            "pair_payload_manifest": pair_payload_manifest,
            "decoy_match": dict(sorted(matches.items())),
        },
        "calibration": {
            "documents": len(calibration_pairs),
            "candidate_order": candidate_order,
            "trajectory_identity": calibration_trajectory_identity,
            "summaries": calibration_summaries,
            "selected_fixed": selected_fixed,
            "selected_adaptive": selected_adaptive,
            "rows": calibration_rows,
        },
        "test": {
            "documents": len(test_pairs),
            "trajectory_identity": test_trajectory_identity,
            "summaries": test_summaries,
            "rows": test_rows,
        },
        "quality_gate": {
            "reference": V3_REPLAY_REFERENCE,
            "rule": V3_REPLAY_RULE,
            "near_tie_max_ulps": NEAR_TIE_MAX_ULPS,
            "scored_policy_documents": scored_policy_documents,
            "scored_tokens": scored_tokens,
            "exact_argmax_tokens": exact_argmax_tokens,
            "certified_near_tie_tokens": certified_near_tie_tokens,
            "failed_tokens": 0,
            "branch_replay_passes": branch_replay_passes,
            "cross_policy_trajectory_matches": (
                calibration_trajectory_identity["matching_documents"]
                + test_trajectory_identity["matching_documents"]
            ),
            "passed": True,
        },
        "cached_ar_diagnostic": {
            "scored_policy_documents": scored_policy_documents,
            "exact_output_matches": cached_ar_exact,
            "different_cached_ar_branches": (
                scored_policy_documents - cached_ar_exact
            ),
            "claim_scope": (
                "diagnostic_only; release quality is established by "
                "branch-local replay, not cached-AR byte identity"
            ),
        },
        "primary_endpoint": {
            "comparison": endpoint["comparison"],
            "metric": metric,
            "adaptive_policy": selected_adaptive["policy"],
            "fixed_policy": selected_fixed["policy"],
            "difference": difference,
            "ci95": [lo, hi],
            "documents": len(test_pairs),
            "verdict": verdict,
        },
        "secondary_target_forward_endpoint": companion_results[
            "output_tokens_per_target_forward"
        ],
        "secondary_draft_issued_proxy_endpoint": companion_results[
            "drafts_issued_per_output_token"
        ],
        "secondary_draft_work_endpoint": companion_results[
            "draft_recursions_per_output_token"
        ],
        "wall_clock_note": (
            "This reference recomputes full prefixes and has no rollback-capable "
            "KV cache. Wall time is recorded for audit, not claimed as deployment "
            "latency."
        ),
        "endpoint_scope_note": (
            "The primary endpoint measures accepted drafts per verification. "
            "It does not establish verification-width cost or deployment "
            "latency. Draft recursions per output token is the registered "
            "drafter-work companion; issued drafts per output token is retained "
            "only as an issuance proxy. Target forwards exclude the added "
            "post-hoc branch-replay forward and independent verification pass."
        ),
    }
    if exact_argmax_tokens + certified_near_tie_tokens != scored_tokens:
        raise RuntimeError("branch replay aggregate token counts do not close")
    if branch_replay_passes != scored_policy_documents:
        raise RuntimeError("not every rollout row passed branch-local replay")
    if (
        calibration_trajectory_identity["matching_documents"]
        != len(calibration_pairs)
        or test_trajectory_identity["matching_documents"] != len(test_pairs)
    ):
        raise RuntimeError("cross-policy trajectory identity did not pass")
    if (
        file_sha256(meta_path) != checkpoint_hashes["meta_sha256"]
        or file_sha256(master_path) != checkpoint_hashes["master_sha256"]
    ):
        raise RuntimeError("checkpoint changed during rollout evaluation")
    if (
        validate_rollout_source_manifest(receipt, source_root)
        != receipt_evidence["implementation"]
    ):
        raise RuntimeError("registered rollout sources changed during evaluation")
    written = atomic_write_json(cli.out, report)
    print(
        f"primary adaptive minus fixed difference {difference:.4f} "
        f"[{lo:.4f}, {hi:.4f}] -> {verdict}",
        flush=True,
    )
    print(f"wrote {written}")


if __name__ == "__main__":
    main()