"""
llm_guard/mechanism_layer.py — Mechanism-layer feature extractor for universal LLM monitoring.
==============================================================================================
Computes 12 text-level signals from (question, answer, context) triples.
No label parameter. No dataset-specific code. No external corpus dependencies.

Phase A signals (12): text-only, zero API cost.
Phase B signals (3): require LLM traces (f13, f14, f15) — see extract_phase_b_features().

Usage:
    extractor = MechanismFeatureExtractor()
    features = extractor.extract_features(question, answer, context="")  # -> np.ndarray shape (12,)
    batch    = extractor.extract_features_batch(records)                  # -> np.ndarray shape (N, 12)
Purpose:    Computes 12 text-level signals from (question, answer, context) triples with zero API cost. Phase A achieves AUROC ≥ 0.60 gate on 4/5 datasets; Phase B adds 3 SC signals.
Exports:    MechanismFeatureExtractor, MECHANISM_FEATURE_NAMES, bootstrap_ci, run_cv_experiment
Added:      v0.46.0
Status:     Validated
"""

from __future__ import annotations

import json
import re
import math
from pathlib import Path
from typing import List, Optional
import numpy as np

# ── Word lists (dependency-free polarity) ─────────────────────────────────────

CERTAINTY_WORDS = frozenset({
    "definitely", "certainly", "always", "clearly", "obviously",
    "undoubtedly", "absolutely", "unquestionably", "surely", "necessarily",
})

HEDGE_WORDS = frozenset({
    "might", "perhaps", "could", "unclear", "possibly", "probably",
    "seemingly", "approximately", "around", "roughly", "maybe", "sometimes",
    "generally", "often", "typically",
})

POSITIVE_WORDS = frozenset({
    "good", "correct", "true", "right", "yes", "accurate", "valid", "agree",
    "confirm", "support", "positive", "beneficial", "helpful", "effective",
})

NEGATIVE_WORDS = frozenset({
    "wrong", "false", "no", "incorrect", "invalid", "disagree", "bad",
    "harmful", "negative", "misleading", "inaccurate", "dangerous",
})

FEATURE_NAMES = [
    "sem_sim_qa",         # f1
    "answer_len_norm",    # f2
    "certainty_density",  # f3
    "hedge_density",      # f4
    "lexical_overlap",    # f5
    "answer_repetition",  # f6
    "question_type",      # f7
    "polarity_shift",     # f8
    "type_token_ratio",   # f9
    "numeric_density",    # f10
    "entity_proxy",       # f11
    "context_answer_sim", # f12
]

PHASE_B_FEATURE_NAMES = [
    "self_consistency_score",  # f13
    "response_entropy_proxy",  # f14
    "answer_stability",        # f15
]


# ── Helper functions ──────────────────────────────────────────────────────────

def tokenize(text: str) -> list[str]:
    """Public — used by Phase B script and tests."""
    return re.findall(r'\b\w+\b', text.lower())

# Private aliases kept for internal use only
_tokenize = tokenize


def _bigrams(tokens: list[str]) -> set[tuple]:
    return set(zip(tokens[:-1], tokens[1:])) if len(tokens) > 1 else set()


def _trigrams(tokens: list[str]) -> list[tuple]:
    return list(zip(tokens[:-2], tokens[1:-1], tokens[2:])) if len(tokens) > 2 else []


def _polarity(text: str) -> float:
    tokens = _tokenize(text)
    if not tokens:
        return 0.0
    pos = sum(1 for t in tokens if t in POSITIVE_WORDS)
    neg = sum(1 for t in tokens if t in NEGATIVE_WORDS)
    return (pos - neg) / len(tokens)


def _question_type(question: str) -> float:
    """0.0=factoid (who/what/when), 1.0=yes-no (is/are/does), 2.0=open."""
    q = question.lower().strip()
    if re.match(r'^(is|are|does|do|was|were|has|have|can|could|will|would|should)\b', q):
        return 1.0
    if re.match(r'^(who|what|when|where|which|how many|how much)\b', q):
        return 0.0
    return 2.0


def vocab_entropy(tokens: list[str]) -> float:
    """Shannon entropy of token unigram distribution. Public — used by Phase B."""
    if not tokens:
        return 0.0
    from collections import Counter
    counts = Counter(tokens)
    n = len(tokens)
    return -sum((c/n) * math.log2(c/n) for c in counts.values())

