"""
LLMNoveltyDetector — Out-of-distribution detection for deployed LLM agents.
============================================================================

Detects when an incoming chain is structurally dissimilar to the calibration
distribution, signalling that the guard's risk estimate may be unreliable.

Four components (can be used individually or as an ensemble):

  A. Mahalanobis distance (chain features)
     Computes the statistical distance between a chain's BEHAVIORAL feature
     vector and the calibration centroid, weighted by the inverse covariance.
     Normalised to [0,1] via the 95th percentile of calibration distances.
     Requires n_cal ≥ 2 × n_features for reliable covariance estimation.
     SCOPE: measures agent failure-mode novelty (unusual reasoning patterns),
     NOT domain novelty.  Spearman(score, error)=0.37 within-domain.

  B. Conformal p-value (chain features)
     Nonconformity = Mahalanobis distance.  Conformal p-value = proportion
     of calibration points with distance ≥ new point.  Under exchangeability
     (new chain from the same distribution as calibration), p is Uniform[0,1].
     A low p-value (< 0.05) is a statistically valid OOD flag at the 5% level.
     No labels required.  Guarantee holds exactly under exchangeability.

  C. Isolation Forest (chain features, sklearn, optional)
     Partitions the feature space with random hyperplanes.  Points that are
     isolated quickly (shallow average path length) are anomalous.  Provides
     an independent, non-parametric OOD signal.  contamination='auto' avoids
     assuming a known fraction of outliers in the calibration set.

  D. Question-level kNN distance  ← PRIMARY domain novelty signal
     Embeds calibration questions with all-MiniLM-L6-v2, then computes the
     mean cosine distance from the new question to its k nearest calibration
     neighbours (default k=5).  This is the correct signal for domain novelty:
     an HP-calibrated detector will score TriviaQA questions as far from
     HP questions in embedding space, regardless of chain structure.
     Requires: sentence-transformers / all-MiniLM-L6-v2 (used by mini_judge).
     Falls back gracefully when embedder is unavailable.

Routing output
--------------
  ensemble_score > 0.80  →  "force_judge"   (likely OOD; override phase ctrl)
  ensemble_score > 0.50  →  "uncertain"     (may be OOD; flag for monitoring)
  ensemble_score ≤ 0.50  →  "trusted"       (likely in-domain; trust guard)

Honest limitations
------------------
  - Does NOT score whether a chain is correct or incorrect.
  - Components A/B/C measure AGENT FAILURE-MODE novelty (unusual reasoning),
    not domain novelty.  exp_novelty_detector showed these scored TriviaQA
    LOWER than HotpotQA (domain-separation AUROC = 0.42, below random).
    Root cause: TriviaQA is simpler → shorter chains → look more "normal".
  - Component D (question kNN) is the correct domain-novelty signal and is
    the PRIMARY driver of ensemble_score for domain separation.
  - Mahalanobis covariance estimate noisy below n_cal ≈ 30.
  - IsolationForest requires scikit-learn (already a project dependency).
  - kNN embedding requires sentence-transformers; falls back to A/B/C only.

Validated application
---------------------
  Calibrate on HotpotQA-200 (HP, in-domain).
  Score HP-60 test chains (expect low novelty) and TriviaQA-60 chains
  (expect high novelty) → AUROC at domain-separation should be measurable.
  See experiments/exp_novelty_detector.py.

Quick start
-----------
    from llm_guard import LLMNoveltyDetector

    detector = LLMNoveltyDetector()
    detector.fit(cal_chains)        # list of {question, steps, final_answer}

    result = detector.score(new_chain)
    print(result.ensemble_score)    # 0.0 (in-domain) → 1.0 (novel)
    print(result.routing)           # "trusted" | "uncertain" | "force_judge"
    print(result.conformal_pval)    # small → statistically OOD
Purpose:    Detects structurally dissimilar chains using Mahalanobis distance, conformal p-values, embedding novelty, and lexical OOD signals. Signals when the guard's risk estimate may be unreliable.
Exports:    LLMNoveltyDetector, NoveltyResult
Added:      v0.1.0
Status:     Experimental
"""

