#!/usr/bin/env python3
"""Run a pinned special-token-family boundary control on DeepSeek-V2-Lite.

The official tokenizer at revision 604d5664 has exactly two special token IDs:

    bos: 100000, <｜begin▁of▁sentence｜>
    eos: 100001, <｜end▁of▁sentence｜> (also used as pad_token)

This experiment keeps the audited repeated-token history and changes exactly
one completed-assistant boundary ID:

    system off/on x official EOS / counterfactual BOS / x / period

EOS and BOS exhaust the tokenizer's special-token inventory. X and period are
ordinary one-token controls. The three counterfactuals are not valid official
chat serializations. This design can identify differences among these four
pinned IDs inside one fixed BF16 batch contract; it cannot establish a general
"specialness" property from two special and two selected ordinary tokens.
"""

from __future__ import annotations

import contextlib
import hashlib
import json
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np

import v2_lite_routing_history_boundary_token_control as prior


base = prior.base

BOUNDARY_LEVELS = ("eos", "bos", "x", "period")
REFERENCE_LEVEL = "eos"
BOUNDARY_TEXT = {
    "x": "x",
    "period": ".",
}
CONDITIONS = tuple(
    f"s{system}_{boundary}"
    for boundary in BOUNDARY_LEVELS
    for system in (0, 1)
)
FACTORS = {
    condition: {
        "system": int(condition[1]),
        "history": "filler",
        "boundary": condition.split("_", 1)[1],
        "token_class": (
            "special"
            if condition.split("_", 1)[1] in {"eos", "bos"}
            else "ordinary_control"
        ),
    }
    for condition in CONDITIONS
}
SYSTEM_CELLS = {
    boundary: (f"s0_{boundary}", f"s1_{boundary}")
    for boundary in BOUNDARY_LEVELS
}
SYSTEM_EDGE_CONTRASTS = {
    f"{boundary}_minus_eos": ("eos", boundary)
    for boundary in BOUNDARY_LEVELS
    if boundary != "eos"
}
COMPARISONS = tuple(
    [
        (
            f"system_{boundary}",
            f"s0_{boundary}",
            f"s1_{boundary}",
        )
        for boundary in BOUNDARY_LEVELS
    ]
    + [
        (
            f"{boundary}_at_s{system}",
            f"s{system}_eos",
            f"s{system}_{boundary}",
        )
        for boundary in BOUNDARY_LEVELS
        if boundary != "eos"
        for system in (0, 1)
    ]
)
ALIGNMENT_COMPARISONS = COMPARISONS
ORIGINAL_BOUNDARY_DOMAIN = prior.boundary_control_domain


def special_family_token_ids(tokenizer: Any) -> dict[str, int]:
    """Resolve and validate the complete pinned special-token inventory."""
    if tokenizer.bos_token_id is None or tokenizer.eos_token_id is None:
        raise RuntimeError("tokenizer must expose both BOS and EOS")
    special_ids = [int(value) for value in tokenizer.all_special_ids]
    expected_special = {
        int(tokenizer.bos_token_id),
        int(tokenizer.eos_token_id),
    }
    if set(special_ids) != expected_special or len(special_ids) != 2:
        raise RuntimeError(
            "pinned special-token inventory changed: "
            f"observed={special_ids} expected={sorted(expected_special)}"
        )
    if tokenizer.pad_token_id != tokenizer.eos_token_id:
        raise RuntimeError(
            "pinned tokenizer no longer aliases pad_token_id to eos_token_id"
        )

    ids = {
        "eos": int(tokenizer.eos_token_id),
        "bos": int(tokenizer.bos_token_id),
    }
    for name, text in BOUNDARY_TEXT.items():
        encoded = list(
            tokenizer(text, add_special_tokens=False).input_ids
        )
        if len(encoded) != 1:
            raise RuntimeError(
                f"{name} boundary control is not one token: {encoded}"
            )
        if encoded[0] in expected_special:
            raise RuntimeError(
                f"{name} boundary control unexpectedly uses a special token"
            )
        ids[name] = int(encoded[0])
    if len(set(ids.values())) != len(ids):
        raise RuntimeError(f"boundary token IDs are not distinct: {ids}")
    return ids