# Private alias for internal use
_vocab_entropy = vocab_entropy


# ── Feature Extractor ─────────────────────────────────────────────────────────

class MechanismFeatureExtractor:
    """
    Extracts 12 text-level mechanism signals from (question, answer, context) triples.
    No label parameter. No dataset ID parameter. No external corpus dependencies.

    Args:
        model_name: sentence-transformers model name (default: all-MiniLM-L6-v2)
        cache_dir:  optional directory to cache embeddings to disk
        encoder:    "sentence_transformers" (default) or "deberta" (uses
                    microsoft/deberta-v3-small; requires transformers + torch)
    """

    def __init__(
        self,
        model_name: str = "all-MiniLM-L6-v2",
        cache_dir: Optional[Path] = None,
        encoder: str = "sentence_transformers",
    ):
        if encoder not in ("sentence_transformers", "deberta"):
            raise ValueError(f"encoder must be 'sentence_transformers' or 'deberta', got {encoder!r}")
        self.model_name    = model_name
        self.encoder       = encoder
        self.cache_dir     = Path(cache_dir) if cache_dir else None
        self._model        = None  # lazy load
        self._tokenizer    = None  # only used for deberta
        self._domain_clf   = None  # set by fit_domain_calibration()
        self._domain_clf_n = 0

    def _load_model(self):
        if self._model is None:
            if self.encoder == "deberta":
                from transformers import AutoTokenizer, AutoModel
                _dbt_name = "microsoft/deberta-v3-small"
                self._tokenizer = AutoTokenizer.from_pretrained(_dbt_name)
                self._model = AutoModel.from_pretrained(_dbt_name)
                self._model.eval()
            else:
                from sentence_transformers import SentenceTransformer
                self._model = SentenceTransformer(self.model_name)
        return self._model

    def _embed(self, texts: list[str]) -> np.ndarray:
        """Encode texts to L2-normalised embeddings. Returns (N, D) array."""
        model = self._load_model()
        if self.encoder == "deberta":
            return self._embed_deberta(texts)
        return model.encode(texts, normalize_embeddings=True, show_progress_bar=False)

    def _embed_deberta(self, texts: list[str]) -> np.ndarray:
        """DeBERTa-v3-small mean-pool encoder. Returns L2-normalised (N, D) array."""
        import torch
        embs = []
        for text in texts:
            enc = self._tokenizer(
                text, return_tensors="pt", truncation=True, max_length=512, padding=True
            )
            with torch.no_grad():
                out = self._model(**enc)
            # Mean pool over non-padding tokens
            mask = enc["attention_mask"].unsqueeze(-1).float()
            vec = (out.last_hidden_state * mask).sum(1) / mask.sum(1)
            vec = vec.squeeze(0).numpy()
            norm = np.linalg.norm(vec)
            embs.append(vec / max(norm, 1e-9))
        return np.array(embs, dtype=np.float32)

    def _cosine(self, a: np.ndarray, b: np.ndarray) -> float:
        """Cosine similarity of two normalized vectors."""
        if a.ndim == 1 and b.ndim == 1:
            dot = float(np.dot(a, b))
            return float(np.clip(dot, -1.0, 1.0))
        raise ValueError(f"Expected 1D vectors, got {a.shape}, {b.shape}")

    def extract_features(
        self,
        question: str,
        answer: str,
        context: str = "",
    ) -> np.ndarray:
        """
        Extract 12 mechanism signals from a (question, answer, context) triple.

        Args:
            question: the question or prompt (str)
            answer:   the model's response to evaluate (str)
            context:  supporting context/passage (str, optional, default "")

        Returns:
            np.ndarray of shape (12,) — one value per feature
        """
        assert isinstance(question, str), "question must be str"
        assert isinstance(answer, str),   "answer must be str"
        assert isinstance(context, str),  "context must be str"

        q_tokens = _tokenize(question)
        a_tokens = _tokenize(answer)
        n_a = max(1, len(a_tokens))

        # Embed in one call for efficiency
        texts_to_embed = [question, answer]
        if context:
            texts_to_embed.append(context)
        embs = self._embed(texts_to_embed)
        q_emb = embs[0]
        a_emb = embs[1]
        c_emb = embs[2] if context else None

        # f1: sem_sim_qa
        f1 = self._cosine(q_emb, a_emb)

        # f2: answer_len_norm
        f2 = math.log1p(len(a_tokens)) / max(1.0, math.log1p(len(q_tokens)))

        # f3: certainty_density
        f3 = sum(1 for t in a_tokens if t in CERTAINTY_WORDS) / n_a

        # f4: hedge_density
        f4 = sum(1 for t in a_tokens if t in HEDGE_WORDS) / n_a

        # f5: lexical_overlap (bigram)
        q_bigrams = _bigrams(q_tokens)
        a_bigrams = _bigrams(a_tokens)
        f5 = len(q_bigrams & a_bigrams) / max(1, len(q_bigrams))

        # f6: answer_repetition (1 - unique_trigrams/all_trigrams)
        a_trig = _trigrams(a_tokens)
        f6 = 1.0 - (len(set(a_trig)) / max(1, len(a_trig)))

        # f7: question_type_proxy (0=factoid, 1=yes-no, 2=open)
        f7 = _question_type(question)

        # f8: polarity_shift
        f8 = abs(_polarity(question) - _polarity(answer))

        # f9: type_token_ratio
        f9 = len(set(a_tokens)) / n_a

        # f10: numeric_density
        nums = re.findall(r'\b\d+\.?\d*\b', answer)
        f10 = len(nums) / n_a

        # f11: entity_proxy_density (non-start capitalized tokens)
        a_raw_tokens = answer.split()
        entity_count = sum(
            1 for i, t in enumerate(a_raw_tokens)
            if i > 0 and t and t[0].isupper() and t.isalpha()
        )
        f11 = entity_count / max(1, len(a_raw_tokens))

        # f12: context_answer_sim (0.0 when context absent)
        f12 = self._cosine(c_emb, a_emb) if c_emb is not None else 0.0

        result = np.array([f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12],
                          dtype=np.float64)
        if result.shape != (12,):
            raise ValueError(f"Feature vector shape mismatch: expected (12,), got {result.shape}")
        return result

    def fit_domain_calibration(
        self,
        pairs: list[dict],
        labels: list[int],
    ) -> "MechanismFeatureExtractor":
        """
        Fit a domain-specific decision boundary on top of the frozen feature extractor.

        The 12 text-level features are extracted as-is (no retraining of the encoder).
        A logistic regression is trained on the extracted features and the provided labels.
        After calling this, ``predict_domain()`` will use the domain-specific classifier
        instead of the zero-shot sign-of-f1 heuristic.

        Parameters
        ----------
        pairs : list[dict]
            Each dict must have 'question' and 'answer' keys; 'context' is optional.
            Minimum 10 samples required; ≥50 recommended for reliable calibration.
        labels : list[int]
            Binary labels: 1 = failure (hallucination / wrong), 0 = correct.
            Must be the same length as ``pairs``.

        Returns
        -------
        self — for chaining.

        Raises
        ------
        ValueError
            If ``pairs`` and ``labels`` differ in length or fewer than 10 samples.
        """
        from sklearn.linear_model import LogisticRegression as _LR

        if len(pairs) != len(labels):
            raise ValueError(
                f"pairs ({len(pairs)}) and labels ({len(labels)}) must have the same length."
            )
        if len(pairs) < 10:
            raise ValueError(
                f"At least 10 labeled pairs required for domain calibration; got {len(pairs)}."
            )

        X = self.extract_features_batch(pairs)
        y = np.asarray(labels, dtype=int)

        self._domain_clf = _LR(
            C=1.0,
            max_iter=1000,
            class_weight="balanced",
            solver="lbfgs",
            random_state=42,
        )
        self._domain_clf.fit(X, y)
        self._domain_clf_n = len(pairs)
        return self

    def predict_domain(
        self,
        question: str,
        answer: str,
        context: str = "",
    ) -> float:
        """
        Return a failure probability using the domain-calibrated classifier.

        Falls back to zero-shot heuristic (negated f1 similarity) if
        ``fit_domain_calibration()`` has not been called.

        Parameters
        ----------
        question : str
        answer   : str
        context  : str, optional

        Returns
        -------
        float in [0, 1] — probability that the answer is a failure.
        """
        # extract_features returns (12,) array: [f1, f2, ..., f12]
        feats = self.extract_features(question, answer, context)  # shape (12,)

        if not hasattr(self, "_domain_clf") or self._domain_clf is None:
            # Zero-shot fallback: low sem_sim_qa (f1) = likely failure
            f1 = float(feats[0])
            return float(np.clip(1.0 - f1, 0.0, 1.0))

        x = feats.reshape(1, -1)  # (1, 12)
        return float(self._domain_clf.predict_proba(x)[0, 1])

    def extract_features_batch(self, records: list[dict]) -> np.ndarray:
        """
        Truly batched feature extraction — encodes all texts in one sentence-transformers
        call for efficiency (3000 sequential encode() calls would be ~10x slower).
        Each record must have 'question', 'answer', and optionally 'context'.
        Returns (N, 12) array.
        """
        n = len(records)
        questions = [r["question"] for r in records]
        answers   = [r["answer"]   for r in records]
        contexts  = [r.get("context", "") for r in records]

        # Single batch embed — much faster than per-record calls
        all_texts   = questions + answers + [c for c in contexts if c]
        all_embs    = self._embed(all_texts)
        q_embs      = all_embs[:n]
        a_embs      = all_embs[n:2*n]
        ctx_offset  = 2 * n
        ctx_emb_map: dict[int, np.ndarray] = {}
        ctx_idx = 0
        for i, c in enumerate(contexts):
            if c:
                ctx_emb_map[i] = all_embs[ctx_offset + ctx_idx]
                ctx_idx += 1

        out = np.zeros((n, 12), dtype=np.float64)
        for i, rec in enumerate(records):
            q = rec["question"];  a = rec["answer"];  c = rec.get("context", "")
            q_tokens = _tokenize(q);  a_tokens = _tokenize(a);  n_a = max(1, len(a_tokens))
            q_emb = q_embs[i];  a_emb = a_embs[i]
            c_emb = ctx_emb_map.get(i)

            f1  = self._cosine(q_emb, a_emb)
            f2  = math.log1p(len(a_tokens)) / max(1.0, math.log1p(len(q_tokens)))
            f3  = sum(1 for t in a_tokens if t in CERTAINTY_WORDS) / n_a
            f4  = sum(1 for t in a_tokens if t in HEDGE_WORDS) / n_a
            q_bg = _bigrams(q_tokens);  a_bg = _bigrams(a_tokens)
            f5  = len(q_bg & a_bg) / max(1, len(q_bg))
            a_tg = _trigrams(a_tokens)
            f6  = 1.0 - (len(set(a_tg)) / max(1, len(a_tg)))
            f7  = _question_type(q)
            f8  = abs(_polarity(q) - _polarity(a))
            f9  = len(set(a_tokens)) / n_a
            f10 = len(re.findall(r'\b\d+\.?\d*\b', a)) / n_a
            a_raw = a.split()
            f11 = sum(1 for j, t in enumerate(a_raw)
                      if j > 0 and t and t[0].isupper() and t.isalpha()) / max(1, len(a_raw))
            f12 = self._cosine(c_emb, a_emb) if c_emb is not None else 0.0
            out[i] = [f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12]
        return out