from __future__ import annotations

import re
import warnings
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np


# ── NoveltyResult ─────────────────────────────────────────────────────────────

@dataclass
class NoveltyResult:
    """
    Per-chain novelty assessment from LLMNoveltyDetector.score().

    maha_score      Mahalanobis distance normalised to [0,1] via the 95th
                    percentile of calibration distances.  Values > 1.0 are
                    possible for strongly OOD chains.
    conformal_pval  Proportion of calibration chains with distance ≥ this
                    chain's distance.  Low value (< 0.05) → statistically OOD.
    iso_score       Isolation Forest anomaly score normalised to [0,1].
                    NaN when sklearn is unavailable.
    ensemble_score  Mean of available component scores, clipped to [0,1].
    routing         "trusted" | "uncertain" | "force_judge"
    n_features      Number of features used (14 with embedder, 9 without).
    n_cal           Number of calibration chains the detector was fitted on.
    """
    maha_score:        float
    conformal_pval:    float
    iso_score:         float           # NaN if IsoForest not available
    question_knn_score: float          # NaN if embedder unavailable (Component D)
    ensemble_score:    float
    routing:           str
    n_features:        int
    n_cal:             int
    _components:       Dict[str, float] = field(default_factory=dict, repr=False)


# ── Feature extraction ────────────────────────────────────────────────────────

def _extract_vec(chain: Dict) -> np.ndarray:
    """
    Extract a numeric feature vector from a chain dict.

    Tries the 14-feature MiniJudge extractor first (includes embedding
    features f12–f14).  Falls back to 9 pure-behavioural features if the
    MiniJudge module or embedder is unavailable.  The fallback features are
    a subset of SC1–SC9 computable with no external dependencies.

    All values are clipped to [0, 2] to prevent extreme outliers from
    dominating the Mahalanobis distance.
    """
    # ── Try 14-feature extractor from mini_judge ──────────────────────────
    try:
        from llm_guard.scoring.mini_judge import _extract_features
        vec = _extract_features(chain)
        return np.clip(vec.astype(float), 0.0, 2.0)
    except Exception:
        pass

    # ── Fallback: 9 pure-behavioural features (no embedder needed) ────────
    steps   = chain.get("steps") or []
    n       = max(len(steps), 1)
    fa      = str(chain.get("final_answer") or "")
    q       = str(chain.get("question") or "")
    actions = [str(s.get("action_type", s.get("action", "")) or "").lower()
               for s in steps]
    obs     = [str(s.get("observation") or "") for s in steps]
    args    = [str(s.get("action_args", s.get("action_arg", s.get("action", ""))) or "")
               for s in steps]
    thoughts = [str(s.get("thought") or "") for s in steps]

    # SC1 — completion flag: 1 = Finish[] reached
    finished    = any(a in ("finish", "finished") for a in actions)
    sc1_comp    = 1.0 if finished else 0.0

    # SC2 — step count (normalised, capped at 6)
    sc2_steps   = min(1.0, n / 6.0)

    # SC3 — fraction of empty observations
    sc3_empty   = sum(1 for o in obs if len(o.strip()) < 15) / n

    # SC5 — repeated search rate (unique queries / total queries)
    search_args = [a for a, t in zip(args, actions)
                   if t in ("search", "lookup", "tool_call")]
    sc5_repeat  = (1.0 - len(set(search_args)) / max(len(search_args), 1)
                   if search_args else 0.0)

    # SC6 — answer-question length gap
    sc6_gap     = max(0.0, 1.0 - len(fa.split()) / max(len(q.split()) * 0.5, 1))

    # SC8 — consecutive action repetitions
    consec      = sum(1 for i in range(1, len(actions)) if actions[i] == actions[i - 1])
    sc8_back    = consec / max(n - 1, 1)

    # SC9 — observation utilisation (longer obs → more information)
    avg_obs_len = np.mean([len(o.split()) for o in obs]) if obs else 0.0
    sc9_util    = min(1.0, avg_obs_len / 100.0)

    # SC10 — thought length variance (coefficient of variation)
    t_lens = [len(t.split()) for t in thoughts] if thoughts else [0]
    sc10_cv = float(np.std(t_lens) / max(np.mean(t_lens), 1.0))

    # SC11 — answer-last-observation token overlap (inverse)
    def _toks(text):
        import re
        sw = {"the","a","an","is","are","was","were","be","have","has","had",
              "do","does","did","will","would","could","should","of","in","on",
              "at","to","for","with","by","and","or","but","not","it","this"}
        return {w for w in re.findall(r"[a-zA-Z]+", text.lower())
                if w not in sw and len(w) > 1}
    last_obs = obs[-1] if obs else ""
    fa_toks  = _toks(fa)
    lo_toks  = _toks(last_obs)
    union    = fa_toks | lo_toks
    sc11_mm  = 1.0 - (len(fa_toks & lo_toks) / max(len(union), 1))

    vec = np.array([sc1_comp, sc2_steps, sc3_empty, sc5_repeat, sc6_gap,
                    sc8_back, sc9_util, sc10_cv, sc11_mm], dtype=float)
    return np.clip(vec, 0.0, 2.0)


