"""CPU mutation checks for adaptive-depth rollout policy selection."""

import copy
import json
import math
import os
import sys

import numpy as np
from tokenizers import Tokenizer

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

from scripts.rollout_metrics import (  # noqa: E402
    NEAR_TIE_MAX_ULPS,
    bf16_ulp,
    metric_value,
    normalized_entropy,
    paired_mean_difference_ci,
    rollout_pair_manifest,
    rollout_pair_payload_manifest,
    select_policy,
    summarize_policy,
    validate_divergence_evidence,
    validate_pair_payload_manifest,
    validate_rollout_checkpoint,
    validate_rollout_receipt,
)
from scripts.eval_pairs import build_pairs, iter_holdout  # noqa: E402


def rows(accepted, forwards):
    return [
        {
            "document_id": f"doc-{index}",
            "output_matches_ar": True,
            "divergence": None,
            "accepted_drafts": accepted[index],
            "verification_forwards": 10,
            "tokens": 64,
            "target_forwards": forwards[index],
            "elapsed_seconds": 1.0 + index * 0.1,
        }
        for index in range(len(accepted))
    ]


def near_tie_evidence():
    """A certified near-tie one bf16 ulp below the row maximum."""
    row_max = 6.1875
    ulp = bf16_ulp(row_max)
    ar_logit = row_max
    policy_logit = row_max - ulp
    return {
        "position": 10,
        "ar_token": 2776,
        "policy_token": 5187,
        "ar_token_logit": ar_logit,
        "policy_token_logit": policy_logit,
        "row_max_logit": row_max,
        "ulp_at_max": ulp,
        "max_ulps": NEAR_TIE_MAX_ULPS,
        "ar_token_deficit_ulps": (row_max - ar_logit) / ulp,
        "policy_token_deficit_ulps": (row_max - policy_logit) / ulp,
    }