# ══════════════════════════════════════════════════════════════════════════════
# Dataset Loaders — Phase A
# Each returns list of dicts with: question, answer, label (0/1), context
# Labels: 1 = failure (hallucinated / wrong / sycophantic)
# ══════════════════════════════════════════════════════════════════════════════

def _read_jsonl(path: Path) -> list[dict]:
    return [json.loads(l) for l in Path(path).read_text().splitlines() if l.strip()]


def load_halueval_qa(path: Path) -> list[dict]:
    """
    HaluEval 2.0 QA split. Schema: knowledge, question, answer, hallucination.
    Each record is independently labeled — NOT paired.
    label=1 if hallucination=='yes'.
    """
    records = []
    for raw in _read_jsonl(path):
        records.append({
            "question": raw["question"],
            "answer":   raw["answer"],
            "label":    1 if raw["hallucination"] == "yes" else 0,
            "context":  raw.get("knowledge", "")[:500],
            "dataset":  "halueval2_qa",
        })
    return records


def load_halueval_summarization(path: Path) -> list[dict]:
    """
    HaluEval 2.0 summarization split. Schema: document, summary, hallucination.
    No question field — first sentence of document is used as proxy question.
    Primary discriminative signal: f12 (context_answer_sim), not f1.
    """
    records = []
    for raw in _read_jsonl(path):
        doc = raw.get("document", "")
        first_sent = re.split(r'(?<=[.!?])\s+', doc.strip())[0] if doc else ""
        records.append({
            "question": first_sent[:200],
            "answer":   raw.get("summary", ""),
            "label":    1 if raw["hallucination"] == "yes" else 0,
            "context":  doc[:500],
            "dataset":  "halueval2_summary",
        })
    return records


