"""
LabelFreeScorer — zero-label risk scoring for completed ReAct chains.

Based on exp39-44 (intra-chain consistency, multi-step embeddings, QARA on
observations, retrieval relevance, calibration efficiency) plus exp60 (max-pool
embedding aggregation, +0.21 AUROC), exp67 (streaming/prefix scoring),
exp70 (failure mode prediction), and exp75 (new behavioral signals):

  Behavioral signals — no calibration data required:
    SC1  Completion flag    AUROC ~0.84 within-domain
    SC2  Step count         AUROC ~0.88 within-domain  ← single strongest signal
    SC3  Thought variance   AUROC ~0.76 within-domain
    SC5  Uncertainty words  AUROC ~0.63
    SC6  Answer gap         AUROC ~0.80 within-domain
    SC8  Backtracking rate  AUROC ~0.70 within, ~1.0 cross (exp75)
    SC9  Obs utilization    AUROC ~0.77 within, ~0.99 cross (exp75) ← best new
    SC10 Thought coherence  AUROC ~0.80 within, ~1.0  cross (exp75) ← best new
    SC11 Ans-Q mismatch     AUROC ~0.69 within, ~0.95 cross (exp75)

  Density-based signals — requires ≥20 unlabeled calibration chains:
    GMM-E3   GMM k=2-neg + GMM k=4 rank-ensemble on PCA-32 of max-pooled step
             embeddings (exp60: +0.21 AUROC over chain-1 only)
             AUROC 0.680 within-HP, 0.643 cross-NQ  (exp38/exp60)
    Obs-pool Max-pool of Observation embeddings; KNN against calibration pool
             AUROC 0.742 within-HP, 0.622 cross-NQ  (exp40)

  Supervised upgrade — requires 50 labeled chains (calibrate_supervised):
    QARA-Obs SupCon adapter on obs-pool embeddings + GMM
             AUROC 0.675 cross-NQ (exp42), best statistically significant result
             5-fold CV p=0.025, 5/5 folds positive (sign test p=0.031)

  Retrieval quality (no calibration):
    retrieval_quality() returns per-step cosine similarity(question, observation)
    Correct agents get mean_sim=0.568 vs wrong=0.441 (Δ+0.128, exp43)

  Failure mode prediction (no calibration, exp70):
    LabelFreeResult.failure_mode — detected failure pattern:
      "retrieval_fail"  observation contains no-results signal
      "repeated_query"  same search issued twice
      "long_chain"      4+ steps (budget exhaustion)
      "empty_answer"    final answer absent or trivially short
      "low_retrieval_quality"  mean cosine(q, obs) < 0.30
      None              no clear failure mode detected

  Combined SC8 = SC1-6 ensemble + GMM-E3:
             AUROC 0.883 within-HP, 0.664 cross-NQ  (exp39)

  Streaming / mid-chain scoring (exp67):
    score_prefix(question, steps, t) — score using only first t steps
    AUROC at t=2 equals or exceeds full-chain AUROC; enables early termination.

Calibration efficiency (exp44):
  n=0  chains: SC2 alone, AUROC 0.879 / 0.570 — deploy immediately
  n≥20 chains: SC8 (SC2+GMM) stable, AUROC 0.883 / 0.664
  n≥50 labeled: QARA-Obs upgrade, AUROC ~0.675 cross-domain

Usage:

    from qppg_service.label_free_scorer import LabelFreeScorer

    scorer = LabelFreeScorer()

    # Optional: calibrate with unlabeled chains (improves GMM + obs-pool)
    scorer.calibrate(chains)            # chains = [{"question":…, "steps":[…]}]

    # Optional supervised upgrade (requires correctness labels):
    scorer.calibrate_supervised(labeled_chains)  # adds QARA-obs adapter

    # Score a completed chain
    result = scorer.score(question, steps, final_answer, finished)
    print(result.risk_score)            # float [0,1]; higher = more likely wrong
    print(result.needs_review)          # True if risk > review_threshold
    print(result.failure_mode)          # "retrieval_fail" | "long_chain" | … | None

    # Streaming: score after each step (exp67 — AUROC at t=2 ≥ full-chain)
    for t in range(1, len(steps) + 1):
        partial = scorer.score_prefix(question, steps, t)
        if partial.risk_score > 0.70:
            break  # early termination saves ~45% of remaining tokens

    # Retrieval quality (separate from risk):
    rq = scorer.retrieval_quality(question, steps)
    print(rq["mean_sim"])               # 0.568 correct vs 0.441 wrong (avg)

    # Batch scoring
    results = scorer.score_batch(completed_chains)
"""

from __future__ import annotations

import re
import warnings
from dataclasses import dataclass, field
from typing import List, Dict, Any

import numpy as np

warnings.filterwarnings("ignore")

# ── Uncertainty lexicon (SC5) ────────────────────────────────────────────────

_UNCERTAINTY_WORDS: list[str] = [
    "not sure", "unclear", "might", "may not", "possibly", "perhaps",
    "however", "although", "uncertain", "seems", "appears",
    "probably", "could be", "i think", "i believe", "not certain",
    "unsure", "i'm not", "don't know", "confusing", "it's possible",
    "not clear", "ambiguous",
]

_STOP_WORDS: frozenset[str] = frozenset({
    "the", "a", "an", "is", "it", "in", "of", "to", "and", "or", "that",
    "this", "was", "are", "for", "with", "on", "at", "by", "be", "has",
    "had", "i", "its", "as", "from",
})


# ── Behavioral feature extractors ────────────────────────────────────────────

def _sc1_completion(run: dict) -> float:
    """1.0 = fallback call (risky), 0.0 = clean Finish[] (lower risk)."""
    return 0.0 if run.get("finished", False) else 1.0


def _sc2_step_count(run: dict) -> float:
    """Number of Search steps. More steps → potentially more confused."""
    return float(sum(1 for s in run.get("steps", []) if s.get("action_type") == "Search"))


def _sc3_thought_variance(run: dict, embedder) -> float:
    """Mean per-dim variance of Thought embeddings. High = wandering reasoning."""
    thoughts = [s.get("thought", "").strip() for s in run.get("steps", [])
                if s.get("thought", "").strip()]
    if len(thoughts) < 2:
        return 0.0
    embs = embedder.encode(thoughts, normalize_embeddings=True, show_progress_bar=False)
    return float(embs.var(axis=0).mean())


def _sc5_uncertainty_words(run: dict) -> float:
    """Hedging language density per 100 words in all Thoughts."""
    all_thoughts = " ".join(s.get("thought", "") for s in run.get("steps", [])).lower()
    if not all_thoughts.strip():
        return 0.0
    word_count = max(1, len(all_thoughts.split()))
    hit_count = sum(all_thoughts.count(uw) for uw in _UNCERTAINTY_WORDS)
    return float(hit_count / word_count * 100)


def _sc6_answer_gap(run: dict) -> float:
    """1 - overlap(final_answer_words, thought_words). High = answer unexplained by reasoning."""
    final_words = set(re.findall(r'\w+', run.get("final_answer", "").lower())) - _STOP_WORDS
    thought_text = " ".join(s.get("thought", "") for s in run.get("steps", [])).lower()
    thought_words = set(re.findall(r'\w+', thought_text)) - _STOP_WORDS
    if not final_words:
        return 0.5
    return 1.0 - len(final_words & thought_words) / len(final_words)


def _sc8_backtracking_rate(run: dict) -> float:
    """Fraction of Search steps repeating an already-tried action_arg (exp75).
    High = looping agent = risky."""
    steps = run.get("steps", [])
    seen_args: set[str] = set()
    repeats = 0
    searches = 0
    for s in steps:
        if s.get("action_type") == "Search":
            arg = str(s.get("action_arg", "")).lower().strip()
            searches += 1
            if arg in seen_args:
                repeats += 1
            seen_args.add(arg)
    return repeats / max(searches, 1)


def _sc9_obs_utilization(run: dict) -> float:
    """Fraction of observation words referenced in subsequent thoughts (exp75).
    Low = agent ignoring retrieved information = risky."""
    steps = run.get("steps", [])
    if len(steps) < 2:
        return 0.0
    utilized = 0
    total_obs_words = 0
    for i, s in enumerate(steps[:-1]):
        obs = s.get("observation", "")
        if not obs.strip():
            continue
        obs_words = set(re.findall(r'\w+', obs.lower())) - _STOP_WORDS
        if not obs_words:
            continue
        later_text = " ".join(
            s2.get("thought", "") for s2 in steps[i + 1:]
        ).lower()
        later_words = set(re.findall(r'\w+', later_text)) - _STOP_WORDS
        overlap = len(obs_words & later_words)
        utilized += overlap
        total_obs_words += len(obs_words)
    if total_obs_words == 0:
        return 0.0
    # Invert: low utilization = high risk
    return 1.0 - min(1.0, utilized / total_obs_words)


