#!/usr/bin/env python3
"""Run a 2 x 3 message-history control on DeepSeek-V2-Lite.

The system factor is off/on. The history factor has three levels:

    none:   no completed turn before the target user message
    filler: a repeated-token user/assistant turn
    demo:   the fixed meaningful user/assistant turn from the prior probe

The filler and demo histories add exactly the same number of official-template
tokens and keep the same user role, assistant role, assistant EOS, target
position, and generation prompt. This tests whether the prior attenuation of
the system edge requires the specific demo text. It does not identify a pure
distance effect because filler token identity and the added history remain
coupled.

The layer forward path is intentionally reused from the preceding audited
probe. This file owns the six-cell condition renderer, the shared-bootstrap
2 x 3 statistics, the control metadata, and the finalized result contract.
"""

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_factorial_probe as base


SYSTEM_MESSAGE = base.SYSTEM_MESSAGE
DEMO_USER = base.DEMO_USER
DEMO_ASSISTANT = base.DEMO_ASSISTANT
FILLER_USER = "x x x x x x x x x"
FILLER_ASSISTANT = "x"

HISTORY_LEVELS = ("none", "filler", "demo")
CONDITIONS = (
    "s0_none",
    "s1_none",
    "s0_filler",
    "s1_filler",
    "s0_demo",
    "s1_demo",
)
FACTORS = {
    "s0_none": {"system": 0, "history": "none"},
    "s1_none": {"system": 1, "history": "none"},
    "s0_filler": {"system": 0, "history": "filler"},
    "s1_filler": {"system": 1, "history": "filler"},
    "s0_demo": {"system": 0, "history": "demo"},
    "s1_demo": {"system": 1, "history": "demo"},
}
SYSTEM_CELLS = {
    "none": ("s0_none", "s1_none"),
    "filler": ("s0_filler", "s1_filler"),
    "demo": ("s0_demo", "s1_demo"),
}
COMPARISONS = (
    ("system_none", "s0_none", "s1_none"),
    ("system_filler", "s0_filler", "s1_filler"),
    ("system_demo", "s0_demo", "s1_demo"),
    ("filler_at_s0", "s0_none", "s0_filler"),
    ("filler_at_s1", "s1_none", "s1_filler"),
    ("demo_at_s0", "s0_none", "s0_demo"),
    ("demo_at_s1", "s1_none", "s1_demo"),
    ("demo_vs_filler_s0", "s0_filler", "s0_demo"),
    ("demo_vs_filler_s1", "s1_filler", "s1_demo"),
)
ALIGNMENT_COMPARISONS = (
    COMPARISONS[0],
    COMPARISONS[1],
    COMPARISONS[2],
    COMPARISONS[7],
    COMPARISONS[8],
)
SYSTEM_EDGE_CONTRASTS = {
    "filler_minus_none": ("none", "filler"),
    "demo_minus_none": ("none", "demo"),
    "demo_minus_filler": ("filler", "demo"),
}


def condition_messages(
    content: str,
    condition: str,
) -> list[dict[str, str]]:
    factors = FACTORS[condition]
    messages: list[dict[str, str]] = []
    if factors["system"]:
        messages.append({"role": "system", "content": SYSTEM_MESSAGE})
    if factors["history"] == "filler":
        messages.extend(
            [
                {"role": "user", "content": FILLER_USER},
                {"role": "assistant", "content": FILLER_ASSISTANT},
            ]
        )
    elif factors["history"] == "demo":
        messages.extend(
            [
                {"role": "user", "content": DEMO_USER},
                {"role": "assistant", "content": DEMO_ASSISTANT},
            ]
        )
    messages.append({"role": "user", "content": content})
    return messages