def load_halueval_dialogue(path: Path) -> list[dict]:
    """
    HaluEval 2.0 dialogue split. Schema: knowledge, dialogue_history, response, hallucination.
    Last [Human]: utterance is extracted as question.
    """
    records = []
    for raw in _read_jsonl(path):
        history = raw.get("dialogue_history", "")
        human_turns = re.findall(r'\[Human\]:\s*(.+?)(?=\[|$)', history, re.DOTALL)
        q = human_turns[-1].strip() if human_turns else history[:100]
        records.append({
            "question": q[:200],
            "answer":   raw.get("response", ""),
            "label":    1 if raw["hallucination"] == "yes" else 0,
            "context":  raw.get("knowledge", "")[:500],
            "dataset":  "halueval2_dialogue",
        })
    return records


def load_truthfulqa(path: Path) -> list[dict]:
    """
    TruthfulQA MC1 split. Each question expands to N choice records.
    mc1_targets.labels: 1=correct, 0=wrong.
    We invert: label=1 means FAILURE (wrong answer).
    ~20% positive class (1 correct per question, N-1 wrong).
    """
    records = []
    for raw in _read_jsonl(path):
        q = raw["question"]
        mc = raw.get("mc1_targets", {})
        choices = mc.get("choices", [])
        labels  = mc.get("labels", [])
        for choice, correct in zip(choices, labels):
            records.append({
                "question": q,
                "answer":   choice,
                "label":    1 - correct,   # invert: 1=wrong=failure
                "context":  "",
                "dataset":  "truthfulqa",
            })
    return records