# ── LLMNoveltyDetector ────────────────────────────────────────────────────────

class LLMNoveltyDetector:
    """
    Out-of-distribution detector for LLM agent chains.

    Calibrate once on an in-domain chain set, then call score() on each
    new chain to get a novelty score and routing recommendation.

    Parameters
    ----------
    force_judge_threshold : float
        ensemble_score above this → routing="force_judge".  Default 0.80.
    uncertain_threshold : float
        ensemble_score above this (and below force_judge_threshold) →
        routing="uncertain".  Default 0.50.
    use_isolation_forest : bool
        Whether to fit and use an IsolationForest component.  Default True.
    maha_percentile : float
        Calibration-set percentile used to normalise Mahalanobis distances.
        Default 95.
    min_cal_for_maha : int
        Minimum calibration chains required to fit the Mahalanobis component.
        Default 30.
    knn_k : int
        Number of nearest calibration questions to average over in Component D.
        Default 5.  Higher k → smoother but slower.
    knn_percentile : float
        Calibration-set percentile of kNN distances used to normalise Component D
        to [0,1].  Default 95.
    """

    def __init__(
        self,
        force_judge_threshold: float = 0.80,
        uncertain_threshold:   float = 0.50,
        use_isolation_forest:  bool  = True,
        maha_percentile:       float = 95.0,
        min_cal_for_maha:      int   = 30,
        knn_k:                 int   = 5,
        knn_percentile:        float = 95.0,
    ) -> None:
        self._force_threshold     = force_judge_threshold
        self._uncertain_threshold = uncertain_threshold
        self._use_iso             = use_isolation_forest
        self._maha_pct            = maha_percentile
        self._min_cal_maha        = min_cal_for_maha
        self._knn_k               = knn_k
        self._knn_percentile      = knn_percentile

        # Fitted state — chain-feature components (A/B/C)
        self._mu:            Optional[np.ndarray] = None
        self._cov_inv:       Optional[np.ndarray] = None
        self._cal_distances: Optional[np.ndarray] = None
        self._d_normaliser:  float = 1.0
        self._iso_model:     Optional[object] = None
        self._iso_lo:        float = 0.0
        self._iso_hi:        float = 1.0
        # Fitted state — question kNN component (D)
        self._cal_q_embeds:    Optional[np.ndarray] = None  # (n, 384)
        self._knn_normaliser:  float = 1.0                  # 95th-pct kNN dist
        self._embedder:        Optional[object] = None
        self._n_cal:        int   = 0
        self._n_features:   int   = 0
        self._fitted:       bool  = False

    # ── fit ───────────────────────────────────────────────────────────────────

    def fit(self, cal_chains: List[Dict]) -> "LLMNoveltyDetector":
        """
        Fit the novelty detector on in-domain calibration chains.

        Parameters
        ----------
        cal_chains : list of dict
            Each dict: {question, steps, final_answer}.  Labels not needed.
            Recommended n ≥ 50; Mahalanobis component requires n ≥ min_cal_for_maha.

        Returns
        -------
        self
        """
        if not cal_chains:
            raise ValueError("cal_chains must not be empty.")

        # Extract feature matrix
        vecs = []
        for chain in cal_chains:
            try:
                vecs.append(_extract_vec(chain))
            except (ValueError, AttributeError, KeyError, TypeError):
                continue
        if not vecs:
            raise ValueError("Feature extraction failed for all calibration chains.")

        X = np.vstack(vecs)          # (n, p)
        self._n_cal      = len(X)
        self._n_features = X.shape[1]

        # ── Component A: Mahalanobis ──────────────────────────────────────
        if self._n_cal >= self._min_cal_maha:
            self._mu = X.mean(axis=0)
            cov      = np.cov(X.T)
            if cov.ndim == 0:         # single feature (shouldn't happen)
                cov = np.array([[float(cov)]])
            # Regularise: add eps * I so the matrix is always invertible
            eps = max(1e-4, 1e-2 * np.trace(cov) / max(cov.shape[0], 1))
            cov_reg = cov + eps * np.eye(cov.shape[0])
            try:
                self._cov_inv = np.linalg.inv(cov_reg)
            except np.linalg.LinAlgError:
                # Singular even after regularisation — skip Mahalanobis
                self._cov_inv = None

            if self._cov_inv is not None:
                # Compute calibration distances for conformal + normalisation
                dists = np.array([self._maha_dist(x) for x in X])
                self._cal_distances = dists
                self._d_normaliser  = max(
                    float(np.percentile(dists, self._maha_pct)), 1e-6
                )
        else:
            warnings.warn(
                f"LLMNoveltyDetector: only {self._n_cal} calibration chains "
                f"(< min_cal_for_maha={self._min_cal_maha}). "
                "Mahalanobis component disabled; using conformal + IsoForest only.",
                UserWarning,
                stacklevel=2,
            )

        # ── Component C: Isolation Forest ─────────────────────────────────
        if self._use_iso:
            try:
                from sklearn.ensemble import IsolationForest
                iso = IsolationForest(contamination="auto", random_state=42)
                iso.fit(X)
                self._iso_model = iso
                cal_scores = iso.decision_function(X)
                self._iso_lo = float(cal_scores.min())
                self._iso_hi = float(cal_scores.max())
            except ImportError:
                warnings.warn(
                    "LLMNoveltyDetector: scikit-learn not available; "
                    "IsolationForest component disabled.",
                    UserWarning,
                    stacklevel=2,
                )

        # ── Component D: Question-level kNN embedding ─────────────────────
        # This is the PRIMARY domain-novelty signal.  Components A/B/C measure
        # agent failure-mode novelty; Component D measures question-topic novelty.
        questions = [str(c.get("question") or "") for c in cal_chains]
        q_embeds  = self._embed_questions(questions)
        if q_embeds is not None:
            self._cal_q_embeds = q_embeds          # (n, 384)
            # Compute leave-one-out kNN distances for calibration chains
            cal_knn_dists = self._knn_distances(q_embeds, q_embeds,
                                                 k=self._knn_k + 1)  # +1 to exclude self
            # Exclude the closest match (self) by taking [1:] of sorted distances
            # We stored the top-(k+1); drop index 0 (self-match = 0.0)
            cal_knn_dists = cal_knn_dists[:, 1:]   # (n, k)
            cal_mean_dists = cal_knn_dists.mean(axis=1)  # (n,)
            self._knn_normaliser = max(
                float(np.percentile(cal_mean_dists, self._knn_percentile)), 1e-6
            )

        self._fitted = True
        return self

    def _embed_questions(self, questions: List[str]) -> Optional[np.ndarray]:
        """
        Embed a list of questions using all-MiniLM-L6-v2.
        Returns (n, 384) L2-normalised array, or None if embedder unavailable.
        """
        try:
            from sentence_transformers import SentenceTransformer
            if self._embedder is None:
                self._embedder = SentenceTransformer(
                    "sentence-transformers/all-MiniLM-L6-v2"
                )
            embedder = self._embedder
        except ImportError:
            # Try the mini_judge embedder path
            try:
                from llm_guard.scoring.mini_judge import _get_embedder
                embedder = _get_embedder()
            except Exception:
                return None
        except Exception:
            return None

        try:
            embs = embedder.encode(questions, show_progress_bar=False,
                                   normalize_embeddings=True)
            return np.array(embs, dtype=float)
        except Exception:
            return None

    @staticmethod
    def _knn_distances(
        query_embeds: np.ndarray,
        ref_embeds: np.ndarray,
        k: int,
    ) -> np.ndarray:
        """
        Compute cosine distances from each query to its k nearest references.

        Cosine distance = 1 - cosine_similarity (embeddings are L2-normalised,
        so dot product = cosine similarity).

        Returns (n_query, k) array of distances to the k nearest neighbours.
        """
        # Cosine similarity matrix: (n_query, n_ref)
        sim = query_embeds @ ref_embeds.T          # both are L2-normalised
        sim = np.clip(sim, -1.0, 1.0)
        dist = 1.0 - sim                           # cosine distance ∈ [0, 2]

        k = min(k, dist.shape[1])
        # Partial sort: get k smallest distances per row
        idx = np.argpartition(dist, kth=k - 1, axis=1)[:, :k]
        knn_d = np.take_along_axis(dist, idx, axis=1)
        knn_d.sort(axis=1)
        return knn_d                               # (n_query, k)

    # ── score ─────────────────────────────────────────────────────────────────

    def score(self, chain: Dict) -> NoveltyResult:
        """
        Compute novelty scores for a single chain.

        Parameters
        ----------
        chain : dict
            {question, steps, final_answer}

        Returns
        -------
        NoveltyResult
        """
        if not self._fitted:
            raise RuntimeError("Call fit() before score().")

        vec = _extract_vec(chain)
        components: Dict[str, float] = {}

        # ── Component A: Mahalanobis ──────────────────────────────────────
        if self._cov_inv is not None:
            raw_d = self._maha_dist(vec)
            maha_norm = float(np.clip(raw_d / self._d_normaliser, 0.0, 3.0))
            components["maha_raw"]  = float(raw_d)
            components["maha_norm"] = maha_norm
        else:
            maha_norm = float("nan")

        # ── Component B: Conformal p-value ────────────────────────────────
        if self._cal_distances is not None and self._cov_inv is not None:
            raw_d      = components.get("maha_raw", self._maha_dist(vec))
            n_exceed   = int(np.sum(self._cal_distances >= raw_d))
            pval       = (n_exceed + 1) / (self._n_cal + 1)
            # Convert to a novelty score: low p-value → high novelty
            # p=0.05 → novelty=0.95; p=0.50 → novelty=0.50; p=1.0 → novelty=0.0
            conformal_novelty = float(np.clip(1.0 - pval, 0.0, 1.0))
            components["conformal_pval"]    = float(pval)
            components["conformal_novelty"] = conformal_novelty
        else:
            pval              = float("nan")
            conformal_novelty = float("nan")

        # ── Component C: Isolation Forest ─────────────────────────────────
        if self._iso_model is not None:
            dec   = float(self._iso_model.decision_function(vec.reshape(1, -1))[0])
            span  = max(self._iso_hi - self._iso_lo, 1e-6)
            iso_n = float(np.clip((self._iso_hi - dec) / span, 0.0, 1.0))
            components["iso_raw"]  = dec
            components["iso_norm"] = iso_n
        else:
            iso_n = float("nan")

        # ── Component D: Question kNN distance ────────────────────────────
        if self._cal_q_embeds is not None:
            q_text = str(chain.get("question") or "")
            q_emb  = self._embed_questions([q_text])
            if q_emb is not None:
                knn_d = self._knn_distances(q_emb, self._cal_q_embeds,
                                             k=self._knn_k)   # (1, k)
                mean_d = float(knn_d.mean())
                knn_n  = float(np.clip(mean_d / self._knn_normaliser, 0.0, 2.0))
                components["knn_raw"]  = mean_d
                components["knn_norm"] = knn_n
            else:
                knn_n = float("nan")
        else:
            knn_n = float("nan")

        # ── Ensemble ──────────────────────────────────────────────────────
        # Component D (question kNN) has double weight: it is the primary
        # domain-novelty signal validated to separate HP from TV/2Wiki.
        # Components A/B/C have single weight: they measure failure-mode novelty.
        weighted: List[tuple] = []
        for score, weight in [
            (maha_norm,        1.0),
            (conformal_novelty, 1.0),
            (iso_n,            1.0),
            (knn_n,            2.0),   # double weight for question novelty
        ]:
            if not np.isnan(score):
                weighted.append((float(np.clip(score, 0.0, 1.0)), weight))

        if weighted:
            total_w = sum(w for _, w in weighted)
            ensemble = float(sum(s * w for s, w in weighted) / total_w)
        else:
            ensemble = 0.5

        # ── Routing ───────────────────────────────────────────────────────
        if ensemble > self._force_threshold:
            routing = "force_judge"
        elif ensemble > self._uncertain_threshold:
            routing = "uncertain"
        else:
            routing = "trusted"

        return NoveltyResult(
            maha_score         = maha_norm if not np.isnan(maha_norm) else float("nan"),
            conformal_pval     = float(pval) if not np.isnan(pval) else float("nan"),
            iso_score          = iso_n if not np.isnan(iso_n) else float("nan"),
            question_knn_score = knn_n if not np.isnan(knn_n) else float("nan"),
            ensemble_score     = ensemble,
            routing            = routing,
            n_features         = self._n_features,
            n_cal              = self._n_cal,
            _components        = components,
        )

    def score_batch(self, chains: List[Dict]) -> List[NoveltyResult]:
        """Score a list of chains. Returns results in the same order."""
        return [self.score(c) for c in chains]

    # ── internals ─────────────────────────────────────────────────────────────

    def _maha_dist(self, x: np.ndarray) -> float:
        """Mahalanobis distance from x to the calibration centroid."""
        delta = x - self._mu
        return float(np.sqrt(np.clip(delta @ self._cov_inv @ delta, 0.0, None)))

    # ── diagnostics ───────────────────────────────────────────────────────────

    def knn_enabled(self) -> bool:
        return self._cal_q_embeds is not None

    def calibration_summary(self) -> Dict:
        """
        Summary statistics of the calibration distribution.
        Useful for verifying fit quality before deployment.
        """
        if not self._fitted:
            return {"fitted": False}
        summary: Dict = {
            "fitted":           True,
            "n_cal":            self._n_cal,
            "n_features":       self._n_features,
            "maha_enabled":     self._cov_inv is not None,
            "iso_enabled":      self._iso_model is not None,
        }
        if self._cal_distances is not None:
            summary["cal_maha_p50"] = float(np.percentile(self._cal_distances, 50))
            summary["cal_maha_p95"] = float(np.percentile(self._cal_distances, 95))
            summary["cal_maha_p99"] = float(np.percentile(self._cal_distances, 99))
            summary["d_normaliser"] = self._d_normaliser
        return summary