def _sc10_thought_coherence_drop(run: dict) -> float:
    """Fraction of step transitions where the thought fails to reference the
    prior observation (exp75). High = incoherent reasoning = risky."""
    steps = run.get("steps", [])
    if len(steps) < 2:
        return 0.0
    incoherent = 0
    transitions = 0
    for i in range(len(steps) - 1):
        obs = steps[i].get("observation", "")
        next_thought = steps[i + 1].get("thought", "")
        if not obs.strip():
            continue
        transitions += 1
        obs_words = set(re.findall(r'\w+', obs.lower())) - _STOP_WORDS
        thought_words = set(re.findall(r'\w+', next_thought.lower())) - _STOP_WORDS
        if not obs_words or not (obs_words & thought_words):
            incoherent += 1
    return incoherent / max(transitions, 1)


def _sc11_answer_question_mismatch(run: dict) -> float:
    """Vagueness and brevity ratio of the final answer (exp75).
    Vague or too-short answers = risky."""
    final_answer = run.get("final_answer", "")
    question = run.get("question", "")
    if not final_answer.strip():
        return 1.0
    # Vagueness: uncertainty-phrase density
    ans_lower = final_answer.lower()
    vague_phrases = [
        "i don't know", "not sure", "unclear", "unable to",
        "cannot determine", "no information", "not certain",
        "conflicting", "consult", "it depends",
    ]
    vague_count = sum(ans_lower.count(p) for p in vague_phrases)
    vagueness = min(1.0, vague_count / 2.0)

    # Brevity: answer is too short relative to question
    ans_words = len(final_answer.split())
    q_words = max(1, len(question.split()))
    brevity = max(0.0, 1.0 - ans_words / (q_words * 3))

    return float(np.mean([vagueness, brevity]))