def load_sycophancy(path: Path, seed: int = 42) -> list[dict]:
    """
    Anthropic sycophancy eval. Expands each record into 2 records:
      - sycophantic answer (label=1, failure)
      - non-sycophantic answer (label=0, correct)
    answer_matching_behavior is a raw letter (A) — full text extracted from question field.
    question_stem stripped of choices block to prevent f1/f5 contamination.
    """
    rng = np.random.default_rng(seed)
    records = []
    for raw in _read_jsonl(path):
        q_full = raw.get("question", "")
        amb    = raw.get("answer_matching_behavior", "").strip().strip("()")
        anmb   = raw.get("answer_not_matching_behavior", "").strip().strip("()")

        # Strip choices block from question to get clean stem
        q_stem = re.split(
            r'\n\s*(?:choices\s*:\s*|\(A\)\s)',
            q_full,
            flags=re.IGNORECASE,
            maxsplit=1,
        )[0].strip()

        # Parse full choice texts
        choice_pairs = re.findall(r'\(([A-Z])\)\s+([^\n(]+)', q_full)
        choice_dict  = {letter: text.strip() for letter, text in choice_pairs}

        pos_text = choice_dict.get(amb, "")
        neg_text = choice_dict.get(anmb, "")

        # Drop unparseable records
        if not pos_text or not neg_text or not q_stem:
            continue

        # Position shuffle per record
        if rng.integers(2):
            pair = [
                {"question": q_stem, "answer": pos_text, "label": 1, "context": "", "dataset": "sycophancy"},
                {"question": q_stem, "answer": neg_text, "label": 0, "context": "", "dataset": "sycophancy"},
            ]
        else:
            pair = [
                {"question": q_stem, "answer": neg_text, "label": 0, "context": "", "dataset": "sycophancy"},
                {"question": q_stem, "answer": pos_text, "label": 1, "context": "", "dataset": "sycophancy"},
            ]
        records.extend(pair)

    return records


