"""
silent_failure_calibrator.py
==============================
A feature-level trained classifier for detecting wrong LLM-agent chains,
including "silent failures" — chains that look behaviourally correct but
produce wrong answers.

Problem
-------
SC_OLD uses a fixed hand-crafted feature weighting calibrated on loud failures
(chains with obvious behavioural signals). For silent failures it assigns low
risk — AUROC 0.317, actively inverted. exp_decoherence_v2 showed that a
feature-level LogReg dramatically outperforms the fixed blend.

Design
------
- Extract 12 SC_OLD component features from AgentGuard.score_chain()
  (sc1..sc12, mean_sim, search_count_norm)
- Fit L2-regularised LogisticRegression (C=0.1) on labelled training chains
- The LogReg learns the optimal per-feature weight; no assumptions about which
  features matter

Validated (exp_silent_failure_calibrator.py — strict held-out protocol):
  TRAIN: large_dataset/hotpot_qa_chains.json      (n=200, HP questions A)
  TEST:  exp_ptrue_expanded/chains_checkpoint.json (n=200, HP questions B ≠ A)
  OOD:   exp156_live_crossdomain/trivia_chains.json (n=1000, TriviaQA)
  No chain overlap between train and either test set.

Usage
-----
    from llm_guard import AgentGuard, SilentFailureCalibrator

    guard = AgentGuard()
    cal   = SilentFailureCalibrator(guard)
    cal.fit(labeled_train_chains)    # chains must have "correct" field

    # Score a new chain
    risk = cal.score(question, steps, final_answer)

    # Or batch-score
    risks = cal.score_chains(chain_list)
"""

from __future__ import annotations

import warnings
from typing import Dict, List, Optional

import numpy as np


