"""
QPPGService — core confidence scoring and calibration management.

Implements the full QPPG deployment lifecycle:
  COLD_START  (<10 verified correct chains)  — all traffic routed to human review
  WARMING     (10–49 correct chains)         — selective routing, improving AUROC
  DEPLOYED    (≥50 correct chains)           — automated with targeted review

Calibration pool grows organically: every verified correct chain automatically
becomes a free calibration example (online calibration, exp30).

Diversity-maximised sampling (greedy max-distance) reaches AUROC>0.90 at
n=50 correct chains — 33% faster than random insertion (exp30 result).
"""

from __future__ import annotations

import json
import pickle
import time
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

warnings.filterwarnings("ignore")

# ── QARA adapter (Quality-Aware Reasoning Adapter) ───────────────────────────
# Lightweight 2-layer MLP (384→256→64, L2-norm) trained with Supervised
# Contrastive loss using cross-domain positive pairs (exp31–35).
# All weights stored as numpy arrays for portable CPU inference (no torch
# required at runtime; torch is only needed during fit_qara training).


@dataclass
class _QARAWeights:
    """Serialisable numpy weights for the QARA adapter."""
    w1:   np.ndarray   # (384, 256) — Linear 1
    b1:   np.ndarray   # (256,)
    ln_w: np.ndarray   # (256,)     — LayerNorm scale
    ln_b: np.ndarray   # (256,)     — LayerNorm bias
    w2:   np.ndarray   # (256, 64)  — Linear 2
    b2:   np.ndarray   # (64,)


def _supcon_loss_torch(z, labels_correct, temp: float = 0.1):
    """
    Supervised Contrastive loss with cross-domain positive pairs.

    Positive pairs = any two CORRECT chains, regardless of domain.
    Incorrect chains are negatives relative to all correct chains.

    labels_correct : 1-D float tensor, 1.0 = correct, 0.0 = incorrect.
    """
    import torch
    # Outer product: 1 only for correct-correct pairs
    pos_mask = labels_correct.unsqueeze(1) * labels_correct.unsqueeze(0)
    pos_mask.fill_diagonal_(0)
    sim = (z @ z.T).clamp(-20, 20) / temp   # clamp prevents overflow
    sim_masked = sim.clone()
    sim_masked.fill_diagonal_(float("-inf"))
    log_sum   = torch.logsumexp(sim_masked, dim=1, keepdim=True)
    log_probs = sim_masked - log_sum
    # Zero out non-positive positions before summing to avoid 0 × (−inf) = NaN
    log_probs_pos = log_probs.clone()
    log_probs_pos[pos_mask == 0] = 0.0
    n_pos = pos_mask.sum(dim=1).clamp(min=1)
    return (-log_probs_pos.sum(dim=1) / n_pos).mean()


def _train_qara_weights(
    embs:   np.ndarray,
    labels: np.ndarray,
    epochs: int   = 200,
    lr:     float = 3e-4,
    temp:   float = 0.1,
    verbose: bool = True,
) -> _QARAWeights:
    """
    Train a QARA adapter from embeddings + binary correctness labels.
    Returns portable numpy weights (no torch needed after this call).
    """
    try:
        import torch
        import torch.nn as nn
    except ImportError as exc:
        raise ImportError("QARA training requires PyTorch: pip install torch") from exc

    class _Adapter(nn.Module):
        def __init__(self):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(384, 256), nn.ReLU(), nn.LayerNorm(256),
                nn.Linear(256, 64),
            )
        def forward(self, x):
            z = self.net(x)
            return z / (z.norm(dim=-1, keepdim=True) + 1e-8)

    ada   = _Adapter()
    opt   = torch.optim.Adam(ada.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)

    x_t = torch.tensor(embs,  dtype=torch.float32)
    y_t = torch.tensor(labels, dtype=torch.float32)   # 1.0=correct, 0.0=incorrect

    ada.train()
    for ep in range(1, epochs + 1):
        opt.zero_grad()
        loss = _supcon_loss_torch(ada(x_t), y_t, temp)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(ada.parameters(), max_norm=1.0)
        opt.step()
        sched.step()
        if verbose and ep % 50 == 0:
            print(f"      epoch {ep:3d}/{epochs}  loss={loss.item():.4f}")

    ada.eval()
    with torch.no_grad():
        sd = ada.net.state_dict()
    # Transpose Linear weights: torch shape (out, in) → numpy (in, out) for emb @ W
    return _QARAWeights(
        w1   = sd["0.weight"].numpy().T,
        b1   = sd["0.bias"].numpy(),
        ln_w = sd["2.weight"].numpy(),
        ln_b = sd["2.bias"].numpy(),
        w2   = sd["3.weight"].numpy().T,
        b2   = sd["3.bias"].numpy(),
    )