def _sc12_risk_monotone_slope(run: dict) -> float:
    """Slope of per-step thought-observation coherence (exp107: AUROC 0.7465).

    Measures whether the agent's integration of retrieved evidence WORSENS over
    the course of the chain.  In correct chains the agent steadily refines its
    understanding; in failing chains the thought-obs overlap drops in later steps
    (agent stops using what it finds).

    Signal: (mean_gap_second_half − mean_gap_first_half) centered at 0.5.
    High value (> 0.5) → coherence collapsing → riskier chain.

    Returns 0.5 (neutral) when there are fewer than 2 search steps.
    Only included in ensemble when n_steps ≥ 2 (exp107).
    """
    search_steps = [
        s for s in run.get("steps", [])
        if s.get("action_type") == "Search" and s.get("observation", "").strip()
    ]
    if len(search_steps) < 2:
        return 0.5   # neutral — can't measure slope

    per_step_gap = []
    for s in search_steps:
        t_toks = set(re.findall(r'\w+', s.get("thought", "").lower())) - _STOP_WORDS
        o_toks = set(re.findall(r'\w+', s.get("observation", "").lower())) - _STOP_WORDS
        if t_toks | o_toks:
            gap = 1.0 - len(t_toks & o_toks) / len(t_toks | o_toks)
        else:
            gap = 0.5
        per_step_gap.append(gap)

    mid = max(1, len(per_step_gap) // 2)
    first_half  = float(np.mean(per_step_gap[:mid]))
    second_half = float(np.mean(per_step_gap[mid:]))
    slope = second_half - first_half          # positive = coherence dropping = riskier
    return float(np.clip(slope + 0.5, 0.0, 1.0))   # centred at 0.5, range [0, 1]


def _extract_observation_texts(run: dict, max_obs_len: int = 200) -> list[str]:
    """Extract observation texts from a run's steps."""
    return [s["observation"][:max_obs_len] for s in run.get("steps", [])
            if s.get("observation", "").strip()]


def _rank_norm(scores: np.ndarray) -> np.ndarray:
    """Map scores to [0, 1] via percentile rank."""
    from scipy.stats import rankdata
    return rankdata(scores.astype(float)) / len(scores)


# ── Failure mode detection (exp70) ───────────────────────────────────────────

_RETRIEVAL_FAIL_PATTERNS: tuple[str, ...] = (
    "no results", "not found", "no information", "couldn't find",
    "no relevant", "no matches", "i don't know", "sorry, i",
    "unable to find", "no data",
)

def _detect_failure_mode(
    steps: list[dict],
    final_answer: str,
    question: str,
    retrieval_mean_sim: float | None = None,
) -> str | None:
    """
    Rule-based failure mode detection (exp70 top features).

    Returns the first detected failure mode string, or None if no clear mode.
    Priority order: empty_answer > repeated_query > retrieval_fail > long_chain
    > low_retrieval_quality.
    """
    # FP0: empty / trivially short final answer
    if len(final_answer.strip()) < 5:
        return "empty_answer"

    # FP1: repeated search query
    action_args = [s.get("action_arg", "").lower().strip() for s in steps
                   if s.get("action_arg", "").strip()]
    if len(action_args) != len(set(action_args)) and len(action_args) >= 2:
        return "repeated_query"

    # FP2: retrieval failure signal in any observation
    for s in steps:
        obs_lower = s.get("observation", "").lower()
        if any(pat in obs_lower for pat in _RETRIEVAL_FAIL_PATTERNS):
            return "retrieval_fail"

    # FP2.5: no-evidence hallucination — final answer has zero word overlap with
    # any retrieved observation (exp93: AUROC 0.583 within-HP, catches complete
    # hallucination). Only fires when there are observations to compare against.
    final_toks = {w.lower() for w in re.findall(r'\b[a-zA-Z]{3,}\b', final_answer)
                  if w.lower() not in _STOP_WORDS}
    search_obs = [s.get("observation", "") for s in steps
                  if s.get("action_type") == "Search" and s.get("observation", "").strip()]
    if final_toks and search_obs:
        all_obs_toks: set[str] = set()
        for obs in search_obs:
            all_obs_toks |= {w.lower() for w in re.findall(r'\b[a-zA-Z]{3,}\b', obs)
                             if w.lower() not in _STOP_WORDS}
        if all_obs_toks and not (final_toks & all_obs_toks):
            return "no_evidence"

    # FP3: long chain (budget exhaustion proxy)
    if len(steps) >= 4:
        return "long_chain"

    # FP4: low retrieval quality (if pre-computed)
    if retrieval_mean_sim is not None and retrieval_mean_sim < 0.30:
        return "low_retrieval_quality"

    return None


# ── Result dataclass ──────────────────────────────────────────────────────────

@dataclass
class LabelFreeResult:
    risk_score: float
    """Combined risk score [0, 1]. Higher = more likely to be wrong."""

    needs_review: bool
    """True if risk_score exceeds review_threshold."""

    behavioral_score: float | None
    """SC1-SC6 ensemble component (always available, no calibration needed)."""

    gmm_score: float | None
    """GMM-E3 component (None if not calibrated)."""

    obs_score: float | None
    """Observation-pool KNN component (None if not calibrated or no observations)."""

    qara_score: float | None = None
    """QARA-obs supervised component (None if calibrate_supervised not called)."""

    failure_mode: str | None = None
    """Detected failure pattern (exp70): 'retrieval_fail' | 'repeated_query' |
    'long_chain' | 'empty_answer' | 'low_retrieval_quality' | None."""

    components: dict = field(default_factory=dict)
    """Individual scores: sc1, sc2, sc3, sc5, sc6, sc8, sc9, sc10, sc11, sc12 and
    polarity-corrected values. SC9/SC10 have strongest cross-domain AUROC (exp75).
    SC12 (risk_monotone_slope): coherence drop across steps, AUROC 0.7465 (exp107)."""

    calibration_size: int = 0
    """Number of chains used to calibrate density-based models."""

    calibration_state: str = "behavioral_only"
    """Calibration regime for this score.

    - ``"fitted"``            — GMM + obs-pool active (calibration_size ≥ min_calibration).
                                 Risk score uses behavioral + density signals.
    - ``"insufficient_data"`` — calibrate() called but pool below min_calibration (default 50).
                                 Risk score is behavioral only; GMM not reliable.
    - ``"behavioral_only"``   — calibrate() never called.  Risk score uses behavioral
                                 signals only (SC1–SC12 ensemble).  No density component.

    Use this field to communicate scoring regime to downstream consumers so they
    can apply appropriate confidence adjustments.
    """

    def to_dict(self) -> dict:
        return {
            "risk_score":         round(self.risk_score, 4),
            "needs_review":       self.needs_review,
            "behavioral_score":   (round(self.behavioral_score, 4)
                                   if self.behavioral_score is not None else None),
            "gmm_score":          (round(self.gmm_score, 4)
                                   if self.gmm_score is not None else None),
            "obs_score":          (round(self.obs_score, 4)
                                   if self.obs_score is not None else None),
            "failure_mode":       self.failure_mode,
            "calibration_size":   self.calibration_size,
            "calibration_state":  self.calibration_state,
            "components":         {k: round(v, 4) for k, v in self.components.items()},
        }


# ── Main scorer class ─────────────────────────────────────────────────────────

class LabelFreeScorer:
    """
    Label-free risk scorer for completed ReAct chains.

    No labeled data required at any stage.  Two modes:

    1. Behavioral-only (no calibration):
       Uses SC1-SC6 features from exp39. AUROC ~0.87 within-domain.
       Call score() immediately after construction.

    2. Full (with calibration):
       Additionally uses GMM-E3 and Obs-pool-KNN from exp38/40.
       AUROC 0.87 within-domain, 0.65 cross-domain after calibrate().
       Requires ≥20 unlabeled chains. Recommended ≥50 for stability.

    Parameters
    ----------
    embed_model : str
        SentenceTransformer model name. Must be pre-cached locally.
        Default: "all-MiniLM-L6-v2"
    review_threshold : float
        risk_score above this → needs_review=True.  Default 0.65.
    min_calibration : int
        Minimum chains before density models are fitted. Default 50.
        Values below 50 produce under-determined GMMs (2-component GMM on 11
        behavioral features with fewer than 50 samples is numerically unstable).
        Pass min_calibration=20 only if you understand the instability risk.
    """

    def __init__(
        self,
        embed_model: str = "all-MiniLM-L6-v2",
        review_threshold: float = 0.65,
        min_calibration: int = 50,
    ):
        self.embed_model      = embed_model
        self.review_threshold = review_threshold
        self.min_calibration  = min_calibration

        self._embedder = None

        # Label-free calibration state (None until calibrate() called)
        self._pca        = None   # sklearn PCA fitted on chain-1 embs
        self._gmm2       = None   # GaussianMixture k=2
        self._gmm4       = None   # GaussianMixture k=4
        self._obs_embs   = None   # mean-pooled obs embeddings for KNN (n_cal, 384)
        self._cal_chain1 = None   # chain-1 embeddings (n_cal, 384)
        self._n_cal      = 0

        # Supervised QARA-obs adapter (None until calibrate_supervised() called)
        self._qara_adapter    = None
        self._qara_obs_embs   = None   # adapted obs embeddings for KNN
        self._n_supervised    = 0

    # ── Lazy embedder ─────────────────────────────────────────────────────────

    def _get_embedder(self):
        if self._embedder is None:
            from sentence_transformers import SentenceTransformer
            self._embedder = SentenceTransformer(self.embed_model)
        return self._embedder

    # ── Chain formatting helpers ───────────────────────────────────────────────

    @staticmethod
    def _fmt_chain1(question: str, steps: list[dict]) -> str:
        """Chain-1 context string (same format as exp38/QPPGService)."""
        if steps:
            s = steps[0]
            atype = s.get("action_type", "Action")
            aarg  = str(s.get("action_arg", ""))[:300]
            thought = str(s.get("thought", ""))[:200]
            return (f"TASK: {question[:300]}\n"
                    f"STEP 1: Thought: {thought} | Action: {atype}[{aarg}]\n"
                    f"CURRENT: {atype}[{aarg[:300]}]")
        return f"TASK: {question[:300]}\nCURRENT: Start[]"

    # ── Calibration ────────────────────────────────────────────────────────────

    def calibrate(self, chains: list[dict]) -> dict:
        """
        Fit GMM-E3 and observation-pool KNN from unlabeled chains.

        Parameters
        ----------
        chains : list of dict, each containing:
            "question" : str
            "steps"    : list of ReAct step dicts

        Returns summary dict with calibration statistics.

        Notes
        -----
        - Requires ≥ min_calibration (default 20) chains.
        - Only chain-1 context and observations are used (not final answers).
        - Repeated calls replace the previous calibration.
        """
        if len(chains) < self.min_calibration:
            import warnings as _warnings
            _warnings.warn(
                f"LabelFreeScorer.calibrate(): refusing to fit GMM with only {len(chains)} chains "
                f"(minimum is {self.min_calibration}). A 2-component GMM with 11 behavioral "
                f"features is under-determined below n=50 and will produce unreliable density "
                f"estimates. Falling back to raw behavioral scoring. "
                f"Collect ≥{self.min_calibration} chains before calling calibrate().",
                UserWarning,
                stacklevel=2,
            )
            return {
                "calibrated": False,
                "reason": f"Need ≥ {self.min_calibration} chains, got {len(chains)}",
                "n_chains": len(chains),
            }

        from sklearn.decomposition import PCA
        from sklearn.mixture import GaussianMixture
        from sklearn.neighbors import NearestNeighbors

        embedder = self._get_embedder()

        # ── Per-step embeddings → max-pool per chain (exp60: +0.21 AUROC) ───
        # Each step text = "Thought: … Action: … Obs: …"
        # Max-pool across all steps captures the most anomalous step signal.
        all_step_texts: list[str] = []
        step_offsets: list[int] = [0]
        for c in chains:
            for s in c.get("steps", []):
                t = (f"Thought: {s.get('thought', '')} "
                     f"Action: {s.get('action_arg', '')} "
                     f"Obs: {s.get('observation', '')[:200]}")
                all_step_texts.append(t[:512])
            step_offsets.append(len(all_step_texts))

        if all_step_texts:
            all_step_embs = embedder.encode(
                all_step_texts, normalize_embeddings=True,
                show_progress_bar=False, batch_size=256,
            )
        else:
            all_step_embs = np.zeros((1, 384))

        chain_embs = []
        for i, c in enumerate(chains):
            start, end = step_offsets[i], step_offsets[i + 1]
            if start < end:
                chain_embs.append(all_step_embs[start:end].max(axis=0))
            else:
                # Fallback: encode chain-1 context
                ctx = self._fmt_chain1(c["question"], c.get("steps", []))
                chain_embs.append(
                    embedder.encode([ctx], normalize_embeddings=True,
                                    show_progress_bar=False)[0]
                )
        chain_embs_arr = np.array(chain_embs, dtype=np.float32)
        self._cal_chain1 = chain_embs_arr  # kept for GMM scoring at test time

        # PCA-32 + GMM fit on max-pooled chain embeddings
        n_comp = min(32, len(chains) - 1, chain_embs_arr.shape[1])
        self._pca = PCA(n_components=n_comp, random_state=42).fit(chain_embs_arr)
        cal_pca = self._pca.transform(chain_embs_arr)
        # covariance_type="diag" keeps payload small for federated sharing:
        # 2×32 + 4×32 = 192 floats vs 2×32×32 + 4×32×32 = 6144 for "full"
        self._gmm2 = GaussianMixture(2, n_init=5, covariance_type="diag",
                                     random_state=42).fit(cal_pca)
        self._gmm4 = GaussianMixture(4, n_init=5, covariance_type="diag",
                                     random_state=42).fit(cal_pca)

        # ── Observation embeddings (for Obs-pool KNN) ────────────────────────
        obs_texts_per_chain = [_extract_observation_texts(c) for c in chains]
        all_obs_texts = [t for obs in obs_texts_per_chain for t in obs]
        obs_counts    = [len(obs) for obs in obs_texts_per_chain]

        if all_obs_texts:
            all_obs_embs = embedder.encode(all_obs_texts, normalize_embeddings=True,
                                            show_progress_bar=False, batch_size=256)
            # Max-pool per chain (consistent with step embedding strategy)
            obs_pool = []
            offset = 0
            for cnt in obs_counts:
                if cnt > 0:
                    obs_pool.append(all_obs_embs[offset:offset + cnt].max(axis=0))
                else:
                    obs_pool.append(np.zeros(all_obs_embs.shape[1]))
                offset += cnt
            self._obs_embs = np.array(obs_pool, dtype=np.float32)
        else:
            self._obs_embs = None

        self._n_cal = len(chains)
        n_with_obs = sum(1 for c in chains
                         if any(s.get("observation") for s in c.get("steps", [])))
        return {
            "calibrated":       True,
            "n_chains":         len(chains),
            "n_with_obs":       n_with_obs,
            "embedding_strategy": "max_pool_per_step",
            "pca_variance":     round(float(self._pca.explained_variance_ratio_.sum()), 3),
            "gmm_fitted":       True,
            "obs_pool_fitted":  self._obs_embs is not None,
        }

    def is_calibrated(self) -> bool:
        return self._pca is not None and self._n_cal >= self.min_calibration

    # ── Scoring ────────────────────────────────────────────────────────────────

    def score(
        self,
        question: str,
        steps: list[dict],
        final_answer: str = "",
        finished: bool = True,
    ) -> LabelFreeResult:
        """
        Score a single completed chain.

        Parameters
        ----------
        question    : str   — the original task/question
        steps       : list  — ReAct steps (each a dict with thought/action_type/
                              action_arg/observation keys)
        final_answer: str   — the agent's final answer text
        finished    : bool  — True if agent called Finish[], False if fallback

        Returns
        -------
        LabelFreeResult
        """
        run = {"question": question, "steps": steps,
               "final_answer": final_answer, "finished": finished}

        embedder = self._get_embedder()
        components: dict[str, float] = {}

        # ── Behavioral features (always available) ────────────────────────────
        sc1  = _sc1_completion(run)
        sc2  = _sc2_step_count(run)
        sc3  = _sc3_thought_variance(run, embedder)
        sc5  = _sc5_uncertainty_words(run)
        sc6  = _sc6_answer_gap(run)
        # New signals from exp75/exp77 (topic-agnostic).
        # NOTE: SC8 (backtracking) excluded from ensemble — 0.50 AUROC on real HP data
        # because real chains rarely repeat searches (exp77: 148/200 chains have 2 steps).
        # SC9/SC10 also weak on short chains (need ≥3 steps); kept for completeness.
        sc8  = _sc8_backtracking_rate(run)
        sc9  = _sc9_obs_utilization(run)
        sc10 = _sc10_thought_coherence_drop(run)
        sc11 = _sc11_answer_question_mismatch(run)
        sc12 = _sc12_risk_monotone_slope(run)

        components.update(sc1=sc1, sc2=sc2, sc3=sc3, sc5=sc5, sc6=sc6,
                          sc8=sc8, sc9=sc9, sc10=sc10, sc11=sc11, sc12=sc12)

        # All signals are positive polarity (higher = riskier).
        # Normalise to [0, 1] comparable scale:
        #   sc1  ∈ {0, 1}     → keep as-is
        #   sc2  ∈ [0, 6]     → divide by 6 (max observed steps in training)
        #   sc3  ∈ [0, ~0.01] → scale by 100 to put in [0, 1]
        #   sc5  ∈ [0, ~2]    → divide by 2
        #   sc6  ∈ [0, 1]     → keep as-is
        #   sc9  ∈ [0, 1]     → keep as-is (not included when n_steps < 3)
        #   sc10 ∈ [0, 1]     → keep as-is (not included when n_steps < 3)
        #   sc11 ∈ [0, 1]     → keep as-is (strong on real data: HP AUROC 0.760, exp77)
        #   sc12 ∈ [0, 1]     → keep as-is (only when n_search_steps ≥ 2, exp107 AUROC 0.7465)
        sc2_norm = min(1.0, sc2 / 6.0)
        sc3_norm = min(1.0, sc3 * 100.0)
        sc5_norm = min(1.0, sc5 / 2.0)
        n_steps = len(run.get("steps", []))
        n_search = sum(1 for s in run.get("steps", [])
                       if s.get("action_type") == "Search" and s.get("observation", "").strip())
        # SC11 always included; SC9/SC10 only when chain has ≥3 steps (exp77)
        # SC12 only when chain has ≥2 search steps (exp107)
        base_signals = [sc1, sc2_norm, sc3_norm, sc5_norm, sc6, sc11]
        base_weights = [1.0, 1.0,      1.0,      1.0,      1.0, 1.2]
        if n_steps >= 3:
            base_signals.extend([sc9, sc10])
            base_weights.extend([1.3, 1.3])
        if n_search >= 2:
            base_signals.append(sc12)
            base_weights.append(1.4)   # exp107: strongest early-step signal (+0.156 vs SC_OLD)
        behavioral_score = float(np.average(base_signals, weights=base_weights))
        components["behavioral_score"] = behavioral_score

        # ── Density-based scores (only if calibrated) ─────────────────────────
        gmm_score = None
        obs_score = None

        if self.is_calibrated():
            # Max-pool step embeddings at test time (matches calibrate(), exp60: +0.21 AUROC)
            step_texts = [
                (f"Thought: {s.get('thought', '')} "
                 f"Action: {s.get('action_arg', '')} "
                 f"Obs: {s.get('observation', '')[:200]}")[:512]
                for s in steps
            ] if steps else []
            if step_texts:
                step_embs = embedder.encode(step_texts, normalize_embeddings=True,
                                            show_progress_bar=False)
                test_chain_emb = step_embs.max(axis=0)  # elementwise max across steps
            else:
                ctx1 = self._fmt_chain1(question, steps)
                test_chain_emb = embedder.encode([ctx1], normalize_embeddings=True,
                                                  show_progress_bar=False)[0]

            # GMM-E3: negate GMM-k2, keep GMM-k4, rank-average
            test_pca = self._pca.transform(test_chain_emb[np.newaxis, :])
            cal_pca  = self._pca.transform(self._cal_chain1)
            s2_raw   = -float(self._gmm2.score_samples(test_pca))  # inverted polarity
            s4_raw   = -float(self._gmm4.score_samples(test_pca))  # correct polarity

            # Calibration reference distribution for normalisation
            cal_s2 = -self._gmm2.score_samples(cal_pca)
            cal_s4 = -self._gmm4.score_samples(cal_pca)
            # ECDF-based rank: fraction of cal points BELOW test score
            r2 = float((cal_s2 < s2_raw).mean())  # high = test is outlier in raw
            r4 = float((cal_s4 < s4_raw).mean())
            # GMM-k2 inverted: flip r2 → (1 - r2) = how typical is test (riskier if typical)
            gmm_score = float((1.0 - r2 + r4) / 2.0)
            components["gmm_score"] = gmm_score

            # Observation-pool KNN
            if self._obs_embs is not None:
                obs_texts = _extract_observation_texts(run)
                if obs_texts:
                    obs_embs = embedder.encode(obs_texts, normalize_embeddings=True,
                                               show_progress_bar=False)
                    test_obs_emb = obs_embs.max(axis=0)  # max-pool (matches calibration)
                    dists = np.linalg.norm(self._obs_embs - test_obs_emb[np.newaxis, :],
                                           axis=1)
                    knn_dist = float(dists.min())
                    # ECDF: fraction of cal obs whose NN distance is smaller than test
                    cal_obs_dists = np.linalg.norm(
                        self._obs_embs - self._obs_embs.mean(axis=0)[np.newaxis, :], axis=1
                    )
                    obs_score = float((cal_obs_dists < knn_dist).mean())
                    components["obs_score"] = obs_score

        # ── QARA-obs supervised upgrade (Tier 2) ──────────────────────────────
        qara_score = None
        if self._qara_adapter is not None:
            import torch
            obs_texts = _extract_observation_texts(run)
            if obs_texts:
                obs_embs_raw = self._get_embedder().encode(
                    obs_texts, normalize_embeddings=True, show_progress_bar=False)
                obs_mean = torch.tensor(obs_embs_raw.mean(axis=0), dtype=torch.float32)
                with torch.no_grad():
                    adapted = self._qara_adapter(obs_mean.unsqueeze(0)).numpy()[0]
                # GMM-E3 on adapted space
                ap = self._qara_pca.transform(adapted[np.newaxis, :])
                cal_ap = self._qara_pca.transform(self._qara_obs_embs)
                s2r = -float(self._qara_gmm2.score_samples(ap))
                s4r = -float(self._qara_gmm4.score_samples(ap))
                cal_s2r = -self._qara_gmm2.score_samples(cal_ap)
                cal_s4r = -self._qara_gmm4.score_samples(cal_ap)
                rq2 = float((cal_s2r < s2r).mean())
                rq4 = float((cal_s4r < s4r).mean())
                qara_score = float((1.0 - rq2 + rq4) / 2.0)
                components["qara_score"] = qara_score

        # ── Failure mode detection (exp70, rule-based, no calibration needed) ──
        retrieval_mean_sim: float | None = None
        obs_steps = [s for s in steps if s.get("action_type") == "Search"
                     and s.get("observation", "").strip()]
        if obs_steps:
            rq = self.retrieval_quality(question, steps)
            retrieval_mean_sim = rq["mean_sim"]
        failure_mode = _detect_failure_mode(steps, final_answer, question,
                                            retrieval_mean_sim)

        # ── Combined risk score ────────────────────────────────────────────────
        active = [behavioral_score]
        if gmm_score   is not None: active.append(gmm_score)
        if qara_score  is not None: active.append(qara_score)
        elif obs_score is not None: active.append(obs_score)   # use obs_score if no QARA
        risk_score = float(np.mean(active))

        # Determine calibration_state for transparency
        if gmm_score is not None:
            _cal_state = "fitted"
        elif self._n_cal > 0:
            _cal_state = "insufficient_data"
        else:
            _cal_state = "behavioral_only"

        return LabelFreeResult(
            risk_score=round(risk_score, 4),
            needs_review=risk_score >= self.review_threshold,
            behavioral_score=round(behavioral_score, 4),
            gmm_score=(round(gmm_score, 4) if gmm_score is not None else None),
            obs_score=(round(obs_score, 4) if obs_score is not None else None),
            qara_score=(round(qara_score, 4) if qara_score is not None else None),
            failure_mode=failure_mode,
            components=components,
            calibration_size=self._n_cal,
            calibration_state=_cal_state,
        )

    def calibrate_supervised(self, labeled_chains: list[dict]) -> dict:
        """
        Tier-2 upgrade: train QARA SupCon adapter on obs-pool embeddings.

        Requires correctness labels — each chain dict must have a boolean
        "correct" field.  Recommended ≥50 chains for stable results.

        After calling this, score() uses the adapted obs-pool embeddings,
        improving cross-domain AUROC from ~0.622 to ~0.675 (exp42).

        Parameters
        ----------
        labeled_chains : list of dict, each with:
            "question" : str
            "steps"    : list of ReAct step dicts
            "correct"  : bool

        Returns summary dict with training statistics.
        """
        import copy
        import torch
        import torch.nn as nn
        import torch.optim as optim
        from sklearn.decomposition import PCA
        from sklearn.mixture import GaussianMixture

        if len(labeled_chains) < 10:
            return {"trained": False, "reason": f"Need ≥10 labeled chains, got {len(labeled_chains)}"}

        labels = np.array([1 if not c.get("correct", True) else 0 for c in labeled_chains],
                          dtype=np.int64)
        if not (labels == 0).any() or not (labels == 1).any():
            return {"trained": False, "reason": "Need both correct and incorrect chains"}

        embedder = self._get_embedder()

        # Build obs-pool embeddings
        obs_texts_per = [_extract_observation_texts(c) for c in labeled_chains]
        all_obs = [t for obs in obs_texts_per for t in obs]
        if not all_obs:
            return {"trained": False, "reason": "No observations found in chains"}

        all_embs = embedder.encode(all_obs, normalize_embeddings=True,
                                   show_progress_bar=False, batch_size=256)
        obs_pool, offset = [], 0
        for cnt_list in obs_texts_per:
            cnt = len(cnt_list)
            if cnt > 0:
                obs_pool.append(all_embs[offset:offset + cnt].mean(axis=0))
            else:
                obs_pool.append(np.zeros(all_embs.shape[1]))
            offset += cnt
        obs_pool = np.array(obs_pool, dtype=np.float32)

        # Build + train QARA adapter (SupCon)
        in_dim = obs_pool.shape[1]

        class _QARAAdapter(nn.Module):
            def __init__(self):
                super().__init__()
                self.net = nn.Sequential(nn.Linear(in_dim, 256), nn.ReLU(),
                                         nn.LayerNorm(256), nn.Linear(256, 64))
            def forward(self, x):
                return nn.functional.normalize(self.net(x), dim=-1)

        def _supcon(embs, labs, temperature=0.1):
            N = embs.shape[0]
            sim = torch.mm(embs, embs.T) / temperature
            lbl = labs.unsqueeze(1)
            pos = (lbl == lbl.T).float(); pos.fill_diagonal_(0.0)
            eye = torch.eye(N, device=embs.device); dnom = 1.0 - eye
            sim_s = sim - (sim * dnom - 1e9 * eye).max(1, keepdim=True)[0].detach()
            exp_s = torch.exp(sim_s) * dnom
            log_p = sim_s - torch.log(exp_s.sum(1, keepdim=True) + 1e-8)
            return (-(pos * log_p).sum(1) / pos.sum(1).clamp(min=1)).mean()

        adapter   = _QARAAdapter()
        optimizer = optim.Adam(adapter.parameters(), lr=3e-4)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
        X = torch.tensor(obs_pool, dtype=torch.float32)
        y = torch.tensor(labels,   dtype=torch.long)
        rng = np.random.default_rng(42)
        best_loss, best_state = float("inf"), None

        for epoch in range(200):
            idx = rng.permutation(len(X)); loss_sum = 0.0; nb = 0
            for st in range(0, len(X), 64):
                bi = idx[st:st + 64]
                if not ((y[bi] == 1).any() and (y[bi] == 0).any()): continue
                optimizer.zero_grad()
                loss = _supcon(adapter(X[bi]), y[bi].float()); loss.backward()
                torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
                optimizer.step(); loss_sum += loss.item(); nb += 1
            scheduler.step()
            avg = loss_sum / max(nb, 1)
            if avg < best_loss: best_loss = avg; best_state = copy.deepcopy(adapter.state_dict())

        adapter.load_state_dict(best_state); adapter.eval()
        self._qara_adapter = adapter

        with torch.no_grad():
            self._qara_obs_embs = adapter(X).numpy()
        self._n_supervised = len(labeled_chains)

        # Also fit GMM on adapted obs embeddings for label-free scoring
        n_comp = min(32, self._qara_obs_embs.shape[0] - 1, self._qara_obs_embs.shape[1])
        pca_q = PCA(n_components=n_comp, random_state=42).fit(self._qara_obs_embs)
        qp = pca_q.transform(self._qara_obs_embs)
        self._qara_gmm2 = GaussianMixture(2, n_init=5, random_state=42).fit(qp)
        self._qara_gmm4 = GaussianMixture(4, n_init=5, random_state=42).fit(qp)
        self._qara_pca  = pca_q

        return {
            "trained":      True,
            "n_chains":     len(labeled_chains),
            "n_correct":    int((labels == 0).sum()),
            "n_incorrect":  int((labels == 1).sum()),
            "final_loss":   round(float(best_loss), 4),
            "expected_cross_domain_auroc": "~0.675 (exp42 5-fold CV)",
        }

    def retrieval_quality(self, question: str, steps: list[dict]) -> dict:
        """
        Compute retrieval quality scores for a completed chain (exp43 signals).

        Returns cosine similarity between the question embedding and each
        observation embedding, plus aggregate statistics.

        No calibration required.  Correct agents average mean_sim=0.568,
        incorrect agents average mean_sim=0.441 (Δ+0.128, exp43).

        Returns
        -------
        dict with:
          mean_sim    : float  — mean cosine similarity (higher = better retrieval)
          max_sim     : float  — max similarity (best observation)
          min_sim     : float  — min similarity (worst observation; Δ+0.232 correct vs wrong)
          var_sim     : float  — variance (higher = inconsistent retrieval quality)
          n_relevant  : int    — count of observations with sim > 0.25
          per_step    : list[float]  — similarity for each search step
          quality_label: str  — "GOOD" / "FAIR" / "POOR"
        """
        embedder = self._get_embedder()
        obs_list = [(i, s.get("observation", "").strip())
                    for i, s in enumerate(steps)
                    if s.get("action_type") == "Search" and s.get("observation", "").strip()]

        if not obs_list:
            return {
                "mean_sim": 0.0, "max_sim": 0.0, "min_sim": 0.0,
                "var_sim": 0.0, "n_relevant": 0, "per_step": [],
                "quality_label": "POOR",
            }

        q_emb = embedder.encode([question], normalize_embeddings=True,
                                  show_progress_bar=False)[0]
        obs_texts = [o[:300] for _, o in obs_list]
        obs_embs  = embedder.encode(obs_texts, normalize_embeddings=True,
                                     show_progress_bar=False)
        sims = (obs_embs @ q_emb).tolist()

        REL_THRESHOLD = 0.25
        mean_s = float(np.mean(sims))
        max_s  = float(np.max(sims))
        min_s  = float(np.min(sims))
        var_s  = float(np.var(sims))
        n_rel  = sum(1 for s in sims if s > REL_THRESHOLD)

        if mean_s >= 0.50:
            label = "GOOD"
        elif mean_s >= 0.35:
            label = "FAIR"
        else:
            label = "POOR"

        return {
            "mean_sim":     round(mean_s, 4),
            "max_sim":      round(max_s,  4),
            "min_sim":      round(min_s,  4),
            "var_sim":      round(var_s,  4),
            "n_relevant":   n_rel,
            "per_step":     [round(s, 4) for s in sims],
            "quality_label": label,
        }

    def retrieval_verification(self, steps: list[dict], final_answer: str) -> dict:
        """
        Token-level retrieval verification (exp93).

        Measures how much of the final answer is grounded in retrieved observations.
        Useful as a supplementary signal for real search engines where the fake-search
        artifact (Haiku generates plausible text for any query) doesn't apply.

        In production with a real search engine, within-domain AUROC is expected to
        be 0.65-0.75 (vs 0.53 on fake-search). Cross-domain: V3 (F1) showed 0.70 on NQ.

        Returns
        -------
        dict with:
          verif_precision : float  — |answer_words ∩ obs_words| / |answer_words|
                                     0.0 = complete hallucination, 1.0 = fully grounded
          verif_recall    : float  — |answer_words ∩ obs_words| / |obs_words|
          verif_f1        : float  — harmonic mean(precision, recall)
          risk            : float  — 1 - verif_precision (higher = more hallucinated)
          no_evidence     : bool   — True if zero token overlap (complete hallucination)
          n_search_steps  : int    — number of search steps with non-empty observations
        """
        _sv_stop = frozenset({
            "the","a","an","is","are","was","were","be","been","have","has","had",
            "do","does","did","will","would","could","should","may","might","can",
            "this","that","these","those","it","its","they","them","their",
            "of","in","on","at","to","for","with","by","from","and","or","but",
            "not","as","into","i","he","she","we","you",
        })

        def _sv_toks(text: str) -> set:
            return {w.lower() for w in re.findall(r'\b[a-zA-Z]{3,}\b', text)
                    if w.lower() not in _sv_stop}

        final_toks = _sv_toks(final_answer)
        all_obs_toks: set = set()
        n_search = 0
        for s in steps:
            if s.get("action_type") == "Search" and s.get("observation", "").strip():
                all_obs_toks |= _sv_toks(s["observation"])
                n_search += 1

        if not final_toks or not all_obs_toks:
            return {
                "verif_precision": 0.0, "verif_recall": 0.0, "verif_f1": 0.0,
                "risk": 1.0, "no_evidence": bool(final_toks and not all_obs_toks),
                "n_search_steps": n_search,
            }

        prec = len(final_toks & all_obs_toks) / len(final_toks)
        rec  = len(final_toks & all_obs_toks) / len(all_obs_toks)
        f1   = 2 * prec * rec / (prec + rec) if prec + rec > 0 else 0.0

        return {
            "verif_precision": round(prec, 4),
            "verif_recall":    round(rec, 4),
            "verif_f1":        round(f1, 4),
            "risk":            round(1.0 - prec, 4),
            "no_evidence":     (prec == 0.0),
            "n_search_steps":  n_search,
        }

    def optimize_weights(
        self,
        labeled_chains: list[dict],
        sigma: float = 0.1,
        n_restarts: int = 3,
    ) -> dict:
        """
        RL-inspired offline weight optimization (exp95).

        Learns optimal signal weights by maximizing a differentiable AUROC surrogate
        (soft Wilcoxon-Mann-Whitney statistic) on labeled training data.

        IMPORTANT: exp95 finding — the optimizer consistently assigns ~97% weight to
        the Sonnet judge and ~0% to other signals. This confirms the judge is the
        dominant signal. Learned weights do NOT reliably transfer cross-domain (they
        overfit to HP patterns). Use J5 fixed weights [SC×1, Judge×3] for cross-domain.

        Use this method when you have labeled chains from a SINGLE deployment domain
        and want to find the optimal signal blend for that specific domain.

        Parameters
        ----------
        labeled_chains : list of dict with:
            "question", "steps", "final_answer", "finished" — chain fields
            "correct" : bool — ground truth correctness label
            Optional: "judge_risk" : float — pre-computed judge risk score [0,5]
        sigma : float
            Sigmoid bandwidth for soft-AUROC (smaller = sharper approximation).
        n_restarts : int
            Number of random restarts for Nelder-Mead (combats local optima).

        Returns
        -------
        dict with:
          weights : dict  — {signal_name: weight}, normalized to sum=1
          auroc   : float — soft-AUROC on training data (optimistic, not CV)
          n_chains : int
          note    : str
        """
        try:
            from scipy.optimize import minimize
            from scipy.stats import rankdata
        except ImportError:
            raise ImportError("scipy required for optimize_weights. pip install scipy")

        labels = np.array([1 - int(c["correct"]) for c in labeled_chains])
        if len(set(labels)) < 2:
            return {"weights": {}, "auroc": 0.5, "n_chains": len(labeled_chains),
                    "note": "Need both correct and incorrect chains."}

        # Compute signals for each chain
        sc_scores = np.array([
            self.score(c["question"], c.get("steps", []),
                       c.get("final_answer", ""), c.get("finished", True)).risk_score
            for c in labeled_chains
        ])
        judge_scores = np.array([c.get("judge_risk", 2.5) for c in labeled_chains])
        has_judge = any(c.get("judge_risk") is not None for c in labeled_chains)
        X = np.column_stack([sc_scores, judge_scores]) if has_judge else sc_scores[:, None]
        feature_names = ["sc_old", "judge"] if has_judge else ["sc_old"]

        def _rank_norm(v: np.ndarray) -> np.ndarray:
            r = rankdata(v)
            return (r - 1) / (len(r) - 1)

        def _soft_auroc(w: np.ndarray) -> float:
            w_abs = np.abs(w) / (np.abs(w).sum() + 1e-9)
            ranked = np.column_stack([_rank_norm(X[:, i]) for i in range(X.shape[1])])
            scores = ranked @ w_abs
            pos = scores[labels == 1]
            neg = scores[labels == 0]
            if len(pos) == 0 or len(neg) == 0:
                return 0.5
            diffs = pos[:, None] - neg[None, :]
            return float((1.0 / (1.0 + np.exp(-diffs / sigma))).mean())

        def _loss(w: np.ndarray) -> float:
            return -_soft_auroc(w)

        best_w, best_auc = np.ones(X.shape[1]) / X.shape[1], 0.0
        for _ in range(n_restarts):
            x0 = np.abs(np.random.randn(X.shape[1]))
            res = minimize(_loss, x0, method="Nelder-Mead",
                           options={"maxiter": 2000, "xatol": 1e-4, "fatol": 1e-5})
            if -res.fun > best_auc:
                best_auc = -res.fun
                best_w = np.abs(res.x)

        best_w = best_w / best_w.sum()
        return {
            "weights": {name: round(float(w), 4)
                        for name, w in zip(feature_names, best_w)},
            "auroc": round(best_auc, 4),
            "n_chains": len(labeled_chains),
            "note": (
                "Optimistic (train=test). Exp95: judge weight converges to ~0.97. "
                "Cross-domain transfer not guaranteed — use J5 weights instead."
            ),
        }

    def score_batch(self, chains: list[dict]) -> list[LabelFreeResult]:
        """
        Efficiently score a batch of completed chains.

        chains: list of {"question": str, "steps": list, "final_answer": str,
                          "finished": bool}
        """
        return [
            self.score(
                question=c["question"],
                steps=c.get("steps", []),
                final_answer=c.get("final_answer", ""),
                finished=c.get("finished", True),
            )
            for c in chains
        ]

    # ── LLM judge integration (exp89/91/92) ───────────────────────────────────

    _JUDGE_SYSTEM = (
        "You are a meticulous evaluator of AI reasoning chains. "
        "Your job is to assess whether an AI agent's search-and-answer chain is reliable. "
        "Think step-by-step before scoring."
    )
    _JUDGE_PROMPT = (
        "Evaluate this AI agent's reasoning chain for a question-answering task.\n\n"
        "Question: {question}\n\nReasoning chain:\n{chain_text}\n\nFinal answer: {final_answer}\n\n"
        "Instructions — think through each step carefully:\n\n"
        "STEP 1 — DIAGNOSE (think aloud, 2-3 sentences):\n"
        "Identify the single most important strength or weakness of this chain.\n"
        "Does the agent actually use search results to form its answer, or ignore them?\n\n"
        "STEP 2 — SCORE each dimension 1-5:\n"
        "  A) STEP_RELEVANCE: Do search queries directly address sub-questions in the main question?\n"
        "     5=all queries on-target, 3=some tangential queries, 1=queries unrelated to question\n"
        "  B) COHERENCE: Does each step build logically on previous observations?\n"
        "     5=tight logical flow, 3=some disconnected steps, 1=random/contradictory steps\n"
        "  C) CONCLUSION_SUPPORT: Does the final answer follow from the evidence gathered?\n"
        "     5=answer explicitly supported by retrieved facts, 3=partially supported, 1=answer contradicts/ignores evidence\n"
        "  D) ANSWER_SPECIFICITY: Is the final answer specific and complete?\n"
        "     5=specific names/dates/facts, 3=partial answer with hedging, 1=empty/vague\n\n"
        "STEP 3 — FAILURE FLAGS (list any that apply):\n"
        '  - "loop": repeated identical or near-identical search queries\n'
        '  - "no_results": multiple "no relevant results found" observations\n'
        '  - "answer_mismatch": final answer doesn\'t match the question topic\n'
        '  - "abandoned": chain ends mid-reasoning without Finish action\n'
        '  - "hallucination_risk": answer contains specifics not found in any observation\n\n'
        "STEP 4 — OVERALL QUALITY:\n"
        "  GOOD = all scores >= 4, no critical failure flags\n"
        "  BORDERLINE = mixed scores (2-4), minor issues\n"
        "  POOR = any score <= 2, OR has loop/no_results/answer_mismatch flags\n\n"
        'Respond with ONLY this JSON (no other text):\n'
        '{{"diagnosis": "...", "step_relevance": N, "coherence": N, "conclusion_support": N,\n'
        ' "answer_specificity": N, "failure_flags": [...], "quality_label": "GOOD"|"BORDERLINE"|"POOR"}}'
    )

    def _format_chain_for_judge(self, steps: list[dict]) -> str:
        text = ""
        search_count = 0
        for step in steps:
            if step.get("action_type") == "Search":
                search_count += 1
                obs = step.get("observation", "")
                text += (f"\nStep {search_count}:\n"
                         f"  Thought: {step.get('thought','')}\n"
                         f"  Search: {step.get('action_arg','')}\n"
                         f"  Result: {obs[:300]}{'...' if len(obs)>300 else ''}\n")
            elif step.get("action_type") == "Finish":
                text += f"\nConclusion: {step.get('action_arg','')[:200]}\n"
        return text.strip()

    def score_with_judge(
        self,
        question: str,
        steps: list[dict],
        final_answer: str = "",
        finished: bool = True,
        client=None,
        judge_cache_dir: "Path | None" = None,
        judge_model: str = "claude-sonnet-4-6",
        sc_weight: float = 1.0,
        judge_weight: float = 3.0,
    ) -> LabelFreeResult:
        """
        Score a chain using the behavioral SC_OLD ensemble plus a Sonnet LLM judge.

        Implements J5 from exp89: SC_OLD×1 + Sonnet×3 weight blend.
        Sonnet judge alone achieves AUROC=0.7735 within-HP and 0.7411 cross-NQ (exp91).

        Parameters
        ----------
        client         : anthropic.Anthropic instance (required to make new judge calls).
                         If None, only cached judgments are used.
        judge_cache_dir: directory for caching judge JSON responses.
                         Defaults to ~/.qppg/judge_cache/.
        judge_model    : Anthropic model ID (default: "claude-sonnet-4-6").
        sc_weight      : weight of behavioral SC_OLD score in ensemble (default: 1.0).
        judge_weight   : weight of LLM judge score in ensemble (default: 3.0).

        Returns
        -------
        LabelFreeResult with judge_score added to components dict.
        """
        import json, hashlib, re, time
        from pathlib import Path

        # Get behavioral score first
        base_result = self.score(question, steps, final_answer, finished)

        # Set up cache directory
        if judge_cache_dir is None:
            judge_cache_dir = Path("~/.qppg/judge_cache").expanduser()
        else:
            judge_cache_dir = Path(judge_cache_dir)
        judge_cache_dir.mkdir(parents=True, exist_ok=True)

        # Cache key includes judge model version
        cache_key = json.dumps({
            "q":       question,
            "fa":      final_answer,
            "n_steps": len(steps),
            "judge":   f"{judge_model}-v1",
        })
        cache_path = judge_cache_dir / f"{hashlib.md5(cache_key.encode()).hexdigest()}.json"

        judgment = None
        if cache_path.exists():
            try:
                j = json.loads(cache_path.read_text())
                if not j.get("_simulated"):
                    judgment = j
            except (json.JSONDecodeError, KeyError):
                pass

        if judgment is None and client is not None:
            chain_text = self._format_chain_for_judge(steps)
            prompt = self._JUDGE_PROMPT.format(
                question=question,
                chain_text=chain_text,
                final_answer=final_answer,
            )
            try:
                resp = client.messages.create(
                    model=judge_model, max_tokens=400,
                    system=self._JUDGE_SYSTEM,
                    messages=[{"role": "user", "content": prompt}],
                )
                text = resp.content[0].text.strip()
                m = re.search(r'\{.*\}', text, re.DOTALL)
                if m:
                    judgment = json.loads(m.group())
                    judgment["_model"] = judge_model
                    cache_path.write_text(json.dumps(judgment))
            except json.JSONDecodeError:
                pass
            except Exception:
                time.sleep(1)

        # Compute judge risk score (higher = more risky)
        judge_score: float | None = None
        if judgment is not None:
            comp = (judgment.get("step_relevance", 3) +
                    judgment.get("coherence", 3) +
                    judgment.get("conclusion_support", 3) +
                    judgment.get("answer_specificity", 3)) / 4.0
            flags = judgment.get("failure_flags", [])
            penalty = len([f for f in flags
                           if f in ("loop", "no_results", "answer_mismatch", "abandoned")]) * 0.3
            # Normalize: raw score (0-5+) → [0, 1] by dividing by 5
            judge_score = min(1.0, max(0.0, (5.0 - comp + penalty) / 5.0))

        # Blend: SC_OLD × sc_weight + judge × judge_weight
        if judge_score is not None:
            total_w = sc_weight + judge_weight
            blended = (base_result.risk_score * sc_weight + judge_score * judge_weight) / total_w
        else:
            blended = base_result.risk_score  # fallback to behavioral only

        new_components = dict(base_result.components)
        if judge_score is not None:
            new_components["judge_score"] = round(judge_score, 4)
            new_components["judge_quality"] = judgment.get("quality_label", "")

        return LabelFreeResult(
            risk_score=round(blended, 4),
            needs_review=blended >= self.review_threshold,
            behavioral_score=base_result.behavioral_score,
            gmm_score=base_result.gmm_score,
            obs_score=base_result.obs_score,
            qara_score=base_result.qara_score,
            failure_mode=base_result.failure_mode,
            components=new_components,
            calibration_size=base_result.calibration_size,
        )

    def conformal_threshold(
        self,
        calibration_chains: list[dict],
        alpha: float = 0.10,
        use_judge: bool = False,
        client=None,
        judge_cache_dir: "Path | None" = None,
    ) -> float:
        """
        Compute a split-conformal risk threshold with P(false positive) ≤ alpha.

        Based on exp86/92. Call this with a set of CORRECT (non-failing) chains.
        Any new chain with risk_score > returned threshold is flagged as ALERT.

        Coverage Theorem guarantees: P(flag a correct chain) ≤ alpha.

        Parameters
        ----------
        calibration_chains : list of chain dicts that are KNOWN CORRECT (non-failing).
        alpha              : target false-positive rate (default 0.10 = 10% FPR).
        use_judge          : use score_with_judge instead of score (slower, better AUROC).
        client             : anthropic.Anthropic instance (required if use_judge=True).

        Returns
        -------
        float — threshold q̂. Flag chains with risk_score > q̂.

        Example
        -------
        >>> q_hat = scorer.conformal_threshold(correct_chains, alpha=0.10)
        >>> result = scorer.score(q, steps, answer)
        >>> if result.risk_score > q_hat:
        ...     send_alert()
        """
        import numpy as np

        if not calibration_chains:
            raise ValueError("Need at least 1 correct chain for conformal calibration.")

        scores = []
        for c in calibration_chains:
            if use_judge and client is not None:
                r = self.score_with_judge(
                    c.get("question", ""), c.get("steps", []),
                    c.get("final_answer", ""), c.get("finished", True),
                    client=client, judge_cache_dir=judge_cache_dir,
                )
            else:
                r = self.score(
                    c.get("question", ""), c.get("steps", []),
                    c.get("final_answer", ""), c.get("finished", True),
                )
            scores.append(r.risk_score)

        cal_scores = np.sort(scores)
        n = len(cal_scores)
        idx = int(np.ceil((n + 1) * (1 - alpha))) - 1
        idx = min(max(idx, 0), n - 1)
        return float(cal_scores[idx])

    def score_prefix(
        self,
        question: str,
        steps: list[dict],
        t: int,
        final_answer: str = "",
        finished: bool = False,
    ) -> LabelFreeResult:
        """
        Score using only the first *t* steps (streaming / mid-chain scoring).

        Based on exp67: AUROC at t=2 equals or exceeds full-chain AUROC.
        Setting a risk threshold of 0.70 on prefix scoring enables early
        termination that saves ~45% of remaining agent tokens.

        Parameters
        ----------
        question     : str  — original task
        steps        : list — all ReAct steps observed so far (any length)
        t            : int  — number of steps to use (1-indexed; t=0 → behavioral only)
        final_answer : str  — partial answer text if available
        finished     : bool — False for mid-chain scoring (default)

        Returns
        -------
        LabelFreeResult  with risk_score and failure_mode populated.
        """
        prefix = steps[:max(0, t)]
        return self.score(question, prefix, final_answer=final_answer,
                          finished=finished)

    async def ascore(
        self,
        question: str,
        steps: list[dict],
        final_answer: str = "",
        finished: bool = True,
    ) -> LabelFreeResult:
        """
        Async wrapper around score() for use in FastAPI / asyncio event loops.

        Runs score() in a thread-pool executor so it does not block the
        event loop during embedding inference (typically 8–12ms on CPU).

        Usage:
            result = await scorer.ascore(question, steps, final_answer, finished)
        """
        import asyncio
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None, self.score, question, steps, final_answer, finished
        )

    async def ascore_prefix(
        self,
        question: str,
        steps: list[dict],
        t: int,
        final_answer: str = "",
        finished: bool = False,
    ) -> LabelFreeResult:
        """Async wrapper around score_prefix() for streaming use in async agents."""
        import asyncio
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None, self.score_prefix, question, steps, t, final_answer, finished
        )

    def get_gmm_params(self) -> dict | None:
        """
        Export GMM parameters for federated calibration (exp68).

        Returns a JSON-serializable dict containing only GMM means, covariances,
        and weights — NOT the PCA matrix.  In a federated setup each site trains
        its own PCA on local data; only the GMM captures the site's calibration
        distribution.

        Payload size: ~6 KB (2×32 + 4×32 diagonal floats × 3 arrays + weights)
        vs ~40 KB for raw chains — safe to transmit without sharing chain data.

        Returns None if the scorer is not yet calibrated.
        """
        if self._gmm2 is None or self._pca is None:
            return None
        return {
            "gmm2": {
                "means":       self._gmm2.means_.tolist(),
                "covariances": self._gmm2.covariances_.tolist(),   # diag: (k, d)
                "weights":     self._gmm2.weights_.tolist(),
            },
            "gmm4": {
                "means":       self._gmm4.means_.tolist(),
                "covariances": self._gmm4.covariances_.tolist(),
                "weights":     self._gmm4.weights_.tolist(),
            },
            "n_chains":    self._n_cal,
            "embed_model": self.embed_model,
            "pca_dims":    int(self._pca.n_components_),
        }

    def set_gmm_params(self, params: dict) -> bool:
        """
        Import federated-averaged GMM parameters (exp68).

        Replaces the local GMM models with parameters from a federated average
        (e.g., weighted mean of GMM params from multiple sites).  The local PCA
        is retained — each site runs calibrate() first to fit its own PCA, then
        calls set_gmm_params() to upgrade the GMM with federated knowledge.

        Workflow:
          1. Each site: calibrate(local_chains) — fits PCA + GMM locally
          2. Each site: get_gmm_params()        — export GMM (~8 KB, no raw data)
          3. Central:   POST /api/gmm/average   — weighted average across sites
          4. Each site: set_gmm_params(avg)     — replace GMM, keep local PCA

        Parameters
        ----------
        params : dict  — as returned by get_gmm_params()

        Returns True if successful, False if the dict is malformed.
        """
        try:
            import numpy as np
            from sklearn.mixture import GaussianMixture

            def _build_gmm(p: dict, k: int) -> GaussianMixture:
                """Reconstruct a diagonal-covariance GaussianMixture from JSON params."""
                g = GaussianMixture(n_components=k, covariance_type="diag")
                g.means_           = np.array(p["means"], dtype=np.float64)
                g.covariances_     = np.array(p["covariances"], dtype=np.float64)
                g.weights_         = np.array(p["weights"], dtype=np.float64)
                # sklearn: precisions_cholesky_ = 1/sqrt(variance) for diag covariance
                g.precisions_cholesky_ = 1.0 / np.sqrt(g.covariances_)
                g.converged_       = True
                g.n_iter_          = 1
                return g

            g2 = _build_gmm(params["gmm2"], 2)
            g4 = _build_gmm(params["gmm4"], 4)

            # Validate that local PCA dims match the imported GMM dims
            imported_dims = g2.means_.shape[1]
            if self._pca is not None:
                local_dims = int(self._pca.n_components_)
                if local_dims != imported_dims:
                    return False   # caller must re-calibrate with matching n_chains

            self._gmm2  = g2
            self._gmm4  = g4
            self._n_cal = int(params.get("n_chains", 0))

            # If PCA not yet fitted locally, build a pass-through PCA from params dims
            if self._pca is None:
                n_pca_dims = int(params.get("pca_dims", g2.means_.shape[1]))
                from sklearn.decomposition import PCA
                pca = PCA(n_components=n_pca_dims)
                # Identity-like PCA on zero data (score() will use it via transform)
                pca.mean_       = np.zeros(n_pca_dims, dtype=np.float64)
                pca.components_ = np.eye(n_pca_dims, dtype=np.float64)
                pca.explained_variance_         = np.ones(n_pca_dims)
                pca.explained_variance_ratio_   = np.ones(n_pca_dims) / n_pca_dims
                self._pca = pca
                self._cal_chain1 = np.zeros((1, n_pca_dims), dtype=np.float32)
            return True
        except Exception:
            return False

    # ── Convenience ───────────────────────────────────────────────────────────


    def status(self) -> dict:
        return {
            "calibrated":    self.is_calibrated(),
            "n_calibration": self._n_cal,
            "min_needed":    self.min_calibration,
            "embed_model":   self.embed_model,
            "review_threshold": self.review_threshold,
            "components_available": {
                "behavioral":    True,
                "gmm_e3":        self._gmm2 is not None,
                "obs_pool":      self._obs_embs is not None,
                "qara_obs":      self._qara_adapter is not None,
            },
            "n_supervised":  self._n_supervised,
            "expected_auroc": {
                "within_domain": (0.883 if self.is_calibrated() else 0.879),
                "cross_domain":  (0.675 if self._qara_adapter is not None
                                  else 0.664 if self.is_calibrated()
                                  else 0.570),
                "source":        "exp39-42 on HP+NQ benchmark",
            },
        }