def special_family_domain(
    loads: dict[str, np.ndarray],
    mode: str,
    replicates: int,
    seed: int,
    scope: str,
) -> dict[str, Any]:
    """Add a paired descriptive 2-special versus 2-ordinary summary."""
    result = ORIGINAL_BOUNDARY_DOMAIN(
        loads,
        mode,
        replicates,
        seed,
        scope,
    )
    shapes = {value.shape for value in loads.values()}
    if len(shapes) != 1:
        raise ValueError(
            f"special-family shape mismatch: {sorted(shapes)}"
        )
    rows = next(iter(loads.values())).shape[0]
    rng = np.random.default_rng(base.scoped_seed(seed, scope))
    sampled = rng.integers(
        0,
        rows,
        size=(replicates, rows),
        endpoint=False,
    )
    point = {
        condition: base.distribution(value, mode)
        for condition, value in loads.items()
    }
    boot = {
        condition: base.bootstrap_distributions(value, mode, sampled)
        for condition, value in loads.items()
    }

    point_tv: dict[str, float] = {}
    boot_tv: dict[str, np.ndarray] = {}
    point_jsd: dict[str, float] = {}
    boot_jsd: dict[str, np.ndarray] = {}
    for level, (before, after) in SYSTEM_CELLS.items():
        point_tv[level] = float(
            0.5 * np.abs(point[after] - point[before]).sum()
        )
        boot_tv[level] = 0.5 * np.abs(
            boot[after] - boot[before]
        ).sum(axis=1)
        point_jsd[level] = float(
            base.js_divergence(point[before], point[after])[0]
        )
        boot_jsd[level] = base.js_divergence(
            boot[before],
            boot[after],
        )

    def paired_family(
        point_metric: dict[str, float],
        boot_metric: dict[str, np.ndarray],
        unit: str | None = None,
    ) -> dict[str, Any]:
        special_point = 0.5 * (
            point_metric["eos"] + point_metric["bos"]
        )
        ordinary_point = 0.5 * (
            point_metric["x"] + point_metric["period"]
        )
        special_boot = 0.5 * (
            boot_metric["eos"] + boot_metric["bos"]
        )
        ordinary_boot = 0.5 * (
            boot_metric["x"] + boot_metric["period"]
        )
        payload = {
            "special_mean": {
                "point": special_point,
                "ci95": base.interval(special_boot),
            },
            "ordinary_control_mean": {
                "point": ordinary_point,
                "ci95": base.interval(ordinary_boot),
            },
            "ordinary_minus_special": {
                "point": ordinary_point - special_point,
                "ci95": base.interval(ordinary_boot - special_boot),
            },
        }
        if unit is not None:
            for value in payload.values():
                value["unit"] = unit
        return payload

    result["descriptive_family_summary"] = {
        "definition": (
            "mean(EOS,BOS) versus mean(x,period) inside the same shared "
            "source bootstrap; descriptive for these four IDs only"
        ),
        "total_variation": paired_family(point_tv, boot_tv),
        "js_divergence": paired_family(
            point_jsd,
            boot_jsd,
            unit="nats",
        ),
    }
    return result


def install_control_contract() -> None:
    """Install four pinned boundary IDs into the audited eight-cell runner."""
    prior.BOUNDARY_LEVELS = BOUNDARY_LEVELS
    prior.REFERENCE_LEVEL = REFERENCE_LEVEL
    prior.BOUNDARY_TEXT = BOUNDARY_TEXT
    prior.CONDITIONS = CONDITIONS
    prior.FACTORS = FACTORS
    prior.SYSTEM_CELLS = SYSTEM_CELLS
    prior.SYSTEM_EDGE_CONTRASTS = SYSTEM_EDGE_CONTRASTS
    prior.COMPARISONS = COMPARISONS
    prior.ALIGNMENT_COMPARISONS = ALIGNMENT_COMPARISONS
    prior.RENDER_AUDIT.clear()
    prior.BOUNDARY_TOKEN_IDS.clear()
    prior.boundary_token_ids = special_family_token_ids
    prior.boundary_control_domain = special_family_domain
    prior.install_control_contract()