class SilentFailureCalibrator:
    """
    Feature-level trained classifier that outperforms SC_OLD on silent failures.

    Parameters
    ----------
    guard : AgentGuard, optional
        Existing AgentGuard instance (used for feature extraction via
        score_chain()).  If None, creates a fresh instance ($0, no API calls).
    C : float
        Inverse regularisation strength for LogisticRegression.  Default 0.1
        (strong regularisation — important on small n to prevent overfitting).
    """

    # 12 SC_OLD component features extracted from behavioral_components
    FEAT_KEYS: List[str] = [
        "sc1",                  # completion (binary: reached Finish)
        "sc2",                  # step count (normalised)
        "sc3",                  # thought embedding variance
        "sc5",                  # uncertainty words count
        "sc6",                  # answer gap (chars)
        "sc8",                  # backtracking rate
        "sc9",                  # observation utilisation
        "sc10",                 # thought coherence drop
        "sc11",                 # answer-question mismatch
        "sc12",                 # risk monotone slope
        "mean_sim",             # mean retrieval similarity
        "search_count_norm",    # normalised search count
    ]

    def __init__(self, guard=None, C: float = 0.1):
        from sklearn.linear_model import LogisticRegression
        from sklearn.preprocessing import StandardScaler

        if guard is None:
            from llm_guard.agent_guard import AgentGuard
            guard = AgentGuard()
        self._guard = guard
        self._C = C
        self._scaler = StandardScaler()
        self._lr = LogisticRegression(
            C=C,
            max_iter=1000,
            class_weight="balanced",
            solver="lbfgs",
            random_state=42,
        )
        self._is_fitted: bool = False
        self._n_train: int = 0
        self._n_train_wrong: int = 0
        self._n_train_correct: int = 0
        self._feature_coefs: Dict[str, float] = {}

    # ── Public API ─────────────────────────────────────────────────────────────

    def fit(self, chains: List[Dict]) -> "SilentFailureCalibrator":
        """
        Fit on labelled training chains.

        Parameters
        ----------
        chains : list of dict
            Each dict must contain "question", "steps", "final_answer", and
            "correct" (bool).  A minimum of 20 correct + 20 wrong chains is
            required; 50+ of each is recommended for stable estimates.

        Returns
        -------
        self  (for chaining)

        Raises
        ------
        ValueError
            If fewer than 2 distinct labels are present after extraction.
        RuntimeError
            If fewer than 10 chains score successfully.
        """
        X, y = self._extract_matrix(chains, require_labels=True)
        if len(set(y)) < 2:
            raise ValueError(
                "SilentFailureCalibrator.fit(): training chains must contain "
                "both correct and wrong examples."
            )
        if len(X) < 10:
            raise RuntimeError(
                f"SilentFailureCalibrator.fit(): only {len(X)} chains "
                "scored successfully — need at least 10."
            )

        X_s = self._scaler.fit_transform(X)
        self._lr.fit(X_s, y)
        self._is_fitted = True
        self._n_train = len(X)
        self._n_train_wrong   = int(y.sum())
        self._n_train_correct = int((y == 0).sum())

        # Store signed coefficients for interpretability
        for k, coef in zip(self.FEAT_KEYS, self._lr.coef_[0]):
            self._feature_coefs[k] = round(float(coef), 5)

        return self

    def score(
        self,
        question: str,
        steps: list,
        final_answer: str,
    ) -> float:
        """
        Score a single chain.

        Parameters
        ----------
        question, steps, final_answer : chain components

        Returns
        -------
        float : risk score in [0, 1].  Higher = higher probability of wrong.

        Raises
        ------
        RuntimeError  if fit() has not been called.
        """
        self._check_fitted()
        feats = self._extract_features_single(question, steps, final_answer)
        X = np.array([[feats.get(k, 0.0) for k in self.FEAT_KEYS]])
        X_s = self._scaler.transform(X)
        return float(self._lr.predict_proba(X_s)[0, 1])

    def score_chains(self, chains: List[Dict]) -> np.ndarray:
        """
        Batch score a list of chains.

        Parameters
        ----------
        chains : list of dict with "question", "steps", "final_answer".

        Returns
        -------
        np.ndarray of shape (n,) — risk scores in [0, 1].
        """
        self._check_fitted()
        X, _ = self._extract_matrix(chains, require_labels=False)
        X_s = self._scaler.transform(X)
        return self._lr.predict_proba(X_s)[:, 1]

    def feature_importances(self) -> Dict[str, float]:
        """
        Return a dict mapping feature name → logistic coefficient (signed).
        Positive = feature increases risk (predicts wrong).
        Negative = feature decreases risk (predicts correct).

        Raises
        ------
        RuntimeError  if fit() has not been called.
        """
        self._check_fitted()
        return dict(sorted(
            self._feature_coefs.items(),
            key=lambda x: abs(x[1]),
            reverse=True,
        ))

    # ── Properties ─────────────────────────────────────────────────────────────

    @property
    def is_fitted(self) -> bool:
        return self._is_fitted

    @property
    def n_train(self) -> int:
        return self._n_train

    # ── Internal helpers ───────────────────────────────────────────────────────

    def _check_fitted(self) -> None:
        if not self._is_fitted:
            raise RuntimeError(
                "SilentFailureCalibrator has not been fitted yet. "
                "Call fit(labeled_chains) first."
            )

    def _extract_features_single(
        self,
        question: str,
        steps: list,
        final_answer: str,
    ) -> Dict[str, float]:
        """Score one chain and return its behavioral_components dict."""
        try:
            r = self._guard.score_chain(question, steps, final_answer)
            return dict(r.behavioral_components)
        except Exception as exc:
            warnings.warn(
                f"SilentFailureCalibrator: feature extraction failed: {exc}",
                RuntimeWarning,
                stacklevel=3,
            )
            return {}

    def _extract_matrix(
        self,
        chains: List[Dict],
        require_labels: bool = True,
    ):
        """
        Extract feature matrix X and label vector y from a chain list.

        Labels: 1 = wrong (failed), 0 = correct.  Inferred from
        chain["correct"] (bool) when present.  If require_labels=False,
        returns y filled with -1 for chains lacking the "correct" field.
        """
        rows: List[List[float]] = []
        labels: List[int] = []
        n_failed = 0

        for i, chain in enumerate(chains):
            try:
                feats = self._extract_features_single(
                    chain["question"],
                    chain["steps"],
                    chain["final_answer"],
                )
                row = [feats.get(k, 0.0) for k in self.FEAT_KEYS]
                rows.append(row)

                if "correct" in chain:
                    labels.append(0 if chain["correct"] else 1)
                else:
                    if require_labels:
                        raise ValueError(
                            f"Chain {i} missing 'correct' field. "
                            "All training chains must be labelled."
                        )
                    labels.append(-1)

            except Exception as exc:
                n_failed += 1
                warnings.warn(
                    f"SilentFailureCalibrator: chain {i} skipped: {exc}",
                    RuntimeWarning,
                    stacklevel=2,
                )

        if n_failed > 0:
            warnings.warn(
                f"SilentFailureCalibrator: {n_failed}/{len(chains)} chains "
                "failed during feature extraction and were skipped.",
                RuntimeWarning,
                stacklevel=2,
            )

        X = np.array(rows, dtype=float)
        y = np.array(labels, dtype=int)
        return X, y

    def __repr__(self) -> str:
        return (
            f"SilentFailureCalibrator("
            f"C={self._C}, "
            f"fitted={self._is_fitted}, "
            f"n_train={self._n_train}"
            f")"
        )