def load_pubmed_qa(path: Path) -> list[dict]:
    """
    PubMedQA biomedical yes/no/maybe QA.
    Uses first 500 chars of joined context paragraphs as the answer proxy.
    label=1 (failure) when answer is 'no' — hypothesis not supported.
    label=0 (correct) when answer is 'yes' — hypothesis supported.
    Records with 'maybe' are excluded (ambiguous).

    Pre-registered caveat: the 'answer' is context text, not a model-generated
    response. f12 (context_answer_sim) will be near-1.0 for all records (context
    IS the answer). Primary signal is f1 (sem_sim_qa) discriminating yes vs no.
    """
    records = []
    for raw in _read_jsonl(path):
        lbl_str = str(raw.get("label", "")).lower()
        if lbl_str not in ("yes", "no"):
            continue  # skip 'maybe' and missing
        ctx = raw.get("context", "")
        if isinstance(ctx, list):
            ctx = " ".join(str(c) for c in ctx)
        records.append({
            "question": raw["question"],
            "answer":   ctx[:500],
            "label":    1 if lbl_str == "no" else 0,   # no=failure, yes=correct
            "context":  "",
            "dataset":  "pubmed_qa",
        })
    return records


def load_quality_context_drift(path: Path) -> list[dict]:
    """
    QuALITY long-document QA proxy for context drift.
    Each record has: question, options (list[str]), gold_label (int index).
    Expands to N records (one per option).
    label=1 (failure) if option is wrong (index != gold_label).
    label=0 (correct) if option is right (index == gold_label).
    article is truncated to 500 chars as context.

    Follows TruthfulQA MC expansion convention.
    """
    records = []
    for raw in _read_jsonl(path):
        q       = raw.get("question", "")
        options = raw.get("options", [])
        gold    = raw.get("gold_label", -1)
        article = raw.get("article", "")[:500]
        if not options or not q:
            continue
        for idx, opt_text in enumerate(options):
            records.append({
                "question": q,
                "answer":   str(opt_text),
                "label":    0 if idx == gold else 1,  # wrong=failure
                "context":  article,
                "dataset":  "context_drift",
            })
    return records


def load_arc_research(path: Path) -> list[dict]:
    """
    ARC-Challenge dataset (easy+hard split). MC expansion.
    choices is a dict with 'text' (list) and 'label' (list of letters A/B/C/D).
    answerKey is the correct letter.
    label=1 (failure) for wrong choices, label=0 (correct) for the right choice.

    Follows TruthfulQA MC expansion convention.
    """
    records = []
    for raw in _read_jsonl(path):
        q          = raw.get("question", "")
        choices    = raw.get("choices", {})
        answer_key = raw.get("answerKey", "")
        texts      = choices.get("text", [])
        labels     = choices.get("label", [])
        if not texts or not answer_key or not q:
            continue
        for text, letter in zip(texts, labels):
            records.append({
                "question": q,
                "answer":   str(text),
                "label":    0 if letter == answer_key else 1,  # wrong=failure
                "context":  "",
                "dataset":  "arc_research",
            })
    return records


# ══════════════════════════════════════════════════════════════════════════════
# Ensemble with SC_OLD joint features
# ══════════════════════════════════════════════════════════════════════════════

class EnsembleMechanismClassifier:
    """
    Joins mechanism-layer features (12+) with optional SC_OLD behavioral features
    and runs the same stratified CV / 80-20 pipeline.

    Implements Image 2 Layer 2: mechanism layer + SC_OLD scoring + calibration
    signal → weighted ensemble aggregator.

    Args:
        use_scold:          whether to include SC_OLD features (default: True when provided)
        scold_feature_names: names of SC_OLD features (for importance reporting)

    Usage:
        clf = EnsembleMechanismClassifier()
        result = clf.run(X_mechanism, y, X_scold=X_scold)
        # result same shape as run_cv_experiment output, with
        # result['feature_names'] listing all features used
    """

    def __init__(
        self,
        scold_feature_names: Optional[List[str]] = None,
    ):
        self.scold_feature_names = scold_feature_names or []

    def run(
        self,
        X_mechanism: np.ndarray,
        y: np.ndarray,
        X_scold: Optional[np.ndarray] = None,
        seed: int = 42,
    ) -> dict:
        """
        Run CV ensemble on mechanism features optionally joined with SC_OLD features.

        Args:
            X_mechanism: (N, 12+) mechanism layer features
            y:           (N,) binary labels
            X_scold:     (N, K) SC_OLD behavioral features, or None
            seed:        random seed

        Returns:
            run_cv_experiment output dict with extended feature_names
        """
        if X_scold is not None and len(X_scold) == len(X_mechanism):
            X = np.hstack([X_mechanism, X_scold])
            feat_names = list(FEATURE_NAMES) + list(self.scold_feature_names)
        else:
            X = X_mechanism
            feat_names = list(FEATURE_NAMES)

        result = run_cv_experiment(X, y, seed=seed)
        result["feature_names"] = feat_names
        result["n_mechanism_features"] = X_mechanism.shape[1]
        result["n_scold_features"] = X_scold.shape[1] if X_scold is not None else 0
        return result