def output_path_from_argv() -> Path:
    try:
        return Path(sys.argv[sys.argv.index("--output") + 1])
    except (ValueError, IndexError) as error:
        raise ValueError("--output is required") from error


def finalize_result(path: Path) -> dict[str, Any]:
    result = prior.finalize_result(path)
    result["schema_version"] = 4
    result["evidence_identity"] = (
        "X / official BF16 weights and pinned tokenizer; official EOS plus "
        "BOS and two ordinary single-ID boundary counterfactuals"
    )
    boundary = result["boundary"]
    boundary.update(
        {
            "special_token_family_control": True,
            "complete_pinned_special_inventory": True,
            "specialness_generalized": False,
            "causal_boundary": (
                "All replacement contrasts causally intervene on exactly one "
                "prior boundary input ID inside this run. BOS versus EOS "
                "exhausts the pinned tokenizer's two special IDs, while x "
                "and period are selected ordinary controls; four IDs do not "
                "identify a universal special-token category effect."
            ),
        }
    )
    boundary["official_serialization_by_boundary"] = {
        "eos": True,
        "bos": False,
        "x": False,
        "period": False,
    }

    old = result.pop("history_boundary_token_contract")
    result["special_token_family_contract"] = {
        **old,
        "boundary_levels": list(BOUNDARY_LEVELS),
        "boundary_token_ids": dict(prior.BOUNDARY_TOKEN_IDS),
        "boundary_text_controls": BOUNDARY_TEXT,
        "conditions": FACTORS,
        "tokenizer_special_inventory": {
            "all_special_tokens": [
                "<｜begin▁of▁sentence｜>",
                "<｜end▁of▁sentence｜>",
            ],
            "all_special_ids": [100000, 100001],
            "bos_token_id": 100000,
            "eos_token_id": 100001,
            "pad_token_id": 100001,
            "pad_aliases_eos": True,
            "inventory_size": 2,
        },
        "class_comparison_boundary": (
            "EOS/BOS are the complete special inventory, but x/period are "
            "only two chosen ordinary controls; report individual-ID "
            "contrasts and a descriptive 2-vs-2 family summary, not a "
            "population-level specialness claim"
        ),
    }
    result["inference_contract"]["batch_grouping"] = (
        "all eight EOS/BOS/x/period variants of one source prompt execute "
        "in the same 32-row right-padded batch"
    )

    path.write_text(
        json.dumps(result, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    return result


def main() -> None:
    install_control_contract()
    if {"-h", "--help"} & set(sys.argv[1:]):
        base.main()
        return
    output = output_path_from_argv()
    with open(os.devnull, "w", encoding="utf-8") as sink:
        with contextlib.redirect_stdout(sink):
            base.main()
    result = finalize_result(output)
    payload = output.read_bytes()
    print(
        json.dumps(
            {
                "output": str(output),
                "sha256": hashlib.sha256(payload).hexdigest(),
                "bytes": len(payload),
                "source_prompts": result["inference_contract"][
                    "total_source_prompts"
                ],
                "prompt_variants": result["inference_contract"][
                    "total_prompt_variants"
                ],
                "total_routes": result["inference_contract"][
                    "total_routes_all_conditions_all_moe_layers"
                ],
                "token_inventory": result[
                    "special_token_family_contract"
                ]["tokenizer_special_inventory"],
                "render_validation": result[
                    "special_token_family_contract"
                ]["render_validation"],
            },
            indent=2,
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()