"""
deep_chain_scorer.py — Reliable scoring for deep multi-hop agent chains.
=========================================================================

SC_OLD behavioral ensemble degrades on 5–8 step chains (2wiki AUROC 0.545).
This module provides two improvements:

  RetrievalCascadeScorer (label-free, $0, pure numpy)
  ---------------------------------------------------
    Trajectory slope of Jaccard(query, observation) across Search steps.
    A chain whose retrieval quality IMPROVES over time is more likely correct.
    AUROC 0.629 on 2wiki (vs SC_OLD 0.545), 0.742 on TriviaQA.
    Gate: validated in exp_architecture_benchmark across 5 datasets.
    Returns values in [0, 1] — higher = more likely failure.
    Falls back to 0.5 (neutral) when < 2 Search steps.

  DeepChainScorer (composite, label-free)
  ----------------------------------------
    Auto-selects scorer based on chain length:
      n_search < chain_length_threshold  → SC_OLD (validated on short chains)
      n_search >= chain_length_threshold  → RetrievalCascade (better on long)
      n_search >= mamba_threshold         → MambaRiskScorer if torch available
    Returns DeepChainResult with risk_score + which scorer was used.

  MambaRiskScorer integration (optional, requires [qara])
  -------------------------------------------------------
    Pass a fitted MambaRiskScorer instance to DeepChainScorer(mamba_scorer=...).
    Falls back gracefully to RetrievalCascade if torch is not installed.

Validated AUROC (exp_architecture_benchmark + exp_deep_chain_scorer):
  Dataset         SC_OLD   RetrievalCascade   Mamba (research)
  HP (2-3 steps)  0.817    0.592              0.723
  NQ (2-3 steps)  0.810    0.591              0.714
  2wiki (6+ steps) 0.545   0.629              0.708
  Musique         0.571    0.669              0.715
  TriviaQA        0.660    0.742              0.698

Usage
-----
    from llm_guard import DeepChainScorer

    scorer = DeepChainScorer()
    result = scorer.score(question, steps, final_answer)
    print(result.risk_score, result.scorer_used, result.n_search_steps)

    # With optional Mamba (requires pip install llm-guard-kit[qara])
    from llm_guard import DeepChainScorer
    from experiments.exp_mamba_sequence_scorer import MambaRiskScorer
    mamba = MambaRiskScorer()
    mamba.fit(labeled_chains)
    scorer = DeepChainScorer(mamba_scorer=mamba)
Purpose:    Provides RetrievalCascadeScorer (AUROC 0.629 on 2WikiMultiHop) and DeepChainScorer that auto-selects between SC_OLD and cascade scoring based on chain length. Addresses SC_OLD degradation on 5-8 step chains.
Exports:    DeepChainScorer, RetrievalCascadeScorer, DeepChainResult
Added:      v0.1.0
Status:     Validated
"""

from __future__ import annotations

import warnings
from dataclasses import dataclass
from typing import Any, Dict, List, Optional


# ── Jaccard helper (no external deps) ────────────────────────────────────────

def _jaccard(a: str, b: str) -> float:
    """Token-level Jaccard similarity between two strings."""
    if not a or not b:
        return 0.0
    ta, tb = set(a.lower().split()), set(b.lower().split())
    if not ta or not tb:
        return 0.0
    return len(ta & tb) / len(ta | tb)


# ── RetrievalCascadeScorer ────────────────────────────────────────────────────