# ── QuestionComplexityScorer ──────────────────────────────────────────────────

def _complexity_features(question: str) -> np.ndarray:
    """
    Extract 7 structural/linguistic complexity features from a question string.
    These are calibrated to separate multi-hop (HP-style) from single-hop
    (TriviaQA-style) questions.  No LLM or embedder needed.

    Features
    --------
    0  token_count_norm     token count / 20  (HP ≈ 0.9, TV ≈ 0.5)
    1  clause_count_norm    subordinate clauses (which/that/who/where) / 3
    2  multihop_cue         any of: after|before|when|while|also|both / 1
    3  wh_position_norm     position of first WH-word / token_count
                           (high = mid-sentence question = multi-hop)
    4  ne_count_norm        runs of Capitalised words (proxy NE count) / 4
    5  comma_count_norm     comma count / 3
    6  char_len_norm        char length / 120
    """
    q = question.strip()
    tokens = q.split()
    n = max(len(tokens), 1)

    # 0 token count
    tok_norm = n / 20.0

    # 1 subordinate clause markers
    clause_words = re.findall(
        r'\b(which|that|who|where|whose|whom)\b', q, re.IGNORECASE
    )
    clause_norm = len(clause_words) / 3.0

    # 2 multi-hop cue words
    multihop = int(bool(re.search(
        r'\b(after|before|when|while|also|both|same|another|other)\b',
        q, re.IGNORECASE
    )))

    # 3 WH-word position (normalised index in token list)
    wh_pat = re.compile(r'^(what|which|who|where|when|how|why)$', re.IGNORECASE)
    wh_indices = [i for i, t in enumerate(tokens) if wh_pat.match(t)]
    if wh_indices:
        wh_pos_norm = wh_indices[0] / n      # 0 = start, 1 = end
    else:
        wh_pos_norm = 0.0

    # 4 named-entity proxy: runs of consecutive capitalised words (not at sentence start)
    ne_runs = re.findall(r'(?<!\A)(?<!\. )(?:[A-Z][a-z]+(?:\s[A-Z][a-z]+)+)', q)
    ne_norm = len(ne_runs) / 4.0

    # 5 comma count
    comma_norm = q.count(',') / 3.0

    # 6 char length
    char_norm = len(q) / 120.0

    vec = np.array([tok_norm, clause_norm, multihop, wh_pos_norm,
                    ne_norm, comma_norm, char_norm], dtype=float)
    return np.clip(vec, 0.0, 3.0)