def history_control_domain(
    loads: dict[str, np.ndarray],
    mode: str,
    replicates: int,
    seed: int,
    scope: str,
) -> dict[str, Any]:
    """Compute all six cells with one shared source-bootstrap matrix."""
    shapes = {value.shape for value in loads.values()}
    if len(shapes) != 1:
        raise ValueError(f"history-control 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_metrics = {
        condition: base.metric_vector(value)
        for condition, value in point.items()
    }
    boot_metrics = {
        condition: base.metric_vector(value)
        for condition, value in boot.items()
    }

    def system_edges(
        values: dict[str, np.ndarray],
    ) -> dict[str, np.ndarray]:
        return {
            history: values[after] - values[before]
            for history, (before, after) in SYSTEM_CELLS.items()
        }

    def edge_contrasts(
        edges: dict[str, np.ndarray],
    ) -> dict[str, np.ndarray]:
        return {
            name: edges[after] - edges[before]
            for name, (before, after) in SYSTEM_EDGE_CONTRASTS.items()
        }

    metric_system_edges: dict[str, Any] = {}
    metric_system_edge_contrasts: dict[str, Any] = {}
    for metric in point_metrics["s0_none"]:
        point_edges = system_edges(
            {
                condition: point_metrics[condition][metric]
                for condition in CONDITIONS
            }
        )
        boot_edges = system_edges(
            {
                condition: boot_metrics[condition][metric]
                for condition in CONDITIONS
            }
        )
        metric_system_edges[metric] = {
            history: {
                "point": float(point_edges[history][0]),
                "ci95": base.interval(boot_edges[history]),
            }
            for history in HISTORY_LEVELS
        }
        point_contrasts = edge_contrasts(point_edges)
        boot_contrasts = edge_contrasts(boot_edges)
        metric_system_edge_contrasts[metric] = {
            name: {
                "point": float(point_contrasts[name][0]),
                "ci95": base.interval(boot_contrasts[name]),
            }
            for name in SYSTEM_EDGE_CONTRASTS
        }

    point_vectors = system_edges(point)
    boot_vectors = system_edges(boot)
    point_vector_contrasts = edge_contrasts(point_vectors)
    boot_vector_contrasts = edge_contrasts(boot_vectors)
    distribution_system_edge_contrasts = {}
    for name in SYSTEM_EDGE_CONTRASTS:
        point_magnitude = 0.5 * np.abs(
            point_vector_contrasts[name]
        ).sum()
        boot_magnitude = 0.5 * np.abs(
            boot_vector_contrasts[name]
        ).sum(axis=1)
        distribution_system_edge_contrasts[name] = {
            "half_l1_magnitude": {
                "point": float(point_magnitude),
                "ci95": base.interval(boot_magnitude),
            },
            "expert_share_difference_in_system_edges": (
                point_vector_contrasts[name].tolist()
            ),
            "expert_share_difference_in_system_edges_ci95": (
                base.interval(boot_vector_contrasts[name])
            ),
        }

    system_edge_distances: dict[str, Any] = {}
    point_tv: dict[str, float] = {}
    boot_tv: dict[str, np.ndarray] = {}
    point_jsd: dict[str, float] = {}
    boot_jsd: dict[str, np.ndarray] = {}
    for history, (before, after) in SYSTEM_CELLS.items():
        point_delta = point[after] - point[before]
        point_tv[history] = float(0.5 * np.abs(point_delta).sum())
        boot_tv[history] = 0.5 * np.abs(
            boot[after] - boot[before]
        ).sum(axis=1)
        point_jsd[history] = float(
            base.js_divergence(point[before], point[after])[0]
        )
        boot_jsd[history] = base.js_divergence(
            boot[before],
            boot[after],
        )
        system_edge_distances[history] = {
            "total_variation": {
                "point": point_tv[history],
                "ci95": base.interval(boot_tv[history]),
            },
            "js_divergence": {
                "point": point_jsd[history],
                "ci95": base.interval(boot_jsd[history]),
                "unit": "nats",
            },
        }

    system_edge_distance_contrasts = {}
    for name, (before, after) in SYSTEM_EDGE_CONTRASTS.items():
        system_edge_distance_contrasts[name] = {
            "total_variation_delta": {
                "point": point_tv[after] - point_tv[before],
                "ci95": base.interval(boot_tv[after] - boot_tv[before]),
            },
            "js_divergence_delta": {
                "point": point_jsd[after] - point_jsd[before],
                "ci95": base.interval(boot_jsd[after] - boot_jsd[before]),
                "unit": "nats",
            },
        }

    lexical_replacement = {}
    for system in (0, 1):
        before = f"s{system}_filler"
        after = f"s{system}_demo"
        point_delta = point[after] - point[before]
        tv_boot = 0.5 * np.abs(
            boot[after] - boot[before]
        ).sum(axis=1)
        jsd_boot = base.js_divergence(boot[before], boot[after])
        lexical_replacement[f"at_s{system}"] = {
            "total_variation": {
                "point": float(0.5 * np.abs(point_delta).sum()),
                "ci95": base.interval(tv_boot),
            },
            "js_divergence": {
                "point": float(
                    base.js_divergence(point[before], point[after])[0]
                ),
                "ci95": base.interval(jsd_boot),
                "unit": "nats",
            },
        }

    return {
        "metric_system_edges": metric_system_edges,
        "metric_system_edge_contrasts": metric_system_edge_contrasts,
        "system_edge_distances": system_edge_distances,
        "system_edge_distance_contrasts": (
            system_edge_distance_contrasts
        ),
        "distribution_system_edge_contrasts": (
            distribution_system_edge_contrasts
        ),
        "lexical_replacement": lexical_replacement,
    }


