"""
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):

  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

  Density-based signals — requires ≥5 unlabeled calibration chains:
    GMM-E3   GMM k=2-neg + GMM k=4 rank-ensemble on PCA-32 of chain-1 embeddings
             AUROC 0.680 within-HP, 0.643 cross-NQ  (exp38)
    Obs-pool Mean-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)

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

Calibration efficiency (exp44):
  n=0  chains: SC2 alone, AUROC 0.879 / 0.570 — deploy immediately
  n≥5  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

    # 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 _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)


# ── 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)."""

    components: dict = field(default_factory=dict)
    """Individual scores: sc1, sc2, sc3, sc5, sc6 and polarity-corrected values."""

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

    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),
            "calibration_size":  self.calibration_size,
            "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 20.
    """

    def __init__(
        self,
        embed_model: str = "all-MiniLM-L6-v2",
        review_threshold: float = 0.65,
        min_calibration: int = 20,
    ):
        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:
            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()

        # ── Chain-1 embeddings (for GMM-E3) ─────────────────────────────────
        chain1_texts = [self._fmt_chain1(c["question"], c.get("steps", []))
                        for c in chains]
        chain1_embs = embedder.encode(chain1_texts, normalize_embeddings=True,
                                       show_progress_bar=False, batch_size=256)
        self._cal_chain1 = chain1_embs

        # PCA-32 + GMM fit
        n_comp = min(32, len(chains) - 1, chain1_embs.shape[1])
        self._pca = PCA(n_components=n_comp, random_state=42).fit(chain1_embs)
        cal_pca = self._pca.transform(chain1_embs)
        self._gmm2 = GaussianMixture(2, n_init=5, random_state=42).fit(cal_pca)
        self._gmm4 = GaussianMixture(4, n_init=5, 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)
            # Mean-pool per chain
            obs_pool = []
            offset = 0
            for cnt in obs_counts:
                if cnt > 0:
                    obs_pool.append(all_obs_embs[offset:offset + cnt].mean(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,
            "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)

        components.update(sc1=sc1, sc2=sc2, sc3=sc3, sc5=sc5, sc6=sc6)

        # SC1, SC2, SC3, SC5, SC6 are all positive polarity (higher = riskier)
        # We can't rank-normalise a single sample, so we return raw values
        # and the combined behavioral_score is their simple mean (0→1 normalised).
        # To put on comparable scales:
        #   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
        sc2_norm = min(1.0, sc2 / 6.0)
        sc3_norm = min(1.0, sc3 * 100.0)
        sc5_norm = min(1.0, sc5 / 2.0)
        behavioral_score = float(np.mean([sc1, sc2_norm, sc3_norm, sc5_norm, sc6]))
        components["behavioral_score"] = behavioral_score

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

        if self.is_calibrated():
            ctx1 = self._fmt_chain1(question, steps)
            emb1 = 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(emb1[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.mean(axis=0)
                    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

        # ── 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))

        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),
            components=components,
            calibration_size=self._n_cal,
        )

    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 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
        ]

    # ── 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",
            },
        }
