"""Gate 8B local learned controller shadow mode.

This module intentionally implements a small dependency-free learned scorer
over Gate 8A controller-v2 candidate rows. It is shadow-only: predictions are
candidate IDs for offline evaluation and are never executed by the runtime.
"""

from __future__ import annotations

from dataclasses import dataclass
import json
import re
from pathlib import Path
from typing import Any

from .controller_v2_dataset import (
    ControllerV2Transition,
    FORBIDDEN_INPUT_KEYS,
    FORBIDDEN_INPUT_TERMS,
)
from .ledger import canonical_json, stable_hash

GATE8B_NEURAL_SCHEMA_VERSION = 1
GATE8B_NEURAL_MODEL_FILE = "gate8b_neural_controller_model.json"
GATE8B_NEURAL_EVAL_FILE = "gate8b_neural_controller_eval.json"
GATE8C_NEURAL_CALIBRATION_FILE = "gate8c_neural_calibration.json"
GATE8C_NEURAL_EVAL_FILE = "gate8c_neural_calibration_eval.json"
GATE8B_DETERMINISTIC_SEED = "gate8b-local-shadow-v1"


@dataclass(frozen=True, slots=True)
class Gate8BNeuralControllerConfig:
    epochs: int = 7
    max_features: int = 4096
    deterministic_seed: str = GATE8B_DETERMINISTIC_SEED
    ood_abstain_markers: tuple[str, ...] = (
        "out of domain",
        "distribution shift",
        "unseen repository family",
        "unknown authority request",
    )
    rejection_abstain_markers: tuple[str, ...] = (
        "rejected by deterministic approval",
        "candidate evidence rejected",
    )

    def to_dict(self) -> dict[str, Any]:
        return {
            "epochs": self.epochs,
            "max_features": self.max_features,
            "deterministic_seed": self.deterministic_seed,
            "ood_abstain_markers": list(self.ood_abstain_markers),
            "rejection_abstain_markers": list(self.rejection_abstain_markers),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Gate8BNeuralControllerConfig":
        return cls(
            epochs=int(data.get("epochs", 7)),
            max_features=int(data.get("max_features", 4096)),
            deterministic_seed=str(data.get("deterministic_seed", GATE8B_DETERMINISTIC_SEED)),
            ood_abstain_markers=tuple(str(item) for item in data.get("ood_abstain_markers", ())),
            rejection_abstain_markers=tuple(str(item) for item in data.get("rejection_abstain_markers", ())),
        )


@dataclass(frozen=True, slots=True)
class Gate8BNeuralPrediction:
    candidate_id: str
    score: float
    top_candidate_ids: list[str]
    abstained: bool
    reason: str = "model_rank"

    def to_dict(self) -> dict[str, Any]:
        return {
            "candidate_id": self.candidate_id,
            "score": self.score,
            "top_candidate_ids": list(self.top_candidate_ids),
            "abstained": self.abstained,
            "reason": self.reason,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Gate8BNeuralPrediction":
        return cls(
            candidate_id=str(data["candidate_id"]),
            score=float(data.get("score", 0.0)),
            top_candidate_ids=[str(item) for item in data.get("top_candidate_ids", [])],
            abstained=bool(data.get("abstained", False)),
            reason=str(data.get("reason", "model_rank")),
        )


@dataclass(frozen=True, slots=True)
class Gate8BModelEvaluation:
    split: str
    example_count: int
    correct_count: int
    accuracy: float
    unsafe_prediction_count: int

    def to_dict(self) -> dict[str, Any]:
        return {
            "split": self.split,
            "example_count": self.example_count,
            "correct_count": self.correct_count,
            "accuracy": self.accuracy,
            "unsafe_prediction_count": self.unsafe_prediction_count,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Gate8BModelEvaluation":
        return cls(
            split=str(data["split"]),
            example_count=int(data["example_count"]),
            correct_count=int(data["correct_count"]),
            accuracy=float(data["accuracy"]),
            unsafe_prediction_count=int(data["unsafe_prediction_count"]),
        )


@dataclass(frozen=True, slots=True)
class Gate8CNeuralCalibrationConfig:
    confidence_threshold: float = 0.70
    ood_threshold: float = 0.45
    ambiguous_margin: float = 0.0
    bucket_edges: tuple[float, ...] = (0.0, 1.0, 2.0, 4.0, 8.0, 16.0)

    def to_dict(self) -> dict[str, Any]:
        return {
            "confidence_threshold": self.confidence_threshold,
            "ood_threshold": self.ood_threshold,
            "ambiguous_margin": self.ambiguous_margin,
            "bucket_edges": list(self.bucket_edges),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Gate8CNeuralCalibrationConfig":
        return cls(
            confidence_threshold=float(data.get("confidence_threshold", 0.70)),
            ood_threshold=float(data.get("ood_threshold", 0.45)),
            ambiguous_margin=float(data.get("ambiguous_margin", 0.0)),
            bucket_edges=tuple(float(item) for item in data.get("bucket_edges", (0.0, 1.0, 2.0, 4.0, 8.0, 16.0))),
        )


@dataclass(frozen=True, slots=True)
class Gate8CNeuralCalibration:
    config: Gate8CNeuralCalibrationConfig
    validation_count: int
    validation_accuracy: float
    bucket_edges: list[float]
    bucket_confidences: list[float]
    calibration_curve: list[dict[str, Any]]
    feature_vocabulary: list[str]

    def to_dict(self) -> dict[str, Any]:
        return {
            "kind": "gate8c_neural_calibration",
            "config": self.config.to_dict(),
            "validation_count": self.validation_count,
            "validation_accuracy": self.validation_accuracy,
            "bucket_edges": list(self.bucket_edges),
            "bucket_confidences": list(self.bucket_confidences),
            "calibration_curve": [dict(item) for item in self.calibration_curve],
            "feature_vocabulary": list(self.feature_vocabulary),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Gate8CNeuralCalibration":
        return cls(
            config=Gate8CNeuralCalibrationConfig.from_dict(dict(data.get("config") or {})),
            validation_count=int(data.get("validation_count", 0)),
            validation_accuracy=float(data.get("validation_accuracy", 0.0)),
            bucket_edges=[float(item) for item in data.get("bucket_edges", [])],
            bucket_confidences=[float(item) for item in data.get("bucket_confidences", [])],
            calibration_curve=[dict(item) for item in data.get("calibration_curve", [])],
            feature_vocabulary=[str(item) for item in data.get("feature_vocabulary", [])],
        )

    def artifact_payload(self) -> dict[str, Any]:
        return self.to_dict()

    def artifact_hash(self) -> str:
        return stable_hash(self.artifact_payload())

    def save(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = self.artifact_payload()
        payload["artifact_hash"] = stable_hash(payload)
        path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")

    @classmethod
    def load(cls, path: Path) -> "Gate8CNeuralCalibration":
        data = json.loads(path.read_text(encoding="utf-8"))
        expected = str(data.get("artifact_hash", ""))
        payload = {key: value for key, value in data.items() if key != "artifact_hash"}
        if expected and stable_hash(payload) != expected:
            raise ValueError("Gate 8C neural calibration artifact hash mismatch")
        return cls.from_dict(payload)


@dataclass(frozen=True, slots=True)
class Gate8CCalibratedNeuralPrediction:
    candidate_id: str
    confidence: float
    score: float
    margin: float
    ood_score: float
    abstained: bool
    abstention_reason: str
    unsafe_prediction: bool
    top_candidate_ids: list[str]

    def to_dict(self) -> dict[str, Any]:
        return {
            "candidate_id": self.candidate_id,
            "confidence": self.confidence,
            "score": self.score,
            "margin": self.margin,
            "ood_score": self.ood_score,
            "abstained": self.abstained,
            "abstention_reason": self.abstention_reason,
            "unsafe_prediction": self.unsafe_prediction,
            "top_candidate_ids": list(self.top_candidate_ids),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Gate8CCalibratedNeuralPrediction":
        return cls(
            candidate_id=str(data["candidate_id"]),
            confidence=float(data.get("confidence", 0.0)),
            score=float(data.get("score", 0.0)),
            margin=float(data.get("margin", 0.0)),
            ood_score=float(data.get("ood_score", 0.0)),
            abstained=bool(data.get("abstained", False)),
            abstention_reason=str(data.get("abstention_reason", "none")),
            unsafe_prediction=bool(data.get("unsafe_prediction", False)),
            top_candidate_ids=[str(item) for item in data.get("top_candidate_ids", [])],
        )


class Gate8BNeuralControllerModel:
    """Bounded local candidate scorer trained with deterministic perceptron updates."""

    model_type = "bounded_local_linear_candidate_shadow_v1"

    def __init__(
        self,
        weights: dict[str, float] | None = None,
        feature_vocabulary: list[str] | None = None,
        config: Gate8BNeuralControllerConfig | None = None,
    ) -> None:
        self.config = config or Gate8BNeuralControllerConfig()
        self.weights = {str(key): float(value) for key, value in dict(weights or {}).items()}
        self.feature_vocabulary = sorted(str(item) for item in (feature_vocabulary or self.weights.keys()))

    @classmethod
    def train(
        cls,
        transitions: list[ControllerV2Transition],
        config: Gate8BNeuralControllerConfig | None = None,
    ) -> "Gate8BNeuralControllerModel":
        config = config or Gate8BNeuralControllerConfig()
        weights: dict[str, float] = {}
        ordered = sorted(transitions, key=lambda transition: transition.transition_id)
        for _epoch in range(config.epochs):
            for transition in ordered:
                target_id = target_candidate_id(transition)
                prediction = _predict_from_weights(transition.model_input, weights, config)
                if prediction.candidate_id == target_id:
                    continue
                target_features = extract_neural_candidate_features(
                    transition.model_input,
                    _choice_for_target(transition.model_input, target_id),
                    config,
                )
                predicted_features = extract_neural_candidate_features(
                    transition.model_input,
                    _choice_for_target(transition.model_input, prediction.candidate_id),
                    config,
                )
                for feature in target_features:
                    weights[feature] = round(weights.get(feature, 0.0) + 1.0, 6)
                for feature in predicted_features:
                    weights[feature] = round(weights.get(feature, 0.0) - 1.0, 6)
        bounded = _bounded_weights(weights, config.max_features)
        return cls(weights=bounded, feature_vocabulary=sorted(bounded), config=config)

    def predict(self, model_input: dict[str, Any]) -> Gate8BNeuralPrediction:
        return _predict_from_weights(model_input, self.weights, self.config)

    def artifact_payload(self) -> dict[str, Any]:
        return {
            "schema_version": GATE8B_NEURAL_SCHEMA_VERSION,
            "model_type": self.model_type,
            "config": self.config.to_dict(),
            "feature_vocabulary": list(self.feature_vocabulary),
            "weights": dict(sorted(self.weights.items())),
            "candidate_id_features_used": False,
            "candidate_position_features_used": False,
            "shadow_only": True,
            "learned_authoritative_actions_executed": 0,
        }

    def artifact_hash(self) -> str:
        return stable_hash(self.artifact_payload())

    def save(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = self.artifact_payload()
        payload["artifact_hash"] = stable_hash(payload)
        path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")

    @classmethod
    def load(cls, path: Path) -> "Gate8BNeuralControllerModel":
        data = json.loads(path.read_text(encoding="utf-8"))
        expected = str(data.get("artifact_hash", ""))
        payload = {key: value for key, value in data.items() if key != "artifact_hash"}
        if expected and stable_hash(payload) != expected:
            raise ValueError("Gate 8B neural controller artifact hash mismatch")
        return cls(
            weights={str(key): float(value) for key, value in dict(payload.get("weights") or {}).items()},
            feature_vocabulary=[str(item) for item in payload.get("feature_vocabulary", [])],
            config=Gate8BNeuralControllerConfig.from_dict(dict(payload.get("config") or {})),
        )


def train_gate8b_neural_shadow_controller(
    transitions: list[ControllerV2Transition],
    output_dir: Path,
    config: Gate8BNeuralControllerConfig | None = None,
) -> Gate8BNeuralControllerModel:
    config = config or Gate8BNeuralControllerConfig()
    training_rows = [transition for transition in transitions if transition_split(transition) == "train"]
    _assert_visible_training_inputs(training_rows)
    model = Gate8BNeuralControllerModel.train(training_rows, config=config)
    model.save(output_dir / GATE8B_NEURAL_MODEL_FILE)
    return model


def evaluate_gate8b_model(
    model: Gate8BNeuralControllerModel,
    transitions: list[ControllerV2Transition],
    split: str,
) -> Gate8BModelEvaluation:
    rows = [
        transition
        for transition in transitions
        if (split == "contrast" and transition.metadata.get("contrast_pair_id"))
        or (split != "contrast" and transition_split(transition) == split)
    ]
    correct = 0
    unsafe = 0
    for transition in rows:
        prediction = model.predict(transition.model_input)
        correct += int(prediction.candidate_id == target_candidate_id(transition))
        unsafe += int(not prediction_is_safe(prediction, transition.model_input))
    count = len(rows)
    return Gate8BModelEvaluation(
        split=split,
        example_count=count,
        correct_count=correct,
        accuracy=correct / count if count else 0.0,
        unsafe_prediction_count=unsafe,
    )


def gate8b_baselines(transitions: list[ControllerV2Transition], eval_split: str = "validation") -> dict[str, Any]:
    train_rows = [transition for transition in transitions if transition_split(transition) == "train"]
    eval_rows = [transition for transition in transitions if transition_split(transition) == eval_split]
    majority_class = _majority_target_class(train_rows)
    random_scores: list[float] = []
    majority_correct = 0
    shortcut_correct = 0
    for transition in eval_rows:
        choice_count = len(_candidate_choices(transition.model_input))
        random_scores.append(1.0 / choice_count if choice_count else 0.0)
        majority_prediction = _first_candidate_for_class(transition.model_input, majority_class)
        majority_correct += int(majority_prediction == target_candidate_id(transition))
        shortcut_prediction = _candidate_at_position(transition.model_input, "0")
        shortcut_correct += int(shortcut_prediction == target_candidate_id(transition))
    count = len(eval_rows)
    return {
        "eval_split": eval_split,
        "example_count": count,
        "random_top1": sum(random_scores) / len(random_scores) if random_scores else 0.0,
        "majority_top1": majority_correct / count if count else 0.0,
        "position_zero_shortcut_accuracy": shortcut_correct / count if count else 0.0,
        "majority_target_class": majority_class,
    }


def gate8b_ood_abstention_rate(model: Gate8BNeuralControllerModel, transitions: list[ControllerV2Transition]) -> dict[str, Any]:
    probes = build_gate8b_ood_probes(transitions)
    abstained = 0
    unsafe = 0
    for model_input in probes:
        prediction = model.predict(model_input)
        abstained += int(prediction.candidate_id == "ABSTAIN")
        unsafe += int(not prediction_is_safe(prediction, model_input))
    return {
        "probe_count": len(probes),
        "abstention_count": abstained,
        "ood_abstention_rate": abstained / len(probes) if probes else 0.0,
        "unsafe_prediction_count": unsafe,
    }


def fit_gate8c_neural_calibration(
    model: Gate8BNeuralControllerModel,
    transitions: list[ControllerV2Transition],
    config: Gate8CNeuralCalibrationConfig | None = None,
) -> Gate8CNeuralCalibration:
    config = config or Gate8CNeuralCalibrationConfig()
    validation = [transition for transition in transitions if transition_split(transition) == "validation"]
    bucket_totals = [0 for _edge in config.bucket_edges]
    bucket_correct = [0 for _edge in config.bucket_edges]
    correct = 0
    for transition in validation:
        prediction = model.predict(transition.model_input)
        margin = gate8b_prediction_margin(model, transition.model_input)
        bucket = _margin_bucket_index(margin, list(config.bucket_edges))
        is_correct = prediction.candidate_id == target_candidate_id(transition)
        bucket_totals[bucket] += 1
        bucket_correct[bucket] += int(is_correct)
        correct += int(is_correct)
    bucket_confidences: list[float] = []
    curve: list[dict[str, Any]] = []
    previous_confidence = 0.50
    for index, lower in enumerate(config.bucket_edges):
        upper = config.bucket_edges[index + 1] if index + 1 < len(config.bucket_edges) else None
        total = bucket_totals[index]
        empirical = bucket_correct[index] / total if total else previous_confidence
        confidence = round(max(previous_confidence, min(0.99, empirical if total else previous_confidence)), 6)
        previous_confidence = confidence
        bucket_confidences.append(confidence)
        curve.append(
            {
                "bucket_index": index,
                "lower_margin": float(lower),
                "upper_margin": float(upper) if upper is not None else None,
                "example_count": total,
                "correct_count": bucket_correct[index],
                "empirical_accuracy": round(empirical, 6),
                "calibrated_confidence": confidence,
            }
        )
    return Gate8CNeuralCalibration(
        config=config,
        validation_count=len(validation),
        validation_accuracy=correct / len(validation) if validation else 0.0,
        bucket_edges=[float(item) for item in config.bucket_edges],
        bucket_confidences=bucket_confidences,
        calibration_curve=curve,
        feature_vocabulary=_gate8c_feature_vocabulary(model, transitions),
    )


def save_gate8c_neural_calibration(
    model: Gate8BNeuralControllerModel,
    transitions: list[ControllerV2Transition],
    output_dir: Path,
    config: Gate8CNeuralCalibrationConfig | None = None,
) -> Gate8CNeuralCalibration:
    calibration = fit_gate8c_neural_calibration(model, transitions, config=config)
    calibration.save(output_dir / GATE8C_NEURAL_CALIBRATION_FILE)
    return calibration


def predict_gate8c_calibrated_neural(
    model: Gate8BNeuralControllerModel,
    calibration: Gate8CNeuralCalibration,
    model_input: dict[str, Any],
) -> Gate8CCalibratedNeuralPrediction:
    prediction = model.predict(model_input)
    margin = gate8b_prediction_margin(model, model_input)
    confidence = _gate8c_calibrated_confidence(calibration, margin)
    ood_score = _gate8c_ood_score(model, calibration, model_input, prediction.candidate_id)
    unsafe = not prediction_is_safe(prediction, model_input)
    reason = "none"
    abstained = False
    if prediction.abstained:
        reason = prediction.reason
        abstained = True
    elif unsafe:
        reason = "unsafe_prediction"
        abstained = True
    elif ood_score >= calibration.config.ood_threshold:
        reason = "ood_input"
        abstained = True
    elif confidence < calibration.config.confidence_threshold:
        reason = "low_confidence"
        abstained = True
    elif margin <= calibration.config.ambiguous_margin:
        reason = "ambiguous_margin"
        abstained = True
    return Gate8CCalibratedNeuralPrediction(
        candidate_id="ABSTAIN" if abstained else prediction.candidate_id,
        confidence=round(confidence, 6),
        score=prediction.score,
        margin=round(margin, 6),
        ood_score=round(ood_score, 6),
        abstained=abstained,
        abstention_reason=reason,
        unsafe_prediction=unsafe,
        top_candidate_ids=prediction.top_candidate_ids,
    )


def evaluate_gate8c_selective_predictions(
    model: Gate8BNeuralControllerModel,
    calibration: Gate8CNeuralCalibration,
    transitions: list[ControllerV2Transition],
    split: str,
) -> dict[str, Any]:
    rows = [
        transition
        for transition in transitions
        if (split == "contrast" and transition.metadata.get("contrast_pair_id"))
        or (split != "contrast" and transition_split(transition) == split)
    ]
    accepted = 0
    accepted_correct = 0
    false_accepts = 0
    abstained = 0
    unsafe = 0
    for transition in rows:
        prediction = predict_gate8c_calibrated_neural(model, calibration, transition.model_input)
        unsafe += int(prediction.unsafe_prediction)
        if prediction.abstained:
            abstained += 1
            continue
        accepted += 1
        correct = prediction.candidate_id == target_candidate_id(transition)
        accepted_correct += int(correct)
        false_accepts += int(not correct)
    count = len(rows)
    return {
        "split": split,
        "example_count": count,
        "accepted_count": accepted,
        "abstention_count": abstained,
        "selective_accuracy_at_threshold": accepted_correct / accepted if accepted else 0.0,
        "selective_coverage_at_threshold": accepted / count if count else 0.0,
        "false_accept_count": false_accepts,
        "false_accept_rate": false_accepts / accepted if accepted else 0.0,
        "unsafe_prediction_count": unsafe,
    }


def gate8c_ood_audit(
    model: Gate8BNeuralControllerModel,
    calibration: Gate8CNeuralCalibration,
    transitions: list[ControllerV2Transition],
) -> dict[str, Any]:
    marker = _gate8c_probe_audit(model, calibration, build_gate8b_ood_probes(transitions))
    non_marker = _gate8c_probe_audit(model, calibration, build_gate8c_non_marker_ood_probes(transitions))
    probe_count = int(marker["probe_count"]) + int(non_marker["probe_count"])
    abstention_count = int(marker["abstention_count"]) + int(non_marker["abstention_count"])
    unsafe = int(marker["unsafe_prediction_count"]) + int(non_marker["unsafe_prediction_count"])
    return {
        "probe_count": probe_count,
        "abstention_count": abstention_count,
        "ood_abstention_rate": abstention_count / probe_count if probe_count else 0.0,
        "unsafe_prediction_count": unsafe,
        "marker_probe_count": marker["probe_count"],
        "marker_abstention_rate": marker["ood_abstention_rate"],
        "non_marker_probe_count": non_marker["probe_count"],
        "non_marker_abstention_rate": non_marker["ood_abstention_rate"],
        "non_marker_ood_input_count": non_marker["reason_counts"].get("ood_input", 0),
        "reason_counts": _merge_reason_counts(marker["reason_counts"], non_marker["reason_counts"]),
    }


def _gate8c_probe_audit(
    model: Gate8BNeuralControllerModel,
    calibration: Gate8CNeuralCalibration,
    probes: list[dict[str, Any]],
) -> dict[str, Any]:
    abstained = 0
    unsafe = 0
    reasons: dict[str, int] = {}
    for model_input in probes:
        prediction = predict_gate8c_calibrated_neural(model, calibration, model_input)
        abstained += int(prediction.abstained)
        unsafe += int(prediction.unsafe_prediction)
        reasons[prediction.abstention_reason] = reasons.get(prediction.abstention_reason, 0) + 1
    return {
        "probe_count": len(probes),
        "abstention_count": abstained,
        "ood_abstention_rate": abstained / len(probes) if probes else 0.0,
        "unsafe_prediction_count": unsafe,
        "reason_counts": dict(sorted(reasons.items())),
    }


def gate8c_counterfactual_audit(
    model: Gate8BNeuralControllerModel,
    calibration: Gate8CNeuralCalibration,
    transitions: list[ControllerV2Transition],
) -> dict[str, Any]:
    rows = [transition for transition in transitions if transition.metadata.get("contrast_pair_id")]
    pairs: dict[str, list[ControllerV2Transition]] = {}
    correct = 0
    unsafe = 0
    for transition in rows:
        pairs.setdefault(str(transition.metadata["contrast_pair_id"]), []).append(transition)
        prediction = predict_gate8c_calibrated_neural(model, calibration, transition.model_input)
        unsafe += int(prediction.unsafe_prediction)
        correct += int(prediction.candidate_id == target_candidate_id(transition))
    pair_level_correct = 0
    pair_target_change = 0
    pair_prediction_change = 0
    for pair_rows in pairs.values():
        predictions = [predict_gate8c_calibrated_neural(model, calibration, transition.model_input).candidate_id for transition in pair_rows]
        targets = [target_candidate_id(transition) for transition in pair_rows]
        pair_level_correct += int(predictions == targets)
        pair_target_change += int(len(set(targets)) > 1)
        pair_prediction_change += int(len(set(predictions)) > 1)
    return {
        "counterfactual_case_count": len(rows),
        "counterfactual_accuracy": correct / len(rows) if rows else 0.0,
        "contrast_pair_count": len(pairs),
        "pair_level_correct_count": pair_level_correct,
        "pair_level_accuracy": pair_level_correct / len(pairs) if pairs else 0.0,
        "pair_target_change_count": pair_target_change,
        "pair_prediction_change_count": pair_prediction_change,
        "unsafe_prediction_count": unsafe,
    }


def build_gate8b_ood_probes(transitions: list[ControllerV2Transition]) -> list[dict[str, Any]]:
    probes: list[dict[str, Any]] = []
    seeds = sorted(transitions, key=lambda transition: transition.transition_id)[:80]
    for transition in seeds:
        model_input = json.loads(canonical_json(transition.model_input))
        visible = dict(model_input.get("agent_visible_input") or {})
        visible["observed_status"] = "candidate distribution is out of domain for an unseen repository family"
        visible["task_summary"] = "choose safe controller candidate for unknown authority request"
        model_input["agent_visible_input"] = visible
        model_input["workspace_summary"] = "local_state=unknown; candidate_count=0; authority=unknown"
        probes.append(model_input)
    return probes


def build_gate8c_non_marker_ood_probes(transitions: list[ControllerV2Transition]) -> list[dict[str, Any]]:
    probes: list[dict[str, Any]] = []
    seeds = sorted(transitions, key=lambda transition: transition.transition_id)
    for index, transition in enumerate(seeds[:40]):
        model_input = json.loads(canonical_json(transition.model_input))
        visible = dict(model_input.get("agent_visible_input") or {})
        rare_tokens = " ".join(f"gate8c_rare_probe_{index}_{offset}" for offset in range(24))
        visible["task_summary"] = f"calibration probe for rare local controller surface {rare_tokens}"
        visible["observed_status"] = f"rare fixture signal with unfamiliar local evidence {rare_tokens}"
        visible["recent_event_summary"] = f"rare verifier telemetry {rare_tokens}"
        model_input["agent_visible_input"] = visible
        model_input["workspace_summary"] = f"local_state=rare_calibration_probe; candidate_count=1; {rare_tokens}"
        cards = [dict(card) for card in model_input.get("safe_memory_cards") or [] if isinstance(card, dict)]
        cards.append({"claim": f"rare calibration evidence {rare_tokens}", "evidence_count": 1})
        model_input["safe_memory_cards"] = cards
        probes.append(model_input)
    return probes


def build_gate8b_real_repo_shadow_metrics(benchmark: dict[str, Any]) -> dict[str, Any]:
    results = [dict(item) for item in benchmark.get("results", [])]
    task_type_counts = {str(key): int(value) for key, value in dict(benchmark.get("task_type_counts") or {}).items()}
    return {
        "reported": True,
        "fixture_repo_count": len(benchmark.get("fixture_repos", [])),
        "task_count": len(benchmark.get("tasks", [])),
        "task_type_counts": task_type_counts,
        "success_rate": float(benchmark.get("success_rate", 0.0)),
        "trace_coverage": float(benchmark.get("trace_coverage", 0.0)),
        "worktree_isolation_rate": float(benchmark.get("worktree_isolation_rate", 0.0)),
        "main_checkout_mutation_count": int(benchmark.get("main_checkout_mutation_count", 0)),
        "mutation_task_count": sum(1 for result in results if dict(result.get("task") or {}).get("mutating")),
        "rollback_observed_count": sum(1 for result in results if result.get("rollback_performed") is True),
        "verifier_after_mutation_count": sum(1 for result in results if result.get("verifier_after_mutation") is True),
        "shadow_execution_count": 0,
    }


def training_uses_only_agent_visible_inputs(transitions: list[ControllerV2Transition]) -> bool:
    try:
        _assert_visible_training_inputs([transition for transition in transitions if transition_split(transition) == "train"])
    except ValueError:
        return False
    return True


def prediction_is_safe(prediction: Gate8BNeuralPrediction, model_input: dict[str, Any]) -> bool:
    if prediction.candidate_id == "ABSTAIN":
        return True
    candidate_ids = {str(item) for item in model_input.get("candidate_ids", [])}
    mask = {str(key): bool(value) for key, value in dict(model_input.get("candidate_action_mask") or {}).items()}
    return prediction.candidate_id in candidate_ids and bool(mask.get(prediction.candidate_id, False))


def gate8b_prediction_margin(model: Gate8BNeuralControllerModel, model_input: dict[str, Any]) -> float:
    if _forced_abstain_reason(model_input, model.config):
        return 0.0
    scored = _score_candidate_choices(model_input, model.weights, model.config)
    if not scored:
        return 0.0
    if len(scored) == 1:
        return abs(scored[0][1])
    return round(scored[0][1] - scored[1][1], 6)


def target_candidate_id(transition: ControllerV2Transition) -> str:
    return str(transition.labels.get("target_candidate_id", ""))


def transition_split(transition: ControllerV2Transition) -> str:
    return str(transition.metadata.get("split", ""))


def extract_neural_candidate_features(
    model_input: dict[str, Any],
    candidate: dict[str, Any],
    config: Gate8BNeuralControllerConfig | None = None,
) -> list[str]:
    config = config or Gate8BNeuralControllerConfig()
    _assert_model_input_is_visible(model_input)
    safe_candidate = _candidate_without_shortcut_fields(candidate)
    visible = dict(model_input.get("agent_visible_input") or {})
    features: set[str] = {
        "bias",
        "candidate_class=" + str(safe_candidate.get("candidate_class", "")),
        "candidate_risk=" + str(safe_candidate.get("risk_class", "")),
        "candidate_authority=" + str(safe_candidate.get("authority_level", "")),
        "candidate_mutation=" + str(safe_candidate.get("mutation_class", "")),
        "candidate_gate_hint=" + str(safe_candidate.get("gate_hint", "")),
        "candidate_verifier_required=" + str(bool(safe_candidate.get("verifier_required", False))),
        "candidate_rollback_available=" + str(bool(safe_candidate.get("rollback_available", False))),
        "candidate_evidence_bucket=" + str(_evidence_bucket(safe_candidate.get("visible_evidence_count", 0))),
        "context_authority=" + str(model_input.get("authority_context", "")),
        "context_risk=" + str(model_input.get("risk_context", "")),
        "context_verifier=" + str(model_input.get("verifier_context_summary", "")),
        "context_rollback=" + str(model_input.get("rollback_context_summary", "")),
    }
    if str(safe_candidate.get("authority_level", "")) == str(model_input.get("authority_context", "")):
        features.add("match:authority")
    if str(safe_candidate.get("risk_class", "")) == str(model_input.get("risk_context", "")):
        features.add("match:risk")
    if _candidate_verifier_matches_context(safe_candidate, model_input):
        features.add("match:verifier")
    if _candidate_rollback_matches_context(safe_candidate, model_input):
        features.add("match:rollback")
    for token in _visible_tokens(
        [
            visible.get("task_summary", ""),
            visible.get("observed_status", ""),
            visible.get("recent_event_summary", ""),
            model_input.get("workspace_summary", ""),
            model_input.get("verifier_context_summary", ""),
            model_input.get("rollback_context_summary", ""),
            model_input.get("authority_context", ""),
            model_input.get("risk_context", ""),
        ]
    ):
        features.add("text:" + token)
        features.add("candidate_text_pair:" + str(safe_candidate.get("candidate_class", "")) + ":" + token)
    for card in model_input.get("safe_memory_cards") or []:
        if isinstance(card, dict):
            features.add("memory_evidence_bucket=" + str(_evidence_bucket(card.get("evidence_count", 0))))
            for token in _visible_tokens([card.get("claim", "")]):
                features.add("memory:" + token)
    for card in model_input.get("safe_skill_cards") or []:
        if isinstance(card, dict):
            features.add("skill_status=" + str(card.get("status", "")))
            features.add("skill_verifier=" + str(card.get("verifier", "")))
    bounded = sorted(features)
    if len(bounded) > config.max_features:
        return sorted(bounded, key=lambda item: stable_hash(item))[: config.max_features]
    return bounded


def _predict_from_weights(
    model_input: dict[str, Any],
    weights: dict[str, float],
    config: Gate8BNeuralControllerConfig,
) -> Gate8BNeuralPrediction:
    _assert_model_input_is_visible(model_input)
    forced_abstain = _forced_abstain_reason(model_input, config)
    choices = _candidate_choices(model_input)
    if forced_abstain:
        return Gate8BNeuralPrediction(
            candidate_id="ABSTAIN",
            score=0.0,
            top_candidate_ids=["ABSTAIN"] + [str(choice.get("candidate_id", "")) for choice in choices[:2]],
            abstained=True,
            reason=forced_abstain,
        )
    scored = _score_candidate_choices(model_input, weights, config)
    if not scored:
        return Gate8BNeuralPrediction("ABSTAIN", 0.0, ["ABSTAIN"], True, "no_safe_candidates")
    top_id, score = scored[0]
    return Gate8BNeuralPrediction(
        candidate_id=top_id,
        score=score,
        top_candidate_ids=[candidate_id for candidate_id, _score in scored[:3]],
        abstained=top_id == "ABSTAIN",
        reason="model_rank",
    )


def _score_candidate_choices(
    model_input: dict[str, Any],
    weights: dict[str, float],
    config: Gate8BNeuralControllerConfig,
) -> list[tuple[str, float]]:
    scored: list[tuple[str, float]] = []
    for candidate in _candidate_choices(model_input):
        candidate_id = str(candidate.get("candidate_id", ""))
        features = extract_neural_candidate_features(model_input, candidate, config)
        score = sum(weights.get(feature, 0.0) for feature in features)
        scored.append((candidate_id, round(score, 6)))
    scored.sort(key=lambda item: (-item[1], item[0]))
    return scored


def _candidate_choices(model_input: dict[str, Any]) -> list[dict[str, Any]]:
    mask = {str(key): bool(value) for key, value in dict(model_input.get("candidate_action_mask") or {}).items()}
    choices: list[dict[str, Any]] = []
    for candidate in model_input.get("candidate_list") or []:
        if not isinstance(candidate, dict):
            continue
        candidate_id = str(candidate.get("candidate_id", ""))
        if candidate_id and bool(mask.get(candidate_id, False)):
            choices.append(dict(candidate))
    choices.append(
        {
            "candidate_id": "ABSTAIN",
            "candidate_class": "ABSTAIN",
            "risk_class": "abstain",
            "authority_level": "abstain",
            "mutation_class": "none",
            "verifier_required": False,
            "rollback_available": False,
            "visible_evidence_count": 0,
            "gate_hint": "gate8b_shadow_abstain",
        }
    )
    return choices


def _choice_for_target(model_input: dict[str, Any], target_id: str) -> dict[str, Any]:
    for choice in _candidate_choices(model_input):
        if str(choice.get("candidate_id", "")) == target_id:
            return choice
    if target_id == "ABSTAIN":
        return _candidate_choices(model_input)[-1]
    raise ValueError(f"target candidate not present or not unmasked: {target_id}")


def _candidate_without_shortcut_fields(candidate: dict[str, Any]) -> dict[str, Any]:
    payload = dict(candidate)
    payload.pop("candidate_id", None)
    payload.pop("position", None)
    return payload


def _assert_visible_training_inputs(transitions: list[ControllerV2Transition]) -> None:
    for transition in transitions:
        _assert_model_input_is_visible(transition.model_input)


def _assert_model_input_is_visible(model_input: dict[str, Any]) -> None:
    forbidden_keys = _find_forbidden_keys(model_input)
    if forbidden_keys:
        raise ValueError("forbidden Gate 8B model input keys: " + ", ".join(forbidden_keys))
    text = canonical_json(model_input)
    forbidden_terms = [term for term in FORBIDDEN_INPUT_TERMS if term in text]
    if forbidden_terms:
        raise ValueError("forbidden Gate 8B model input terms: " + ", ".join(forbidden_terms))


def _find_forbidden_keys(value: Any, prefix: str = "") -> list[str]:
    found: list[str] = []
    if isinstance(value, dict):
        for key, child in value.items():
            key_text = str(key)
            path = f"{prefix}.{key_text}" if prefix else key_text
            if key_text in FORBIDDEN_INPUT_KEYS:
                found.append(path)
            found.extend(_find_forbidden_keys(child, path))
    elif isinstance(value, list):
        for index, child in enumerate(value):
            found.extend(_find_forbidden_keys(child, f"{prefix}[{index}]"))
    return found


def _forced_abstain_reason(model_input: dict[str, Any], config: Gate8BNeuralControllerConfig) -> str:
    text = canonical_json(model_input).lower()
    for marker in config.ood_abstain_markers:
        if marker.lower() in text:
            return "ood_visible_marker"
    for marker in config.rejection_abstain_markers:
        if marker.lower() in text:
            return "deterministic_rejection_marker"
    return ""


def _bounded_weights(weights: dict[str, float], max_features: int) -> dict[str, float]:
    nonzero = {feature: round(value, 6) for feature, value in weights.items() if abs(value) > 1e-9}
    if len(nonzero) <= max_features:
        return dict(sorted(nonzero.items()))
    ordered = sorted(nonzero.items(), key=lambda item: (-abs(item[1]), item[0]))[:max_features]
    return dict(sorted(ordered))


def _gate8c_calibrated_confidence(calibration: Gate8CNeuralCalibration, margin: float) -> float:
    if not calibration.bucket_confidences:
        return 0.0
    index = _margin_bucket_index(margin, calibration.bucket_edges)
    index = min(index, len(calibration.bucket_confidences) - 1)
    return float(calibration.bucket_confidences[index])


def _gate8c_ood_score(
    model: Gate8BNeuralControllerModel,
    calibration: Gate8CNeuralCalibration,
    model_input: dict[str, Any],
    candidate_id: str,
) -> float:
    if candidate_id == "ABSTAIN":
        return 1.0
    try:
        candidate = _choice_for_target(model_input, candidate_id)
        features = set(extract_neural_candidate_features(model_input, candidate, model.config))
    except ValueError:
        return 1.0
    if not features:
        return 1.0
    vocabulary = set(calibration.feature_vocabulary)
    if not vocabulary:
        return 1.0
    unseen = sum(1 for feature in features if feature not in vocabulary)
    return unseen / len(features)


def _gate8c_feature_vocabulary(
    model: Gate8BNeuralControllerModel,
    transitions: list[ControllerV2Transition],
) -> list[str]:
    vocabulary = set(model.feature_vocabulary)
    for transition in transitions:
        if transition_split(transition) not in {"train", "validation"}:
            continue
        target = target_candidate_id(transition)
        try:
            candidate = _choice_for_target(transition.model_input, target)
        except ValueError:
            continue
        vocabulary.update(extract_neural_candidate_features(transition.model_input, candidate, model.config))
    return sorted(vocabulary)


def _merge_reason_counts(first: dict[str, int], second: dict[str, int]) -> dict[str, int]:
    merged = {str(key): int(value) for key, value in first.items()}
    for key, value in second.items():
        text_key = str(key)
        merged[text_key] = merged.get(text_key, 0) + int(value)
    return dict(sorted(merged.items()))


def _margin_bucket_index(margin: float, bucket_edges: list[float]) -> int:
    if not bucket_edges:
        return 0
    index = 0
    for next_index, edge in enumerate(bucket_edges):
        if margin >= edge:
            index = next_index
    return index


def _evidence_bucket(value: Any) -> int:
    try:
        number = int(value)
    except (TypeError, ValueError):
        return 0
    if number <= 0:
        return 0
    if number <= 2:
        return 1
    if number <= 4:
        return 2
    return 3


def _visible_tokens(values: list[Any]) -> list[str]:
    text = " ".join(str(value).lower() for value in values if value is not None)
    return [token for token in re.findall(r"[a-z0-9_]+", text) if len(token) >= 3][:80]


def _candidate_verifier_matches_context(candidate: dict[str, Any], model_input: dict[str, Any]) -> bool:
    context = str(model_input.get("verifier_context_summary", "")).lower()
    required = bool(candidate.get("verifier_required", False))
    if "not_required" in context:
        return not required
    if "passed" in context or "failed" in context:
        return required
    return False


def _candidate_rollback_matches_context(candidate: dict[str, Any], model_input: dict[str, Any]) -> bool:
    context = str(model_input.get("rollback_context_summary", "")).lower()
    available = bool(candidate.get("rollback_available", False))
    if "not_applicable" in context:
        return not available
    if "available" in context or "succeeded" in context or "rollback" in context:
        return available
    return False


def _majority_target_class(transitions: list[ControllerV2Transition]) -> str:
    counts: dict[str, int] = {}
    for transition in transitions:
        target_id = target_candidate_id(transition)
        if target_id == "ABSTAIN":
            key = "ABSTAIN"
        else:
            choice = _choice_for_target(transition.model_input, target_id)
            key = str(choice.get("candidate_class", ""))
        counts[key] = counts.get(key, 0) + 1
    if not counts:
        return "ABSTAIN"
    return sorted(counts.items(), key=lambda item: (-item[1], item[0]))[0][0]


def _first_candidate_for_class(model_input: dict[str, Any], candidate_class: str) -> str:
    if candidate_class == "ABSTAIN":
        return "ABSTAIN"
    for candidate in _candidate_choices(model_input):
        if str(candidate.get("candidate_class", "")) == candidate_class:
            return str(candidate.get("candidate_id", ""))
    return "ABSTAIN"


def _candidate_at_position(model_input: dict[str, Any], position: str) -> str:
    for candidate in model_input.get("candidate_list") or []:
        if isinstance(candidate, dict) and str(candidate.get("position", "")) == position:
            return str(candidate.get("candidate_id", ""))
    return "ABSTAIN"