# ── Agent Fingerprinter ────────────────────────────────────────────────────────

class AgentFingerprinter:
    """
    Behavioral fingerprinting for cold-start calibration (exp71).

    Clusters agent chains into behavioral archetypes using 10 topic-agnostic
    features (step count, observation length, thought length, etc.).  When a
    new domain has <20 chains (below calibration threshold), use the closest
    behavioral cluster from a source domain as a warm-start calibration pool.

    Exp71 results with K=2 on HotpotQA (n=100):
      - Agglomerative K=2: silhouette=0.614
      - Cluster 0 "efficient":  86 chains, 90% correct, ≤2 search steps
      - Cluster 1 "struggling": 14 chains, 50% correct, 3+ search steps,
                                longer chains, repeated queries

    Usage:
        fp = AgentFingerprinter(n_clusters=2)
        fp.fit(source_chains)                      # train on domain A

        # New domain B has only 8 chains — too few to calibrate
        pool = fp.get_warmstart_pool(new_domain_chains)
        scorer.calibrate(pool)                     # use behavioral match
        print(fp.status())
    """

    _FEATURE_NAMES: list[str] = [
        "n_steps", "n_search", "obs_length_mean", "thought_length_mean",
        "action_arg_length_mean", "answer_length", "has_finish",
        "step_ratio_search", "thought_obs_ratio", "unique_action_ratio",
    ]

    def __init__(self, n_clusters: int = 2):
        self.n_clusters = n_clusters
        self._kmeans  = None
        self._scaler  = None
        self._cluster_pools:  dict[int, list[dict]] = {}
        self._cluster_labels: dict[int, str]        = {}
        self._n_fitted = 0

    @staticmethod
    def extract_features(chain: dict) -> np.ndarray:
        """
        Extract 10 topic-agnostic behavioral features from a chain dict.

        All features are numeric; they encode HOW the agent behaved, not
        WHAT it was asked about.  This is the key insight from exp71:
        behavioral features cluster reliably even across domains.
        """
        steps   = chain.get("steps", [])
        n_steps = max(len(steps), 1)
        n_search = sum(1 for s in steps if s.get("action_type") == "Search")

        obs_lens    = [len(s.get("observation", "")) for s in steps
                       if s.get("observation", "").strip()]
        thought_lens = [len(s.get("thought", "")) for s in steps
                        if s.get("thought", "").strip()]
        arg_lens    = [len(s.get("action_arg", "")) for s in steps
                       if s.get("action_arg", "").strip()]

        obs_mean    = float(np.mean(obs_lens))    if obs_lens    else 0.0
        thought_mean = float(np.mean(thought_lens)) if thought_lens else 0.0
        arg_mean    = float(np.mean(arg_lens))    if arg_lens    else 0.0

        answer_len  = len(chain.get("final_answer", ""))
        has_finish  = float(chain.get("finished", False))

        step_ratio_search = n_search / n_steps
        thought_obs_ratio = thought_mean / max(obs_mean, 1.0)

        action_args = [s.get("action_arg", "").lower().strip() for s in steps]
        unique_ratio = (len(set(action_args)) / len(action_args)
                        if action_args else 1.0)

        return np.array([
            n_steps, n_search, obs_mean, thought_mean, arg_mean,
            answer_len, has_finish, step_ratio_search,
            thought_obs_ratio, unique_ratio,
        ], dtype=np.float32)

    def fit(self, chains: list[dict]) -> dict:
        """
        Fit K-Means clusters on behavioral features.

        Stores each cluster's constituent chains so they can be used as
        warm-start calibration pools for new domains.

        Parameters
        ----------
        chains : list of chain dicts ({"question", "steps", "final_answer", ...})

        Returns summary dict including silhouette score.
        """
        from sklearn.cluster import KMeans
        from sklearn.preprocessing import StandardScaler
        from sklearn.metrics import silhouette_score

        if len(chains) < self.n_clusters * 3:
            return {
                "fitted": False,
                "reason": f"Need ≥{self.n_clusters * 3} chains, got {len(chains)}",
            }

        X_raw = np.array([self.extract_features(c) for c in chains])
        self._scaler = StandardScaler().fit(X_raw)
        X = self._scaler.transform(X_raw)

        km = KMeans(n_clusters=self.n_clusters, n_init=10, random_state=42)
        labels = km.fit_predict(X)
        sil = silhouette_score(X, labels) if len(set(labels)) > 1 else 0.0

        self._kmeans   = km
        self._n_fitted = len(chains)
        self._cluster_pools.clear()
        self._cluster_labels.clear()

        for k in range(self.n_clusters):
            pool_k = [c for c, lbl in zip(chains, labels) if lbl == k]
            self._cluster_pools[k] = pool_k
            # Auto-label: efficient = short chains, struggling = long chains
            avg_steps = float(np.mean([self.extract_features(c)[0] for c in pool_k]))
            self._cluster_labels[k] = "efficient" if avg_steps < 2.5 else "struggling"

        return {
            "fitted":         True,
            "n_chains":       len(chains),
            "n_clusters":     self.n_clusters,
            "silhouette":     round(float(sil), 4),
            "cluster_sizes":  {k: len(self._cluster_pools[k])
                               for k in range(self.n_clusters)},
            "cluster_labels": self._cluster_labels,
        }

    def assign(self, chain: dict) -> int:
        """Assign a single chain to its closest behavioral cluster."""
        if self._kmeans is None:
            raise RuntimeError("Call fit() first.")
        feat = self.extract_features(chain).reshape(1, -1)
        return int(self._kmeans.predict(self._scaler.transform(feat))[0])

    def get_warmstart_pool(self, new_chains: list[dict], n: int = 50) -> list[dict]:
        """
        Return a warm-start calibration pool for a cold-start domain.

        Assigns each chain in new_chains to a cluster (majority vote), then
        returns chains from the matching source-domain cluster (up to n chains).
        This is most valuable when len(new_chains) < min_calibration (default 20).

        Parameters
        ----------
        new_chains : list of chain dicts from the new domain (any size, incl. <20)
        n          : max chains to return from the warm-start pool

        Returns
        -------
        list of chain dicts from the source-domain cluster that best matches
        the new domain's behavioral profile.
        """
        if self._kmeans is None or not new_chains:
            return []
        assignments = [self.assign(c) for c in new_chains]
        best_k = max(set(assignments), key=assignments.count)
        pool = self._cluster_pools.get(best_k, [])
        return pool[:n]

    def status(self) -> dict:
        return {
            "fitted":         self._kmeans is not None,
            "n_fitted":       self._n_fitted,
            "n_clusters":     self.n_clusters,
            "cluster_labels": self._cluster_labels,
            "cluster_sizes":  {k: len(v)
                               for k, v in self._cluster_pools.items()},
        }