# ── Calibration growth curve from exp30 (diversity-maximised ordering) ────────
# Maps n_correct_chains → estimated AUROC on within-domain queries
_GROWTH_CURVE: list[tuple[int, float]] = [
    (0,  0.50), (1,  0.549), (2,  0.554), (3,  0.562),
    (5,  0.583), (7,  0.596), (10, 0.631), (15, 0.672),
    (20, 0.714), (30, 0.802), (50, 0.911), (75, 0.995),
    (85, 1.000),
]

COLD_START_MIN  = 10    # below this: all traffic → review
WARMING_MIN     = 50    # below this: selective routing
DEPLOYED_TARGET = 0.90  # AUROC threshold for "deployed" status


def _estimate_auroc(n: int) -> float:
    """Interpolate estimated AUROC from the exp30 growth curve."""
    xs = [p[0] for p in _GROWTH_CURVE]
    ys = [p[1] for p in _GROWTH_CURVE]
    if n <= 0:
        return 0.50
    if n >= xs[-1]:
        return 1.0
    for i in range(len(xs) - 1):
        if xs[i] <= n <= xs[i + 1]:
            t = (n - xs[i]) / (xs[i + 1] - xs[i])
            return ys[i] + t * (ys[i + 1] - ys[i])
    return 0.50


def _chains_to_deployment(n_current: int, qpd: int = 50,
                           correct_rate: float = 0.50) -> float:
    """Estimate days until WARMING_MIN correct chains at given query rate."""
    needed = max(0, WARMING_MIN - n_current)
    if needed == 0:
        return 0.0
    return (needed / correct_rate) / qpd


# ── Chain context formatting ──────────────────────────────────────────────────

def _fmt_chain1(question: str, steps: list[dict], raw_text: str = "") -> str:
    """
    Build a chain-1 embedding context string.

    If `steps` is provided (ReAct format), uses the first step.
    If `raw_text` is provided, uses that directly.
    Falls back to just the question.
    """
    if raw_text:
        return f"TASK: {question[:300]}\nCURRENT: {raw_text[:400]}"

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


# ── Diversity selection ───────────────────────────────────────────────────────

def _most_diverse_index(pool_embs: np.ndarray, candidate_emb: np.ndarray) -> float:
    """
    Return the minimum distance from candidate_emb to the existing pool.
    Used to decide whether a new chain adds diversity worth keeping.
    """
    if len(pool_embs) == 0:
        return float("inf")
    dists = np.linalg.norm(pool_embs - candidate_emb[None, :], axis=1)
    return float(dists.min())


# ── Core result dataclass ─────────────────────────────────────────────────────

@dataclass
class ScoreResult:
    confidence: float | None
    """Confidence that this chain is CORRECT. Range [0, 1]. None if cold-start."""

    needs_review: bool
    """True if confidence is below threshold OR we're in cold-start."""

    raw_distance: float | None
    """Raw KNN distance to nearest calibration chain. None if cold-start."""

    calibration_size: int
    """Number of verified correct chains in the calibration pool."""

    estimated_auroc: float
    """Estimated AUROC based on calibration pool size (from exp30 growth curve)."""

    deployment_status: str
    """One of: 'COLD_START', 'WARMING', 'DEPLOYED'."""

    chains_to_go: int
    """Verified correct chains still needed to reach DEPLOYED threshold."""

    def to_dict(self) -> dict:
        return {
            "confidence":       self.confidence,
            "needs_review":     self.needs_review,
            "raw_distance":     self.raw_distance,
            "calibration_size": self.calibration_size,
            "estimated_auroc":  round(self.estimated_auroc, 4),
            "deployment_status": self.deployment_status,
            "chains_to_go":     self.chains_to_go,
        }