def layer_statistics(
    prompt_rows: list[dict[str, Any]],
    replicates: int,
    seed: int,
    layer_index: int,
) -> dict[str, Any]:
    scopes = {}
    for load_scope, load_key in (
        ("full_input", "full_load"),
        ("target_content", "content_load"),
    ):
        modes = {}
        for mode in ("token_weighted", "prompt_balanced"):
            conditions: dict[str, Any] = {}
            domain_loads: dict[str, dict[str, np.ndarray]] = {}
            for domain in base.DOMAIN_ORDER:
                rows = [
                    row for row in prompt_rows
                    if row["domain"] == domain
                ]
                domain_loads[domain] = {
                    condition: np.asarray(
                        [
                            row["conditions"][condition][load_key]
                            for row in rows
                        ],
                        dtype=np.int64,
                    )
                    for condition in CONDITIONS
                }
                for condition in CONDITIONS:
                    conditions.setdefault(condition, {})[domain] = (
                        base.bootstrap_domain(
                            domain_loads[domain][condition],
                            mode,
                            replicates,
                            seed,
                            (
                                f"layer={layer_index}|scope={load_scope}|"
                                f"mode={mode}|condition={condition}|"
                                f"domain={domain}"
                            ),
                        )
                    )

            comparisons: dict[str, Any] = {}
            for comparison, before, after in COMPARISONS:
                comparisons[comparison] = {}
                for domain in base.DOMAIN_ORDER:
                    comparisons[comparison][domain] = base.paired_domain(
                        domain_loads[domain][before],
                        domain_loads[domain][after],
                        mode,
                        replicates,
                        seed,
                        (
                            f"layer={layer_index}|scope={load_scope}|"
                            f"mode={mode}|comparison={comparison}|"
                            f"domain={domain}"
                        ),
                    )

            control = {
                domain: history_control_domain(
                    domain_loads[domain],
                    mode,
                    replicates,
                    seed,
                    (
                        f"layer={layer_index}|scope={load_scope}|"
                        f"mode={mode}|history_control|domain={domain}"
                    ),
                )
                for domain in base.DOMAIN_ORDER
            }
            modes[mode] = {
                "conditions": conditions,
                "comparisons": comparisons,
                "history_control": control,
            }
        scopes[load_scope] = {"modes": modes}
    return scopes