# ──────────────────────────────────────────────────────────────────────────────
# SemanticConsistencyChecker — Wave 1 / 29.3 Phase B SC Ablation
# ──────────────────────────────────────────────────────────────────────────────

class SemanticConsistencyChecker:
    """
    Computes semantic self-consistency score from multiple LLM response samples.

    Unlike the pairwise cosine similarity in the existing Phase B f13 feature
    (which compares full response embeddings), this class focuses on the
    extracted answer portion and uses sentence-embedding cosine similarity
    between all response pairs.

    Parameters
    ----------
    sc_temperature : float
        Temperature to use when sampling responses from the LLM API.
        0.0 = deterministic (tests sampling-artifact hypothesis).
        1.0 = stochastic (current Phase B default).
        Default: 1.0.
    model : str
        Haiku model identifier for API sampling.
        Default: "claude-haiku-4-5-20251001".

    Usage
    -----
        checker = SemanticConsistencyChecker(sc_temperature=0.0)
        responses = checker.sample_responses(question, context, api_key, n=3)
        score = checker.compute_semantic_consistency(extractor, responses)
    """

    def __init__(
        self,
        sc_temperature: float = 1.0,
        model: str = "claude-haiku-4-5-20251001",
    ) -> None:
        if not (0.0 <= sc_temperature <= 2.0):
            raise ValueError(f"sc_temperature must be in [0, 2], got {sc_temperature}")
        self.sc_temperature = sc_temperature
        self.model = model

    def compute_semantic_consistency(
        self,
        extractor: "MechanismFeatureExtractor",
        responses: List[str],
    ) -> float:
        """
        Compute mean pairwise cosine similarity between response embeddings.

        Parameters
        ----------
        extractor : MechanismFeatureExtractor
            Used to embed responses via extractor._embed().
        responses : List[str]
            List of N response strings. N=1 returns 1.0 (trivially consistent).

        Returns
        -------
        float in [0, 1] — 1.0 = perfectly consistent, 0.0 = maximally inconsistent.
        """
        if len(responses) < 2:
            return 1.0
        embs = extractor._embed(responses)
        sims = []
        for i in range(len(embs)):
            for j in range(i + 1, len(embs)):
                dot = float(np.dot(embs[i], embs[j]))
                sims.append(float(np.clip(dot, -1.0, 1.0)))
        return float(np.clip(np.mean(sims), 0.0, 1.0))

    def sample_responses(
        self,
        question: str,
        context: str,
        api_key: str,
        n: int = 3,
    ) -> List[str]:
        """
        Sample N responses from Claude Haiku at sc_temperature.

        Parameters
        ----------
        question : str   The question to answer.
        context  : str   Supporting context.
        api_key  : str   Anthropic API key.
        n        : int   Number of responses to sample. Default 3.

        Returns
        -------
        List[str] of N response strings.
        """
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)
        prompt = f"Answer the following question based on the context.\nContext: {context}\nQuestion: {question}\nAnswer concisely."
        responses = []
        for _ in range(n):
            msg = client.messages.create(
                model=self.model,
                max_tokens=64,
                temperature=self.sc_temperature,
                messages=[{"role": "user", "content": prompt}],
            )
            responses.append(msg.content[0].text.strip())
        return responses


# ══════════════════════════════════════════════════════════════════════════════
# Validation Utilities
# ══════════════════════════════════════════════════════════════════════════════

from sklearn.metrics import roc_auc_score, average_precision_score
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, train_test_split