@dataclass
class CalibrateResult:
    added: bool
    """True if the chain was added to the calibration pool."""

    reason: str
    """Why the chain was/wasn't added."""

    calibration_size: int
    estimated_auroc: float
    deployment_status: str
    chains_to_go: int

    def to_dict(self) -> dict:
        return {
            "added":             self.added,
            "reason":            self.reason,
            "calibration_size":  self.calibration_size,
            "estimated_auroc":   round(self.estimated_auroc, 4),
            "deployment_status": self.deployment_status,
            "chains_to_go":      self.chains_to_go,
        }


# ── Main service class ────────────────────────────────────────────────────────

class QPPGService:
    """
    QPPG confidence scoring service for LLM agent deployment.

    Lifecycle:
        COLD_START  → all queries flagged for human review (calibration growing)
        WARMING     → selective review based on QPPG confidence score
        DEPLOYED    → automated answers with targeted review for uncertain chains

    Parameters
    ----------
    domain_name : str
        Identifier for this deployment domain. Used for save file naming.
    save_dir : str | Path
        Directory to persist the calibration pool. Defaults to ~/.qppg/
    embed_model : str
        SentenceTransformer model name. Must be already downloaded.
    review_threshold : float
        Confidence below this → needs_review=True. Default 0.55.
        Confidence is ECDF-based: fraction of calibration pool whose intra-pool
        LOO distance exceeds the test chain's KNN distance. So 0.55 means
        "auto-answer only if this chain is closer to the pool than 55% of
        calibration chains are to their own nearest neighbour".
    min_diversity_distance : float
        Min distance a new chain must be from existing pool to be added
        (diversity-maximised insertion). Set to 0.0 to disable.
    queries_per_day : int
        Assumed query rate, used for deployment time estimates.
    correct_rate : float
        Assumed fraction of queries that are correct (for time estimates).
    """

    def __init__(
        self,
        domain_name: str = "default",
        save_dir: str | Path | None = None,
        embed_model: str = "all-MiniLM-L6-v2",
        review_threshold: float = 0.75,
        min_diversity_distance: float = 0.01,
        queries_per_day: int = 50,
        correct_rate: float = 0.50,
    ):
        self.domain_name            = domain_name
        self.embed_model            = embed_model
        self.review_threshold       = review_threshold
        self.min_diversity_distance = min_diversity_distance
        self.queries_per_day        = queries_per_day
        self.correct_rate           = correct_rate

        # Save path
        if save_dir is None:
            save_dir = Path.home() / ".qppg"
        self._save_path = Path(save_dir) / f"{domain_name}.pkl"
        self._save_path.parent.mkdir(parents=True, exist_ok=True)

        # Calibration state
        self._cal_embs: list[np.ndarray] = []   # correct chain embeddings
        self._cal_meta: list[dict]       = []   # metadata per chain
        self._cal_distances: list[float] = []   # KNN dist of each chain at admission
        self._n_scored  = 0
        self._n_reviewed = 0
        # Cache of intra-pool LOO distances (invalidated when pool changes)
        self._loo_dists_cache: np.ndarray | None = None
        # Running buffer of observed KNN distances (last 200 DEPLOYED queries)
        self._score_buffer: list[float] = []
        self._BUFFER_SIZE = 200

        # QARA adapter (optional; fitted via fit_qara())
        self._qara: _QARAWeights | None = None

        # Lazy-load embedder
        self._embedder = None

        # Load persisted state if it exists
        if self._save_path.exists():
            self._load()
            print(f"[QPPGService] Loaded '{domain_name}' "
                  f"({len(self._cal_embs)} calibration chains)")
        else:
            print(f"[QPPGService] New domain '{domain_name}' — starting cold start")

    # ── Embedder (lazy) ───────────────────────────────────────────────────────

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

    def _embed(self, text: str) -> np.ndarray:
        emb = self._get_embedder().encode(
            [text], normalize_embeddings=True, show_progress_bar=False
        )
        return emb[0]

    # ── Status helpers ────────────────────────────────────────────────────────

    @property
    def n_calibration(self) -> int:
        return len(self._cal_embs)

    def _deployment_status(self) -> str:
        n = self.n_calibration
        if n < COLD_START_MIN:
            return "COLD_START"
        elif n < WARMING_MIN:
            return "WARMING"
        else:
            return "DEPLOYED"

    def _chains_to_go(self) -> int:
        return max(0, WARMING_MIN - self.n_calibration)

    # ── QARA adapter inference (numpy, no torch required) ────────────────────

    def _adapt_embs(self, embs: np.ndarray) -> np.ndarray:
        """
        Apply the QARA adapter to a batch of raw MiniLM embeddings.
        Returns adapted 64-d embeddings if adapter is fitted, else raw 384-d.

        Forward pass:  Linear(384→256) → ReLU → LayerNorm → Linear(256→64) → L2-norm
        """
        if self._qara is None or len(embs) == 0:
            return embs
        q = self._qara
        x = embs @ q.w1 + q.b1                              # (n, 256)
        x = np.maximum(x, 0)                                 # ReLU
        mu  = x.mean(axis=1, keepdims=True)
        var = x.var(axis=1, keepdims=True)
        x   = (x - mu) / np.sqrt(var + 1e-5) * q.ln_w + q.ln_b  # LayerNorm
        x   = x @ q.w2 + q.b2                               # (n, 64)
        nrm = np.linalg.norm(x, axis=1, keepdims=True)
        return x / (nrm + 1e-8)                             # L2-norm

    # ── Confidence calculation ────────────────────────────────────────────────

    def _loo_distances(self) -> np.ndarray:
        """
        Leave-one-out KNN distances for each calibration chain.
        Cached and invalidated whenever the pool changes.

        loo_dist[i] = distance from cal_emb[i] to its nearest neighbour
                      in the pool (excluding itself).

        These represent the typical within-pool distances for "correct" chains.
        """
        if self._loo_dists_cache is not None:
            return self._loo_dists_cache
        n = len(self._cal_embs)
        if n < 2:
            self._loo_dists_cache = np.array([1.0] * n)
            return self._loo_dists_cache
        cal_arr = self._adapt_embs(np.array(self._cal_embs))  # 64-d if QARA fitted
        loo = np.zeros(n)
        for i in range(n):
            others = np.delete(cal_arr, i, axis=0)
            loo[i] = np.linalg.norm(others - cal_arr[i], axis=1).min()
        self._loo_dists_cache = loo
        return loo

    def _cal_dist_confidence(self, knn_dist: float) -> float | None:
        """
        Calibration-admission distance confidence.

        Reference distribution = KNN distances at which known-correct chains
        were admitted to the pool. This directly answers:
          "Is this test chain as close to the pool as typical correct chains?"

        confidence = fraction of admission distances LARGER than knn_dist
                   = "this test chain is at least as close to the pool as X%
                     of the correct chains we have seen."

        Unlike LOO distances (inflated by diversity-maximised selection),
        admission distances are on the same scale as test queries: both measure
        distance to the pool with similar pool sizes. When pool ≥ COLD_START_MIN
        the distances stabilise, giving clean discrimination.

        Returns None when fewer than 10 admission distances are recorded
        (falls back to LOO-ECDF).
        """
        if len(self._cal_distances) < 10:
            return None
        rank = sum(1 for d in self._cal_distances if d <= knn_dist)
        return 1.0 - rank / len(self._cal_distances)

    def _ecdf_confidence(self, knn_dist: float) -> float:
        """
        LOO-ECDF fallback confidence (used only when buffer is too small).

        confidence = fraction of calibration chains whose LOO KNN distance
        exceeds knn_dist (with p90 clipping to reduce diversity-pool inflation).
        """
        n = len(self._cal_embs)
        if n < 2:
            return 0.5
        loo = self._loo_distances()
        p90 = float(np.percentile(loo, 90))
        loo_clipped = np.minimum(loo, p90)
        return float((loo_clipped >= knn_dist).mean())

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

    def score(
        self,
        question: str,
        steps: list[dict] | None = None,
        raw_text: str = "",
    ) -> ScoreResult:
        """
        Score a reasoning chain and return a confidence result.

        Parameters
        ----------
        question : str
            The task/question being answered.
        steps : list[dict], optional
            ReAct-format steps. Each dict may have keys:
            'thought', 'action_type', 'action_arg', 'observation'.
        raw_text : str, optional
            Free-form chain text (used if steps is None/empty).

        Returns
        -------
        ScoreResult
        """
        self._n_scored += 1
        n   = self.n_calibration
        est = _estimate_auroc(n)
        status = self._deployment_status()
        ctg    = self._chains_to_go()

        # Cold start: no calibration data → always review
        if n == 0:
            self._n_reviewed += 1
            return ScoreResult(
                confidence=None, needs_review=True, raw_distance=None,
                calibration_size=0, estimated_auroc=est,
                deployment_status=status, chains_to_go=ctg,
            )

        # Embed the chain; apply QARA adapter if fitted
        ctx = _fmt_chain1(question, steps or [], raw_text)
        emb   = self._embed(ctx)
        emb_q = self._adapt_embs(emb[None, :])[0]    # 64-d if QARA, else 384-d

        # KNN distance to calibration pool (in adapter space if QARA fitted)
        cal_arr  = self._adapt_embs(np.array(self._cal_embs))
        dists    = np.linalg.norm(cal_arr - emb_q[None, :], axis=1)
        knn_dist = float(dists.min())

        # Calibration-admission confidence: compares test distance to the
        # distribution of distances at which correct chains joined the pool.
        # Falls back to LOO-ECDF only in the very early cold-start window
        # before enough correct chains have been admitted.
        cal_conf = self._cal_dist_confidence(knn_dist)
        confidence = cal_conf if cal_conf is not None else self._ecdf_confidence(knn_dist)

        # Only auto-answer once DEPLOYED (AUROC ≥ 0.90 guaranteed)
        # During COLD_START and WARMING, always send to review
        needs_review = (status != "DEPLOYED") or (confidence < self.review_threshold)
        if needs_review:
            self._n_reviewed += 1

        return ScoreResult(
            confidence=round(confidence, 4),
            needs_review=needs_review,
            raw_distance=round(knn_dist, 6),
            calibration_size=n,
            estimated_auroc=round(est, 4),
            deployment_status=status,
            chains_to_go=ctg,
        )

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

    def add_verified(
        self,
        question: str,
        steps: list[dict] | None = None,
        correct: bool = True,
        raw_text: str = "",
        metadata: dict | None = None,
    ) -> CalibrateResult:
        """
        Add a verified chain to the calibration pool.

        Only CORRECT chains are added (the calibration pool is correct-chain-only,
        matching the QPPG KNN design). Incorrect chains are still valuable for
        monitoring but not stored here.

        Uses diversity-maximised insertion: a new chain is only added if it is
        at least `min_diversity_distance` away from all existing pool members,
        ensuring efficient use of the cold-start budget (exp30 finding).

        Parameters
        ----------
        question : str
        steps : list[dict], optional
        correct : bool
            Whether the chain's final answer was verified correct.
        raw_text : str, optional
        metadata : dict, optional
            Arbitrary metadata to store alongside the chain.

        Returns
        -------
        CalibrateResult
        """
        if not correct:
            return CalibrateResult(
                added=False, reason="Only correct chains are added to calibration pool",
                calibration_size=self.n_calibration,
                estimated_auroc=_estimate_auroc(self.n_calibration),
                deployment_status=self._deployment_status(),
                chains_to_go=self._chains_to_go(),
            )

        ctx   = _fmt_chain1(question, steps or [], raw_text)
        emb   = self._embed(ctx)
        emb_a = self._adapt_embs(emb[None, :])[0]   # 64-d if QARA, else 384-d

        # Diversity check (in adapter space if QARA fitted)
        pool_a = (self._adapt_embs(np.array(self._cal_embs))
                  if self._cal_embs else np.empty((0, len(emb_a))))
        div_dist = _most_diverse_index(pool_a, emb_a)
        if div_dist < self.min_diversity_distance and self.n_calibration > 0:
            return CalibrateResult(
                added=False,
                reason=f"Too similar to existing pool (dist={div_dist:.4f} < {self.min_diversity_distance})",
                calibration_size=self.n_calibration,
                estimated_auroc=_estimate_auroc(self.n_calibration),
                deployment_status=self._deployment_status(),
                chains_to_go=self._chains_to_go(),
            )

        # Record admission distance before adding (non-empty pool only)
        if self._cal_embs:
            adm_dist = float(np.linalg.norm(pool_a - emb_a[None, :], axis=1).min())
            self._cal_distances.append(adm_dist)

        # Add to pool and invalidate LOO cache
        self._cal_embs.append(emb)
        self._loo_dists_cache = None
        self._cal_meta.append({
            "question": question[:200],
            "ts": time.time(),
            **(metadata or {}),
        })
        self._save()

        n   = self.n_calibration
        est = _estimate_auroc(n)
        return CalibrateResult(
            added=True,
            reason=f"Added (diversity distance={div_dist:.4f})",
            calibration_size=n,
            estimated_auroc=est,
            deployment_status=self._deployment_status(),
            chains_to_go=self._chains_to_go(),
        )

    # ── Status ────────────────────────────────────────────────────────────────

    def status(self) -> dict:
        """Return a full status dict for dashboards/monitoring."""
        n   = self.n_calibration
        est = _estimate_auroc(n)
        dsp = self._deployment_status()
        days = _chains_to_deployment(n, self.queries_per_day, self.correct_rate)

        review_rate = (self._n_reviewed / self._n_scored
                       if self._n_scored > 0 else 1.0)

        return {
            "domain":             self.domain_name,
            "deployment_status":  dsp,
            "calibration_size":   n,
            "estimated_auroc":    round(est, 4),
            "chains_to_go":       self._chains_to_go(),
            "days_to_deployed":   round(days, 1),
            "queries_scored":     self._n_scored,
            "queries_reviewed":   self._n_reviewed,
            "review_rate":        round(review_rate, 3),
            "embed_model":        self.embed_model,
            "qara_fitted":        self._qara is not None,
            "save_path":          str(self._save_path),
        }

    def progress_bar(self) -> str:
        """ASCII progress bar toward DEPLOYED status."""
        n     = self.n_calibration
        pct   = min(1.0, n / WARMING_MIN)
        width = 40
        filled = int(pct * width)
        bar = "█" * filled + "░" * (width - filled)
        est = _estimate_auroc(n)
        dsp = self._deployment_status()
        return (
            f"Calibration: [{bar}] {n}/{WARMING_MIN}  "
            f"AUROC~{est:.2f}  [{dsp}]"
        )

    # ── Batch operations ──────────────────────────────────────────────────────

    def bulk_calibrate(
        self,
        chains: list[dict],
    ) -> dict:
        """
        Add multiple verified-correct chains at once (e.g. from existing logs).

        chains: list of {"question": str, "steps": list, "correct": bool}
        Returns summary dict.
        """
        n_added = n_skipped = n_wrong = 0
        for c in chains:
            if not c.get("correct", False):
                n_wrong += 1
                continue
            r = self.add_verified(
                question=c["question"],
                steps=c.get("steps", []),
                correct=True,
                raw_text=c.get("raw_text", ""),
                metadata=c.get("metadata"),
            )
            if r.added:
                n_added += 1
            else:
                n_skipped += 1
        return {
            "added": n_added, "skipped_duplicate": n_skipped,
            "skipped_incorrect": n_wrong,
            "calibration_size": self.n_calibration,
            "estimated_auroc": round(_estimate_auroc(self.n_calibration), 4),
            "deployment_status": self._deployment_status(),
        }

    # ── QARA adapter — fitting and persistence ────────────────────────────────

    def fit_qara(
        self,
        domains: list[dict],
        epochs:  int   = 200,
        lr:      float = 3e-4,
        verbose: bool  = True,
    ) -> dict:
        """
        Train a QARA adapter from multi-domain chain data.

        Parameters
        ----------
        domains : list of dicts, each containing:
            "name"   : str  — domain label (e.g. "natural_questions")
            "chains" : list of {"text": str, "correct": bool}
                where "text" is the chain-1 context string
                  (TASK: {question}\\nSTEP 1: Thought: ... | Action: ...\\nCURRENT: ...)

        After fitting, all KNN scoring uses 64-d adapted embeddings.
        The calibration pool's LOO distances are recomputed automatically.
        Previously recorded admission distances are cleared (they were in raw
        384-d space and are no longer comparable).

        Returns a summary dict with training statistics.

        Example
        -------
        service.fit_qara([
            {"name": "natural_questions",
             "chains": [{"text": ctx, "correct": True}, ...]},
            {"name": "hotpotqa",
             "chains": [{"text": ctx, "correct": True},
                        {"text": ctx, "correct": False}, ...]},
        ])
        """
        all_embs, all_labels = [], []
        total = sum(len(d["chains"]) for d in domains)
        if verbose:
            print(f"[fit_qara] Embedding {total} chains across {len(domains)} domains…")

        for d in domains:
            for c in d["chains"]:
                all_embs.append(self._embed(c["text"]))
                all_labels.append(1 if c.get("correct", True) else 0)

        embs   = np.array(all_embs,   dtype=np.float32)
        labels = np.array(all_labels, dtype=np.float32)  # 1.0=correct, 0.0=incorrect

        if verbose:
            nc = int(labels.sum())
            print(f"[fit_qara] {nc} correct / {len(labels)-nc} incorrect chains")
            print("[fit_qara] Training adapter (2-layer MLP, SupCon loss)…")

        self._qara = _train_qara_weights(embs, labels,
                                         epochs=epochs, lr=lr, verbose=verbose)

        # Invalidate caches — admission distances were in raw 384-d space
        self._loo_dists_cache = None
        self._cal_distances   = []
        self._save()

        if verbose:
            print("[fit_qara] Done. Scoring now uses 64-d QARA embeddings.")

        return {
            "n_domains": len(domains),
            "n_chains":  len(all_embs),
            "n_correct": round(float(labels.sum())),
            "epochs":    epochs,
            "adapter":   "QARAAdapter(384\u2192256\u219264)",
        }

    def save_qara(self, path: str | Path) -> None:
        """Save QARA adapter weights to a standalone pickle file."""
        if self._qara is None:
            raise ValueError("No QARA adapter fitted. Call fit_qara() first.")
        with open(path, "wb") as f:
            pickle.dump(self._qara, f)

    def load_qara(self, path: str | Path) -> None:
        """Load QARA adapter weights from a pickle file."""
        with open(path, "rb") as f:
            self._qara = pickle.load(f)
        self._loo_dists_cache = None   # recompute LOO dists in adapted space

    # ── Persistence ───────────────────────────────────────────────────────────

    def _save(self):
        state = {
            "domain_name":    self.domain_name,
            "embed_model":    self.embed_model,
            "cal_embs":       self._cal_embs,
            "cal_meta":       self._cal_meta,
            "cal_distances":  self._cal_distances,
            "n_scored":       self._n_scored,
            "n_reviewed":     self._n_reviewed,
            "qara":           self._qara,   # None or _QARAWeights
        }
        with open(self._save_path, "wb") as f:
            pickle.dump(state, f)

    def _load(self):
        with open(self._save_path, "rb") as f:
            state = pickle.load(f)
        self._cal_embs        = state.get("cal_embs", [])
        self._cal_meta        = state.get("cal_meta", [])
        self._cal_distances   = state.get("cal_distances", [])
        self._n_scored        = state.get("n_scored", 0)
        self._n_reviewed      = state.get("n_reviewed", 0)
        self._qara            = state.get("qara", None)
        self._loo_dists_cache = None   # recompute after load

    def reset(self):
        """Clear the calibration pool and reset all counters."""
        self._cal_embs     = []
        self._cal_meta     = []
        self._cal_distances = []
        self._n_scored     = 0
        self._n_reviewed   = 0
        self._loo_dists_cache = None
        self._score_buffer = []
        if self._save_path.exists():
            self._save_path.unlink()
        print(f"[QPPGService] Reset '{self.domain_name}'")

    def export_pool(self, path: str | Path):
        """Export calibration pool metadata to JSON (embeddings excluded)."""
        out = {
            "domain": self.domain_name,
            "n": self.n_calibration,
            "estimated_auroc": _estimate_auroc(self.n_calibration),
            "chains": self._cal_meta,
        }
        Path(path).write_text(json.dumps(out, indent=2))