class RetrievalCascadeScorer:
    """
    Label-free scorer based on the trajectory of retrieval quality across
    Search steps. Pure numpy — no LLM calls, no external deps beyond numpy.

    Hypothesis: A chain whose Jaccard(query, observation) improves over time
    is more likely to have found relevant evidence. Declining retrieval quality
    (early plateau then drop) correlates with failure.

    Validated on exp_architecture_benchmark:
      2wiki AUROC 0.629 [0.555, 0.702] — +8pp over SC_OLD (0.545)
      Musique AUROC 0.669 [0.585, 0.752]
      TriviaQA AUROC 0.742 [0.666, 0.813]
      HP AUROC 0.592 (below SC_OLD 0.817 — not recommended for short HP chains)

    Parameters
    ----------
    min_search_steps : int
        Minimum Search steps required to compute a meaningful slope.
        Chains with fewer steps return 0.5 (neutral). Default: 2.
    """

    def __init__(self, min_search_steps: int = 2):
        self.min_search_steps = min_search_steps

    def score(self, steps: List[Dict]) -> float:
        """
        Compute cascade risk score from step list.

        Parameters
        ----------
        steps : list of dicts
            Agent steps in canonical format {action_type, action_arg, observation, ...}

        Returns
        -------
        float in [0, 1] — higher = more likely failure.
        Returns 0.5 (neutral) when fewer than min_search_steps Search steps.
        """
        import numpy as _np

        search_steps = [s for s in steps if s.get("action_type", "") == "Search"]
        if len(search_steps) < self.min_search_steps:
            return 0.5

        sims = [_jaccard(s.get("action_arg", ""), s.get("observation", ""))
                for s in search_steps]
        n_s = len(sims)

        # Slope: quality change from first half to second half
        first = float(_np.mean(sims[:n_s // 2])) if n_s >= 2 else sims[0]
        last  = float(_np.mean(sims[n_s // 2:])) if n_s >= 2 else sims[-1]
        slope = last - first  # positive = improving = lower risk

        # Best improvement step (normalised by chain length)
        if len(sims) >= 3:
            best_step   = int(_np.argmax(_np.diff(sims)))
            best_step_n = best_step / max(len(sims) - 1, 1)
        else:
            best_step_n = 0.5

        # Risk: declining slope + late best step = high failure risk
        cascade_risk = (1.0 - slope) * 0.6 + (1.0 - best_step_n) * 0.4
        return float(_np.clip(cascade_risk, 0.0, 1.0))

    def score_chain(self, chain: Dict) -> float:
        """Convenience wrapper accepting a chain dict with 'steps' key."""
        return self.score(chain.get("steps", []))


# ── DeepChainResult ────────────────────────────────────────────────────────────

@dataclass
class DeepChainResult:
    """
    Result from DeepChainScorer.score().

    Fields
    ------
    risk_score : float
        Chain risk in [0, 1]. Higher = more likely failure.
    scorer_used : str
        Which scorer produced this score:
          "sc_old"             — SC_OLD behavioral ensemble (short chains)
          "retrieval_cascade"  — RetrievalCascadeScorer (long chains, ≥ threshold)
          "mamba"              — MambaRiskScorer (very long chains, if torch available)
    n_search_steps : int
        Number of Search-type steps in the chain.
    n_steps : int
        Total number of steps.
    chain_length_threshold : int
        The threshold used to switch from SC_OLD to RetrievalCascade.
    """
    risk_score: float
    scorer_used: str
    n_search_steps: int
    n_steps: int
    chain_length_threshold: int


# ── DeepChainScorer ───────────────────────────────────────────────────────────

class DeepChainScorer:
    """
    Composite scorer that auto-selects the best algorithm based on chain depth.

    Routing logic (default thresholds):
      n_search_steps < 4   → SC_OLD (behavioral ensemble, validated on HP/NQ)
      4 ≤ n_search_steps   → RetrievalCascade (trajectory slope, better on 2wiki/Musique)
      n_search_steps ≥ 6   → Mamba (if a fitted mamba_scorer is provided AND torch available)

    Parameters
    ----------
    chain_length_threshold : int
        Switch from SC_OLD to RetrievalCascade at or above this many Search steps.
        Default: 4. Set to 0 to always use RetrievalCascade.
    mamba_threshold : int
        Use MambaRiskScorer at or above this many Search steps (if available).
        Default: 6. Only used when mamba_scorer is provided.
    mamba_scorer : object, optional
        Fitted MambaRiskScorer instance. If None, Mamba is never used.
        Must have predict_proba(chains: list) -> np.ndarray method.
    review_threshold : float
        Risk threshold passed to LabelFreeScorer for SC_OLD path. Default: 0.65.
    """

    def __init__(
        self,
        chain_length_threshold: int = 4,
        mamba_threshold: int = 6,
        mamba_scorer: Optional[Any] = None,
        review_threshold: float = 0.65,
    ):
        self.chain_length_threshold = chain_length_threshold
        self.mamba_threshold        = mamba_threshold
        self.mamba_scorer           = mamba_scorer
        self.review_threshold       = review_threshold
        self._cascade               = RetrievalCascadeScorer()
        self._sc_old                = None  # lazy init

    def _get_sc_old(self):
        if self._sc_old is None:
            import sys
            from pathlib import Path
            sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
            from qppg_service.label_free_scorer import LabelFreeScorer
            self._sc_old = LabelFreeScorer(review_threshold=self.review_threshold)
        return self._sc_old

    def score(
        self,
        question: str,
        steps: List[Dict],
        final_answer: str,
    ) -> DeepChainResult:
        """
        Score a reasoning chain using the best available scorer for its depth.

        Parameters
        ----------
        question : str
            The original question.
        steps : list of dicts
            Agent steps in canonical {action_type, action_arg, observation, thought} format.
        final_answer : str
            The agent's final answer string.

        Returns
        -------
        DeepChainResult
        """
        search_steps = [s for s in steps if s.get("action_type", "") == "Search"]
        n_search = len(search_steps)
        n_steps  = len(steps)

        # Route to Mamba if fitted scorer provided and chain is long enough
        if (
            self.mamba_scorer is not None
            and n_search >= self.mamba_threshold
        ):
            try:
                chain = {"question": question, "steps": steps,
                         "final_answer": final_answer, "correct": True}
                preds = self.mamba_scorer.predict_proba([chain])
                risk = float(preds[0])
                return DeepChainResult(
                    risk_score=risk,
                    scorer_used="mamba",
                    n_search_steps=n_search,
                    n_steps=n_steps,
                    chain_length_threshold=self.chain_length_threshold,
                )
            except Exception:
                pass  # fall through to RetrievalCascade

        # Route to RetrievalCascade for long chains
        if n_search >= self.chain_length_threshold:
            risk = self._cascade.score(steps)
            return DeepChainResult(
                risk_score=risk,
                scorer_used="retrieval_cascade",
                n_search_steps=n_search,
                n_steps=n_steps,
                chain_length_threshold=self.chain_length_threshold,
            )

        # Short chains: SC_OLD behavioral ensemble
        sc_old = self._get_sc_old()
        risk   = sc_old.score(question, steps, final_answer).risk_score
        return DeepChainResult(
            risk_score=risk,
            scorer_used="sc_old",
            n_search_steps=n_search,
            n_steps=n_steps,
            chain_length_threshold=self.chain_length_threshold,
        )