def install_control_contract() -> None:
    """Install the six-cell renderer/statistics into the audited runner."""
    base.CONDITIONS = CONDITIONS
    base.FACTORS = FACTORS
    base.COMPARISONS = COMPARISONS
    base.ALIGNMENT_COMPARISONS = ALIGNMENT_COMPARISONS
    base.condition_messages = condition_messages
    base.layer_statistics = layer_statistics


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 = json.loads(path.read_text(encoding="utf-8"))
    result["evidence_identity"] = (
        "X / official BF16 weights, official tokenizer chat template, "
        "paired 2x3 history-distance control on local truncated forward"
    )
    boundary = result["boundary"]
    boundary.pop("factorial_claim", None)
    boundary.update(
        {
            "history_control_claim": (
                "filler and demo histories have identical official-template "
                "token increments, roles, assistant EOS, and target position; "
                "their text identities differ"
            ),
            "filler_is_semantics_free": False,
            "pure_distance_isolated": False,
            "lexical_replacement_control": True,
            "causal_boundary": (
                "none-to-filler still couples added history, repeated filler "
                "tokens, and distance; filler-to-demo isolates replacement of "
                "the fixed history text only within this protocol"
            ),
        }
    )

    old_contract = result.pop("message_history_contract")
    selected = result["corpus_contract"]["selected"]
    increments = {}
    for system in (0, 1):
        base_condition = f"s{system}_none"
        for history in ("filler", "demo"):
            condition = f"s{system}_{history}"
            deltas = [
                row["conditions"][condition]["tokens"]
                - row["conditions"][base_condition]["tokens"]
                for row in selected
            ]
            increments[f"{condition}_minus_{base_condition}"] = {
                "min": min(deltas),
                "max": max(deltas),
                "all_equal": len(set(deltas)) == 1,
            }
    result["history_control_contract"] = {
        "chat_template_revision": base.MODEL_REVISION,
        "chat_template": old_contract["chat_template"],
        "chat_template_sha256": old_contract["chat_template_sha256"],
        "system_message": SYSTEM_MESSAGE,
        "system_message_sha256": base.text_sha256(SYSTEM_MESSAGE),
        "demo_user": DEMO_USER,
        "demo_user_sha256": base.text_sha256(DEMO_USER),
        "demo_assistant": DEMO_ASSISTANT,
        "demo_assistant_sha256": base.text_sha256(DEMO_ASSISTANT),
        "filler_user": FILLER_USER,
        "filler_user_sha256": base.text_sha256(FILLER_USER),
        "filler_assistant": FILLER_ASSISTANT,
        "filler_assistant_sha256": base.text_sha256(FILLER_ASSISTANT),
        "target_role": "user",
        "add_generation_prompt": True,
        "conditions": FACTORS,
        "comparisons": [
            {"name": name, "before": before, "after": after}
            for name, before, after in COMPARISONS
        ],
        "system_edge_contrasts": {
            name: {
                "before_history": before,
                "after_history": after,
                "definition": (
                    f"system edge at {after} minus system edge at {before}"
                ),
            }
            for name, (before, after) in SYSTEM_EDGE_CONTRASTS.items()
        },
        "token_increment_validation": increments,
        "scope_split": {
            "full_input": (
                "all rendered BOS, system/history, target, newlines, EOS, "
                "and generation-prompt tokens"
            ),
            "target_content": (
                "exact intersection of (relative character span, token ID) "
                "inside target user content across all six conditions"
            ),
        },
    }

    inference = result["inference_contract"]
    inference["batch_grouping"] = (
        "all six control variants of one source prompt execute in the same "
        "right-padded batch"
    )
    statistical = result["statistical_contract"]
    statistical["paired_indices"] = (
        "one sampled source-prompt index matrix is reused across all six "
        "cells for every history-control contrast within each "
        "domain/layer/scope/mode"
    )
    statistical.pop("interaction_distribution_magnitude", None)
    statistical["system_edge_contrast_distribution_magnitude"] = (
        "0.5 * L1 norm of the signed difference between two system-edge "
        "expert-share vectors; this is not labeled standard TV"
    )
    result["schema_version"] = 2
    path.write_text(
        json.dumps(result, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    return result


def main() -> None:
    install_control_contract()
    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"
                ],
                "input_tokens_by_condition": result[
                    "inference_contract"
                ]["input_tokens_by_condition"],
                "total_routes": result["inference_contract"][
                    "total_routes_all_conditions_all_moe_layers"
                ],
            },
            indent=2,
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()