def main():
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    receipt_path = os.path.join(root, "config", "eval_rollout_receipt.json")
    with open(receipt_path, encoding="utf-8") as f:
        receipt = json.load(f)
    evidence = validate_rollout_receipt(
        receipt,
        receipt_path,
        os.path.join(root, "config", "eval_holdout_receipt.json"),
        os.path.join(root, "data", "eval", "holdout.clean.jsonl"),
        os.path.join(root, "tokenizer", "code32k.json"),
    )
    assert len(evidence["sha256"]) == 64
    tokenizer = Tokenizer.from_file(
        os.path.join(root, "tokenizer", "code32k.json")
    )
    sentinels = {
        "prefix": tokenizer.token_to_id("<|fim_prefix|>"),
        "middle": tokenizer.token_to_id("<|fim_middle|>"),
        "suffix": tokenizer.token_to_id("<|fim_suffix|>"),
    }
    settings = receipt["pair_settings"]
    pairs = build_pairs(
        iter_holdout(
            os.path.join(root, "data", "eval", "holdout.clean.jsonl")
        ),
        tokenizer,
        sentinels,
        settings["examples"],
        settings["prefix_len"],
        settings["span_len"],
        settings["suffix_len"],
        np.random.default_rng(settings["seed"]),
    )
    pair_rows, pair_digest = rollout_pair_manifest(pairs)
    assert len(pair_rows) == settings["examples"]
    assert pair_digest == settings["pair_manifest_sha256"]
    payload_rows, payload_digest = rollout_pair_payload_manifest(pairs)
    assert payload_digest == (
        "9f01a7301f31882428ab56b0d27b44074a515513cdcc1efee9cded8e7a05f830"
    )
    assert validate_pair_payload_manifest(
        payload_rows,
        pair_rows,
        max_tokens=receipt["decoding"]["max_tokens"],
        vocab_size=32768,
    ) == payload_digest
    final_meta = {
        "step": receipt["trained_checkpoint_step"],
        "optimizer_state_included": True,
        "config": {"run_name": receipt["trained_run_name"]},
    }
    validate_rollout_checkpoint(final_meta, receipt)
    incomplete_meta = {
        **final_meta,
        "step": receipt["trained_checkpoint_step"] - 1,
    }
    try:
        validate_rollout_checkpoint(incomplete_meta, receipt)
    except ValueError as exc:
        assert "step 19072 != 19073" in str(exc)
    else:
        raise AssertionError("an incomplete checkpoint passed rollout")

    assert normalized_entropy([1.0, 0.0, 0.0, 0.0]) == 0.0
    assert abs(normalized_entropy([0.25] * 4) - 1.0) < 1e-12
    skewed = normalized_entropy([0.7, 0.1, 0.1, 0.1])
    assert 0.0 < skewed < 1.0

    fixed_rows = rows([10, 11, 9, 10], [20, 19, 21, 20])
    adaptive_rows = rows([14, 15, 13, 14], [16, 15, 17, 16])
    fixed = summarize_policy(fixed_rows)
    adaptive = summarize_policy(adaptive_rows)
    assert fixed["exact_ar_matches"] == 4
    assert fixed["certified_divergences"] == 0
    assert adaptive["mean_accepted_drafts_per_verification"] > (
        fixed["mean_accepted_drafts_per_verification"]
    )
    selected = select_policy(
        {"fixed_d2": fixed, "adaptive_h0.5": adaptive},
        ["fixed_d2", "adaptive_h0.5"],
        "accepted_drafts_per_verification",
    )
    assert selected["policy"] == "adaptive_h0.5"

    adaptive_values = [
        row["accepted_drafts"] / row["verification_forwards"]
        for row in adaptive_rows
    ]
    fixed_values = [
        row["accepted_drafts"] / row["verification_forwards"]
        for row in fixed_rows
    ]
    difference, lo, hi = paired_mean_difference_ci(
        adaptive_values, fixed_values, n_boot=500, seed=3
    )
    assert math.isclose(difference, 0.4)
    assert lo > 0
    assert hi >= lo

    bad = copy.deepcopy(adaptive_rows)
    bad[0]["output_matches_ar"] = False
    try:
        summarize_policy(bad)
    except ValueError as exc:
        assert "differs from greedy AR" in str(exc)
    else:
        raise AssertionError("output mismatch was accepted as equal quality")

    certified_rows = copy.deepcopy(adaptive_rows)
    certified_rows[0]["output_matches_ar"] = False
    certified_rows[0]["divergence"] = near_tie_evidence()
    certified_summary = summarize_policy(certified_rows)
    assert certified_summary["exact_ar_matches"] == 3
    assert certified_summary["certified_divergences"] == 1

    contradictory = copy.deepcopy(adaptive_rows)
    contradictory[0]["divergence"] = near_tie_evidence()
    try:
        summarize_policy(contradictory)
    except ValueError as exc:
        assert "carries divergence evidence" in str(exc)
    else:
        raise AssertionError(
            "a matching row with divergence evidence was accepted"
        )

    validate_divergence_evidence(near_tie_evidence())

    def _expect_evidence_reject(mutate, expected_substring):
        evidence = near_tie_evidence()
        mutate(evidence)
        try:
            validate_divergence_evidence(evidence)
        except ValueError as exc:
            assert expected_substring in str(exc), (
                f"expected {expected_substring!r} in {exc}"
            )
        else:
            raise AssertionError(
                f"divergence evidence that should have failed on "
                f"{expected_substring!r} was accepted"
            )

    _expect_evidence_reject(
        lambda e: e.pop("row_max_logit"), "wrong fields"
    )
    _expect_evidence_reject(
        lambda e: e.update(extra=True), "wrong fields"
    )
    _expect_evidence_reject(
        lambda e: e.update(position=-1), "non-negative integer"
    )
    _expect_evidence_reject(
        lambda e: e.update(policy_token=e["ar_token"]), "do not differ"
    )
    _expect_evidence_reject(
        lambda e: e.update(max_ulps=NEAR_TIE_MAX_ULPS + 1),
        "differs from the registered gate",
    )
    _expect_evidence_reject(
        lambda e: e.update(ulp_at_max=e["ulp_at_max"] * 2),
        "does not recompute from the row maximum",
    )
    _expect_evidence_reject(
        lambda e: e.update(policy_token_logit=e["row_max_logit"] + 1.0),
        "exceeds the row maximum",
    )
    _expect_evidence_reject(
        lambda e: e.update(policy_token_deficit_ulps=0.0),
        "does not recompute from its logit",
    )

    def _wide_gap(evidence):
        wide_logit = evidence["row_max_logit"] - (
            (NEAR_TIE_MAX_ULPS + 1) * evidence["ulp_at_max"]
        )
        evidence["policy_token_logit"] = wide_logit
        evidence["policy_token_deficit_ulps"] = (
            evidence["row_max_logit"] - wide_logit
        ) / evidence["ulp_at_max"]

    _expect_evidence_reject(_wide_gap, "outside the certified near-tie budget")

    duplicate = copy.deepcopy(adaptive_rows)
    duplicate[1]["document_id"] = duplicate[0]["document_id"]
    try:
        summarize_policy(duplicate)
    except ValueError as exc:
        assert "not unique" in str(exc)
    else:
        raise AssertionError("duplicate rollout document was accepted")

    try:
        normalized_entropy(np.asarray([0.5, -0.5]))
    except ValueError:
        pass
    else:
        raise AssertionError("negative probability passed entropy validation")

    changed_contract = copy.deepcopy(receipt)
    changed_contract["decoding"]["max_tokens"] = 63
    try:
        validate_rollout_receipt(
            changed_contract,
            receipt_path,
            os.path.join(root, "config", "eval_holdout_receipt.json"),
            os.path.join(root, "data", "eval", "holdout.clean.jsonl"),
            os.path.join(root, "tokenizer", "code32k.json"),
        )
    except ValueError as exc:
        assert "decoding settings" in str(exc)
    else:
        raise AssertionError("a mutated rollout contract was accepted")

    def _expect_receipt_reject(mutated_receipt, expected_substring):
        try:
            validate_rollout_receipt(
                mutated_receipt,
                receipt_path,
                os.path.join(root, "config", "eval_holdout_receipt.json"),
                os.path.join(root, "data", "eval", "holdout.clean.jsonl"),
                os.path.join(root, "tokenizer", "code32k.json"),
            )
        except ValueError as exc:
            assert expected_substring in str(exc), (
                f"expected {expected_substring!r} in {exc}"
            )
        else:
            raise AssertionError(
                f"a rollout receipt that should have failed on "
                f"{expected_substring!r} was accepted"
            )

    bad = copy.deepcopy(receipt)
    bad["schema_version"] = 2
    _expect_receipt_reject(bad, "must use schema_version 1")

    bad = copy.deepcopy(receipt)
    bad["instrument_version"] = 999
    _expect_receipt_reject(bad, "instrument version does not match code")

    for key in ("acceptance_receipt", "holdout", "tokenizer"):
        bad = copy.deepcopy(receipt)
        bad[key] = dict(bad[key])
        bad[key]["sha256"] = "0" * 64
        _expect_receipt_reject(bad, f"{key} sha256")

    bad = copy.deepcopy(receipt)
    bad["split"] = dict(bad["split"])
    bad["split"]["calibration_documents"] = 1
    _expect_receipt_reject(bad, "invalid frozen split")

    bad = copy.deepcopy(receipt)
    bad["pair_settings"] = dict(bad["pair_settings"])
    bad["pair_settings"]["pair_manifest_sha256"] = "not-hex-zzz"
    _expect_receipt_reject(bad, "pair manifest SHA-256 is invalid")

    bad = copy.deepcopy(receipt)
    bad["policy"] = dict(bad["policy"])
    bad["policy"]["fixed_candidates"] = ["wrong"]
    _expect_receipt_reject(bad, "fixed policy candidates differ from code")

    bad = copy.deepcopy(receipt)
    bad["policy"] = dict(bad["policy"])
    bad["policy"]["adaptive_candidates"] = ["wrong"]
    _expect_receipt_reject(bad, "adaptive policy candidates differ from code")

    bad = copy.deepcopy(receipt)
    bad["policy"] = dict(bad["policy"])
    bad["policy"]["max_depth"] = 5
    _expect_receipt_reject(bad, "rollout max depth must be 4")

    bad = copy.deepcopy(receipt)
    bad["policy"] = dict(bad["policy"])
    bad["policy"]["entropy_temperature"] = 0.5
    _expect_receipt_reject(bad, "entropy temperature must be 1.0")

    bad = copy.deepcopy(receipt)
    bad["policy"] = dict(bad["policy"])
    bad["policy"]["entropy_normalization"] = "wrong"
    _expect_receipt_reject(bad, "entropy normalization differs from code")

    bad = copy.deepcopy(receipt)
    bad["policy"] = dict(bad["policy"])
    bad["policy"]["tie_break"] = "wrong"
    _expect_receipt_reject(bad, "selection rule differs from code")

    bad = copy.deepcopy(receipt)
    bad["test_endpoint"] = dict(bad["test_endpoint"])
    bad["test_endpoint"]["bootstrap_seed"] = 999
    _expect_receipt_reject(bad, "test endpoint differs from code")

    bad = copy.deepcopy(receipt)
    bad["quality_gate"] = dict(bad["quality_gate"])
    bad["quality_gate"]["reference"] = "wrong"
    _expect_receipt_reject(bad, "quality gate differs from code")

    try:
        normalized_entropy([1.0])
    except ValueError as exc:
        assert "must be one-dimensional" in str(exc)
    else:
        raise AssertionError("a single-element entropy distribution was accepted")

    try:
        normalized_entropy(np.asarray([0.5, float("inf")]))
    except ValueError as exc:
        assert "must be finite and non-negative" in str(exc)
    else:
        raise AssertionError("a non-finite entropy distribution was accepted")

    try:
        normalized_entropy([0.0, 0.0, 0.0])
    except ValueError as exc:
        assert "must have positive mass" in str(exc)
    else:
        raise AssertionError("an all-zero entropy distribution was accepted")

    try:
        metric_value({"accepted_drafts": 1, "verification_forwards": 1}, "bogus")
    except ValueError as exc:
        assert "unknown rollout metric" in str(exc)
    else:
        raise AssertionError("an unknown rollout metric name was accepted")

    try:
        metric_value(
            {"accepted_drafts": -1, "verification_forwards": 1},
            "accepted_drafts_per_verification",
        )
    except ValueError as exc:
        assert "invalid counts for rollout metric" in str(exc)
    else:
        raise AssertionError("a negative accepted-drafts count was accepted")

    try:
        summarize_policy([])
    except ValueError as exc:
        assert "at least one document" in str(exc)
    else:
        raise AssertionError("an empty rollout policy was accepted")

    missing_identity = copy.deepcopy(fixed_rows)
    missing_identity[0]["document_id"] = ""
    try:
        summarize_policy(missing_identity)
    except ValueError as exc:
        assert "document identity is missing" in str(exc)
    else:
        raise AssertionError("a rollout row with a missing document id was accepted")

    non_finite_elapsed = copy.deepcopy(fixed_rows)
    non_finite_elapsed[0]["elapsed_seconds"] = 0.0
    try:
        summarize_policy(non_finite_elapsed)
    except ValueError as exc:
        assert "elapsed times must be finite and positive" in str(exc)
    else:
        raise AssertionError("a zero elapsed time was accepted")

    try:
        select_policy(
            {"fixed_d2": fixed, "adaptive_h0.5": adaptive},
            ["fixed_d2", "adaptive_h0.5"],
            "bogus_metric",
        )
    except ValueError as exc:
        assert "unknown selection metric" in str(exc)
    else:
        raise AssertionError("an unknown selection metric was accepted")

    try:
        select_policy(
            {"fixed_d2": fixed},
            ["fixed_d2", "adaptive_h0.5"],
            "accepted_drafts_per_verification",
        )
    except ValueError as exc:
        assert "missing calibration summary for adaptive_h0.5" in str(exc)
    else:
        raise AssertionError(
            "selection over a candidate with no calibration summary was accepted"
        )

    failed_equivalence = dict(adaptive)
    failed_equivalence["certified_divergences"] = (
        failed_equivalence["certified_divergences"] + 1
    )
    try:
        select_policy(
            {"fixed_d2": fixed, "adaptive_h0.5": failed_equivalence},
            ["fixed_d2", "adaptive_h0.5"],
            "accepted_drafts_per_verification",
        )
    except ValueError as exc:
        assert "failed output equivalence" in str(exc)
    else:
        raise AssertionError(
            "selection over a candidate that failed output equivalence was accepted"
        )

    invalid_metric_value = dict(adaptive)
    invalid_metric_value["mean_accepted_drafts_per_verification"] = float("nan")
    try:
        select_policy(
            {"fixed_d2": fixed, "adaptive_h0.5": invalid_metric_value},
            ["fixed_d2", "adaptive_h0.5"],
            "accepted_drafts_per_verification",
        )
    except ValueError as exc:
        assert "has invalid mean_accepted_drafts_per_verification" in str(exc)
    else:
        raise AssertionError(
            "selection over a candidate with a non-finite metric was accepted"
        )

    try:
        paired_mean_difference_ci(
            np.asarray([[1.0, 2.0]]), np.asarray([1.0, 2.0]), n_boot=10, seed=0
        )
    except ValueError as exc:
        assert "must be one-dimensional" in str(exc)
    else:
        raise AssertionError("a non-1D rollout comparison input was accepted")

    try:
        paired_mean_difference_ci([1.0, 2.0], [1.0], n_boot=10, seed=0)
    except ValueError as exc:
        assert "requires equal paired samples" in str(exc)
    else:
        raise AssertionError("mismatched paired-sample sizes were accepted")

    try:
        paired_mean_difference_ci(
            [1.0, float("nan")], [1.0, 2.0], n_boot=10, seed=0
        )
    except ValueError as exc:
        assert "must be finite" in str(exc)
    else:
        raise AssertionError("a non-finite rollout comparison input was accepted")

    print("rollout metrics: PASS")


if __name__ == "__main__":
    main()