@dataclass
class ComplexityResult:
    """Result from QuestionComplexityScorer.score()."""
    complexity_score: float    # 0 = simple (TV-like), 1 = complex (HP-like)
    features: Dict[str, float]
    n_cal: int
    routing: str               # "trusted" | "uncertain" | "force_judge"


class QuestionComplexityScorer:
    """
    Domain novelty detector based on QUESTION STRUCTURAL COMPLEXITY.

    Unlike sentence-embedding kNN (which encodes topic), this scorer uses
    7 syntax/structure features to separate multi-hop (HotpotQA-style)
    from single-hop (TriviaQA-style) questions WITHOUT any embedder.

    Key idea: HP questions are structurally complex (longer, more clauses,
    WH-word mid-sentence, multi-hop cues).  TV questions are structurally
    simple (short, direct, WH-word at start).

    Calibrate on in-domain (HP) chains, then score new questions.  A high
    complexity_score means the question matches the HP distribution;
    a low score means it is likely simpler (TV-like) → potentially OOD
    for an HP-calibrated guard.

    Parameters
    ----------
    ood_low_threshold : float
        complexity_score below this → OOD (domain is SIMPLER than calibration).
        Default 0.30.
    ood_high_threshold : float
        complexity_score above this → OOD (domain is MORE COMPLEX than calibration).
        Default 0.90.

    Usage
    -----
        scorer = QuestionComplexityScorer()
        scorer.fit(hp_calibration_chains)
        result = scorer.score(new_chain)
        print(result.complexity_score, result.routing)
    """

    FEATURE_NAMES = [
        "token_count_norm", "clause_count_norm", "multihop_cue",
        "wh_position_norm", "ne_count_norm", "comma_count_norm", "char_len_norm"
    ]

    def __init__(
        self,
        ood_low_threshold: float = 0.30,
        ood_high_threshold: float = 0.90,
    ) -> None:
        self._ood_low  = ood_low_threshold
        self._ood_high = ood_high_threshold
        self._mu:   Optional[np.ndarray] = None
        self._std:  Optional[np.ndarray] = None
        self._w:    Optional[np.ndarray] = None   # LogReg weights (fitted)
        self._b:    float = 0.0
        self._n_cal: int = 0
        self._fitted: bool = False

    def fit(self, cal_chains: List[Dict]) -> "QuestionComplexityScorer":
        """
        Fit on in-domain calibration chains (labels NOT needed).

        Stores the mean/std of the 7 features for z-score normalisation,
        then fits a single-neuron logistic projection so score() returns
        a calibrated probability rather than a raw sum.

        For unsupervised fitting (no OOD examples), the 'score' is
        simply the sigmoid of the weighted sum: heavier features are
        those with higher variance in the calibration set (i.e. most
        discriminative within HP itself).
        """
        vecs = np.vstack([_complexity_features(str(c.get("question") or ""))
                          for c in cal_chains])
        self._n_cal = len(vecs)
        self._mu    = vecs.mean(axis=0)
        self._std   = np.maximum(vecs.std(axis=0), 1e-6)
        # Unsupervised weight: normalise and use variance-proportional weights
        z = (vecs - self._mu) / self._std
        # Weight each feature by its stddev in calibration set (higher std = more signal)
        self._w = self._std / (self._std.sum() + 1e-9)
        self._b = 0.0
        self._fitted = True
        return self

    def fit_with_ood(
        self,
        in_domain_chains: List[Dict],
        ood_chains: List[Dict],
    ) -> "QuestionComplexityScorer":
        """
        Supervised fit: logistic regression on in-domain (y=1) vs OOD (y=0).
        Requires scikit-learn.  Falls back to unsupervised fit if unavailable.
        """
        try:
            from sklearn.linear_model import LogisticRegression
            from sklearn.preprocessing import StandardScaler
        except ImportError:
            warnings.warn("scikit-learn not found; falling back to unsupervised fit.",
                          UserWarning)
            return self.fit(in_domain_chains)

        in_vecs = np.vstack([_complexity_features(str(c.get("question") or ""))
                              for c in in_domain_chains])
        ood_vecs = np.vstack([_complexity_features(str(c.get("question") or ""))
                               for c in ood_chains])
        X = np.vstack([in_vecs, ood_vecs])
        y = np.array([1] * len(in_vecs) + [0] * len(ood_vecs))

        scaler = StandardScaler()
        X_z = scaler.fit_transform(X)
        lr = LogisticRegression(C=1.0, max_iter=200, random_state=42)
        lr.fit(X_z, y)

        self._mu   = scaler.mean_
        self._std  = scaler.scale_
        self._w    = lr.coef_[0]
        self._b    = float(lr.intercept_[0])
        self._n_cal = len(in_domain_chains)
        self._fitted = True
        return self

    def _sigmoid(self, x: float) -> float:
        return 1.0 / (1.0 + float(np.exp(-x)))

    def score(self, chain: Dict) -> ComplexityResult:
        if not self._fitted:
            raise RuntimeError("Call fit() before score().")
        q   = str(chain.get("question") or "")
        vec = _complexity_features(q)
        z   = (vec - self._mu) / self._std
        raw = float(z @ self._w) + self._b
        score = self._sigmoid(raw)

        if score < self._ood_low or score > self._ood_high:
            routing = "force_judge"
        elif 0.30 <= score <= 0.70:
            routing = "uncertain"
        else:
            routing = "trusted"

        return ComplexityResult(
            complexity_score=round(score, 4),
            features=dict(zip(self.FEATURE_NAMES, vec.tolist())),
            n_cal=self._n_cal,
            routing=routing,
        )

    def score_batch(self, chains: List[Dict]) -> List[ComplexityResult]:
        return [self.score(c) for c in chains]