def bootstrap_ci(
    labels: np.ndarray,
    scores: np.ndarray,
    n_boot: int = 1000,
    seed: int = 42,
) -> dict:
    """
    Compute AUROC + AUPRC with 95% percentile bootstrap CI.
    Consistent with all existing project experiments (np.percentile method, NOT BCa).

    Args:
        labels: binary labels (0/1), shape (N,)
        scores: predicted probabilities, shape (N,)
        n_boot: number of bootstrap iterations
        seed:   random seed for reproducibility

    Returns:
        dict with keys: auroc, auroc_lo, auroc_hi, auprc, auprc_lo, auprc_hi, class_prior
    """
    labels = np.asarray(labels)
    scores = np.asarray(scores)
    rng = np.random.RandomState(seed)

    boot_aurocs = []
    boot_auprcs = []
    for _ in range(n_boot):
        idx = rng.choice(len(labels), len(labels), replace=True)
        l_b, s_b = labels[idx], scores[idx]
        if len(np.unique(l_b)) < 2:
            continue
        boot_aurocs.append(roc_auc_score(l_b, s_b))
        boot_auprcs.append(average_precision_score(l_b, s_b))

    auroc = float(roc_auc_score(labels, scores))
    auprc = float(average_precision_score(labels, scores))

    return {
        "auroc":       auroc,
        "auroc_lo":    float(np.percentile(boot_aurocs, 2.5)),
        "auroc_hi":    float(np.percentile(boot_aurocs, 97.5)),
        "auprc":       auprc,
        "auprc_lo":    float(np.percentile(boot_auprcs, 2.5)),
        "auprc_hi":    float(np.percentile(boot_auprcs, 97.5)),
        "class_prior": float(labels.mean()),
    }


def run_cv_experiment(
    X: np.ndarray,
    y: np.ndarray,
    seed: int = 42,
    cv_threshold: int = 800,
) -> dict:
    """
    Run stratified CV or 80/20 split depending on N.
    Returns OOF scores + metrics + leakage-checked gate result.

    Gate criteria (must ALL pass):
      - AUROC >= 0.60
      - auroc_lo > 0.50
      - AUPRC > class_prior

    Leakage assertions:
    - Every sample scored exactly once (no gaps, no duplicates)
    - Labels only visible inside training fold
    """
    n = len(y)
    clf = LogisticRegression(
        C=1.0, max_iter=1000, class_weight="balanced",
        solver="lbfgs", random_state=seed,
    )

    if n >= cv_threshold:
        # 80/20 stratified split
        X_tr, X_te, y_tr, y_te = train_test_split(
            X, y, test_size=0.2, stratify=y, random_state=seed
        )
        clf.fit(X_tr, y_tr)
        oof_scores = clf.predict_proba(X_te)[:, 1]
        oof_labels = y_te
        n_train, n_test = len(y_tr), len(y_te)
        coef = clf.coef_[0]
        # Leakage assertion for 80/20 path
        if len(oof_scores) != len(y_te):
            raise ValueError(f"OOF score count mismatch: got {len(oof_scores)}, expected {len(y_te)}")
        if len(oof_labels) != len(y_te):
            raise ValueError(f"OOF label count mismatch: got {len(oof_labels)}, expected {len(y_te)}")
    else:
        # 5-fold stratified CV
        cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=seed)
        oof_scores  = np.zeros(n)
        oof_indices = []
        coefs       = []
        for train_idx, test_idx in cv.split(X, y):
            clf.fit(X[train_idx], y[train_idx])
            oof_scores[test_idx] = clf.predict_proba(X[test_idx])[:, 1]
            oof_indices.extend(test_idx.tolist())
            coefs.append(clf.coef_[0].copy())

        # Leakage assertions
        if len(oof_scores) != n:
            raise ValueError(f"OOF score count mismatch: got {len(oof_scores)}, expected {n}")
        if sorted(oof_indices) != list(range(n)):
            raise ValueError("OOF samples not scored exactly once — each sample must appear exactly once")

        oof_labels = y
        n_train = int(n * 0.8)
        n_test  = n - n_train
        coef = np.mean(coefs, axis=0)

    metrics = bootstrap_ci(oof_labels, oof_scores, seed=seed)

    # Gate criteria (spec Section 5.5)
    gate_pass = (
        metrics["auroc"] >= 0.60 and
        metrics["auroc_lo"] > 0.50 and
        metrics["auprc"] > metrics["class_prior"]
    )

    return {
        **metrics,
        "oof_scores":  oof_scores,
        "oof_labels":  oof_labels,
        "n_train":     n_train,
        "n_test":      n_test,
        "gate_pass":   gate_pass,
        "coef":        coef,
        "feature_names": FEATURE_NAMES,
    }
