"""
DensityMatrixScorer — Quantum-inspired reasoning chain risk scorer.

Patent 7 (QPPG-QI) ported into llm-guard-kit.

Two modes
---------
Jaccard mode (default, zero deps):
    graph edges use Jaccard overlap for all weights.
    No additional dependencies beyond numpy.

Semantic mode (embedder= provided):
    adds E_sem edges using cosine similarity between thought embeddings.
    Requires: pip install sentence-transformers

Validated (Patent 7 experiments):
    Jaccard mode:   expected AUROC ~SC_OLD + 3-5pp on HP within-domain
    Semantic mode:  expected AUROC ~SC_OLD + 7-8pp on HP within-domain

Quick start
-----------
    from llm_guard import DensityMatrixScorer

    scorer = DensityMatrixScorer()
    scorer.fit(labeled_chains)   # list of {"question", "steps", "final_answer", "correct": bool}
    result = scorer.score(question, steps, final_answer)
    print(result.risk_score)      # 0.0-1.0
    print(result.sc_old_risk)     # SC_OLD behavioral baseline for comparison
Purpose:    Models agent reasoning steps as a density matrix graph and computes risk from trace coherence. Jaccard mode needs no extra deps; semantic mode with sentence-transformers gains ~7-8pp AUROC over SC_OLD.
Exports:    DensityMatrixScorer
Added:      v0.44.0
Status:     Experimental
"""
from __future__ import annotations

import re
from collections import Counter
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np


# ── Text utilities ──────────────────────────────────────────────────────────

def _tokenize(text: str) -> List[str]:
    return re.findall(r'\b\w+\b', text.lower())

def _jaccard(a: str, b: str) -> float:
    sa, sb = set(_tokenize(a)), set(_tokenize(b))
    if not sa and not sb:
        return 0.0
    return len(sa & sb) / len(sa | sb)

def _tfidf_overlap(obs_i: str, obs_j: str) -> float:
    ti, tj = _tokenize(obs_i), _tokenize(obs_j)
    if not ti or not tj:
        return 0.0
    cf = Counter(ti)
    shared = sum(min(cf[w], 1) for w in tj if w in cf)
    return shared / max(len(set(ti) | set(tj)), 1)


def _encode_amplitude(phi: np.ndarray) -> np.ndarray:
    """Normalize feature vector to unit amplitude: |ψ⟩ = φ / ‖φ‖."""
    norm = np.linalg.norm(phi)
    if norm < 1e-10:
        d = len(phi)
        return np.ones(d) / np.sqrt(max(d, 1))
    return phi / norm


# ── PatternGraph ─────────────────────────────────────────────────────────────

@dataclass
class PatternGraph:
    """Lightweight graph with nodes (dicts with 'id', 'text') and weighted edges."""
    nodes: List[Dict[str, Any]]
    edges: List[Tuple[str, str, float]]  # (src_id, dst_id, weight)

    def adjacency(self) -> np.ndarray:
        ids = [n["id"] for n in self.nodes]
        idx = {nid: i for i, nid in enumerate(ids)}
        n = len(ids)
        A = np.zeros((n, n))
        for src, dst, w in self.edges:
            if src in idx and dst in idx:
                A[idx[src], idx[dst]] = w
                A[idx[dst], idx[src]] = w  # undirected
        return A

    def laplacian(self) -> np.ndarray:
        A = self.adjacency()
        D = np.diag(A.sum(axis=1))
        return D - A

    def spectral_gap(self) -> float:
        """λ₂ of Laplacian. Small gap = disconnected graph = suspicious chain."""
        if len(self.nodes) < 2:
            return 0.0
        L = self.laplacian()
        eigs = np.linalg.eigvalsh(L)
        return float(np.sort(eigs)[1]) if len(eigs) > 1 else 0.0

    def pagerank(self, alpha: float = 0.85, n_iter: int = 50) -> np.ndarray:
        A = self.adjacency()
        n = A.shape[0]
        if n == 0:
            return np.array([])
        row_sums = A.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1
        P = A / row_sums
        r = np.ones(n) / n
        for _ in range(n_iter):
            r = alpha * P.T @ r + (1 - alpha) / n
        r = r / (r.sum() + 1e-10)
        return r

    def pagerank_entropy(self) -> float:
        """Shannon entropy of PageRank distribution. High entropy = uniform = suspicious."""
        pr = self.pagerank()
        if len(pr) == 0:
            return 0.0
        pr = np.clip(pr, 1e-10, 1.0)
        return float(-np.sum(pr * np.log(pr)))


def build_pattern_graph(
    steps: List[Dict],
    embedder: Optional[Any] = None,
    sem_threshold: float = 0.6,
    ref_threshold: float = 0.3,
) -> PatternGraph:
    """
    Build a PatternGraph from ReAct chain steps.

    Edge types
    ----------
    E_seq : thought→action, action→observation, observation→next_thought
    E_ref : TF-IDF overlap between observations if overlap > ref_threshold
    E_sem : cosine sim between thought embeddings if sim > sem_threshold (embedder only)
    """
    nodes: List[Dict] = []
    edges: List[Tuple[str, str, float]] = []

    thoughts = [s.get("thought", "") for s in steps]
    actions  = [s.get("action_arg", s.get("action", "")) for s in steps]
    obs      = [s.get("observation", "") for s in steps]

    # 1. Nodes
    for i in range(len(steps)):
        nodes.extend([
            {"id": f"t{i}", "text": thoughts[i]},
            {"id": f"a{i}", "text": actions[i]},
            {"id": f"o{i}", "text": obs[i]},
        ])

    # 2. Sequential edges (E_seq)
    for i in range(len(steps)):
        edges.append((f"t{i}", f"a{i}", _jaccard(thoughts[i], actions[i]) + 1e-3))
        edges.append((f"a{i}", f"o{i}", 1.0))
        if i < len(steps) - 1:
            ig = max(0.0, 1.0 - _jaccard(obs[i], thoughts[i + 1]))
            edges.append((f"o{i}", f"t{i+1}", ig + 1e-3))

    # 3. Reference edges (E_ref)
    for i in range(len(steps)):
        for j in range(i):
            ov = _tfidf_overlap(obs[i], obs[j])
            if ov > ref_threshold:
                edges.append((f"o{j}", f"o{i}", ov))

    # 4. Semantic edges (E_sem) — requires embedder
    if embedder is not None and len(steps) > 1:
        try:
            embeds = embedder.encode(thoughts, show_progress_bar=False)
            norms = np.linalg.norm(embeds, axis=1, keepdims=True) + 1e-10
            embeds = embeds / norms
            for i in range(len(steps)):
                for j in range(i + 2, len(steps)):  # skip adjacent (already sequential)
                    sim = float(np.dot(embeds[i], embeds[j]))
                    if sim > sem_threshold:
                        edges.append((f"t{i}", f"t{j}", sim))
        except Exception:
            pass  # graceful degradation if embedder fails

    return PatternGraph(nodes=nodes, edges=edges)


# ── Density matrix ────────────────────────────────────────────────────────────

def _compute_step_phi(
    step: Dict,
    step_idx: int,
    steps: List[Dict],
    graph: PatternGraph,
    _A_mat: np.ndarray,
    _pr: np.ndarray,
    _sg: float,
) -> np.ndarray:
    """16-dim behavioral + graph feature vector for step i.

    Parameters
    ----------
    _A_mat, _pr, _sg : pre-computed graph-level properties (adjacency, pagerank,
        spectral_gap).  Callers MUST compute these once and pass them in to avoid
        redundant O(n²) recomputation on every step.
    """
    T = step.get("thought", "") or ""
    A = (step.get("action_arg", "") or "") + " " + (step.get("action_type", "") or "")
    O = step.get("observation", "") or ""
    i = step_idx
    n = len(steps)

    prev_actions = [(steps[k].get("action_arg", "") or "") for k in range(i)]
    f1 = 1.0 - max((_jaccard(A, pa) for pa in prev_actions), default=0.0)
    f2 = _jaccard(T, A)
    obs_lens = [len((steps[k].get("observation", "") or "").split()) for k in range(n)]
    mean_len = max(float(np.mean(obs_lens)), 1.0)
    f3 = min(len(O.split()) / mean_len, 3.0) / 3.0
    prev_O = steps[i - 1].get("observation", "") if i > 0 else ""
    f4 = _jaccard(T, prev_O) if prev_O else 0.5
    T0 = steps[0].get("thought", "") or ""
    f5 = _jaccard(T, T0) if i > 0 else 1.0
    null_phrases = ["no result", "not found", "could not", "error", "none"]
    f6 = 0.1 if any(p in O.lower() for p in null_phrases) else 0.9
    action_types = set(steps[k].get("action_type", "") for k in range(i + 1))
    f7 = min(len(action_types) / max(n, 1), 1.0)
    f8 = i / max(n - 1, 1)

    # Graph structural features (use pre-computed properties)
    A_mat = _A_mat
    pr = _pr
    n_nodes = len(graph.nodes)
    t_idx = 3 * i if 3 * i < n_nodes else 0
    o_idx = 3 * i + 2 if 3 * i + 2 < n_nodes else 0

    sg = _sg
    f9  = float(pr[t_idx]) if len(pr) > t_idx else 0.1
    f10 = float(A_mat[t_idx].sum()) / max(n_nodes, 1) if n_nodes > 0 else 0.0
    f11 = min(sg, 5.0) / 5.0
    degrees = A_mat.sum(axis=1)
    cc_list = []
    for ii in range(n_nodes):
        nbrs = np.where(A_mat[ii] > 0)[0]
        k = len(nbrs)
        if k >= 2:
            links = sum(1 for a in nbrs for b in nbrs if a < b and A_mat[a, b] > 0)
            possible = k * (k - 1) / 2
            cc_list.append(links / possible if possible > 0 else 0.0)
    f12 = float(np.mean(cc_list)) if cc_list else 0.0
    isolated = int((degrees == 0).sum())
    f13 = isolated / max(n_nodes, 1)
    max_edges = n_nodes * (n_nodes - 1) / 2
    f14 = len(graph.edges) / max(max_edges, 1)
    f15 = float(A_mat[o_idx].sum()) / max(n_nodes, 1) if n_nodes > 0 else 0.0
    f16 = _jaccard(T, O)

    phi = np.array([f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15, f16], dtype=np.float64)
    return np.clip(phi, 0.0, 1.0)


def _compute_step_features(
    step: Dict,
    all_steps: List[Dict],
    graph: PatternGraph,
    step_idx: int,
) -> np.ndarray:
    """Public alias for _compute_step_phi (called with pre-computed graph properties)."""
    _A_mat = graph.adjacency()
    _pr = graph.pagerank()
    _sg = graph.spectral_gap()
    return _compute_step_phi(step, step_idx, all_steps, graph, _A_mat, _pr, _sg)


def _build_density_matrix(
    alphas: List[np.ndarray],
    centralities: np.ndarray,
) -> np.ndarray:
    """
    Build density matrix ρ = Σ_i w_i |ψ_i⟩⟨ψ_i|.
    Normalized so Tr(ρ) = 1.
    """
    if not alphas:
        d = 16
        return np.zeros((d, d))
    d = len(alphas[0])
    rho = np.zeros((d, d))
    for alpha_i, w_i in zip(alphas, centralities):
        rho += float(w_i) * np.outer(alpha_i, alpha_i)
    tr = np.trace(rho)
    if tr > 1e-10:
        rho /= tr
    return rho


def _build_rho_from_steps(steps: List[Dict], graph: PatternGraph) -> np.ndarray:
    """ρ = Σ_i w_i |ψ_i⟩⟨ψ_i| where w_i = PageRank of thought node i."""
    n = len(steps)
    if n == 0:
        return np.zeros((16, 16))
    # Pre-compute graph-level properties ONCE (they are step-invariant).
    _A_mat = graph.adjacency()
    _pr = graph.pagerank()
    _sg = graph.spectral_gap()
    alphas = []
    centralities = []
    for i, step in enumerate(steps):
        phi = _compute_step_phi(step, i, steps, graph, _A_mat, _pr, _sg)
        alpha = _encode_amplitude(phi)
        node_idx = 3 * i if 3 * i < len(_pr) else 0
        w = float(_pr[node_idx]) if len(_pr) > node_idx else 1.0 / n
        alphas.append(alpha)
        centralities.append(w)
    centralities_arr = np.array(centralities)
    # Normalize centralities to sum to 1 before building rho
    total_w = centralities_arr.sum()
    if total_w > 1e-10:
        centralities_arr = centralities_arr / total_w
    return _build_density_matrix(alphas, centralities_arr)


def _sc_old_score(question: str, steps: List[Dict], final_answer: str) -> float:
    """Lightweight SC_OLD proxy (step count + finish signal). $0, no external deps."""
    n = len(steps)
    if n == 0:
        return 0.5
    finish_actions = {"finish", "final_answer", "answer"}
    has_finish = any(
        s.get("action_type", "").lower() in finish_actions for s in steps
    )
    # Step count signal (higher steps → higher risk, normalized)
    step_risk = min(n / 10.0, 1.0)
    # Finish signal
    finish_risk = 0.1 if has_finish else 0.7
    return float(np.clip(0.6 * step_risk + 0.4 * finish_risk, 0.0, 1.0))


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

@dataclass
class DensityMatrixResult:
    """Result of DensityMatrixScorer.score()."""
    risk_score: float
    sc_old_risk: float
    graph_spectral_gap: float   # was spectral_gap
    mode: str                   # "jaccard" | "semantic" — NEW
    density_matrix: Optional[np.ndarray] = None  # NEW, None unless return_density_matrix=True


class DensityMatrixScorer:
    """
    Quantum-inspired density matrix risk scorer (Patent 7 / QPPG-QI).

    Builds a PatternGraph from reasoning chain steps, encodes each step
    as a quantum amplitude vector, and constructs a density matrix ρ.
    Risk is measured via spectral gap and (when fitted) measurement
    operator interference.

    Modes
    -----
    Jaccard (default): all edge weights from Jaccard token overlap.
      Zero extra dependencies. Expected AUROC ~SC_OLD + 3-5pp.

    Semantic (embedder= provided): adds E_sem edges via cosine similarity.
      Requires sentence-transformers. Expected AUROC ~SC_OLD + 7-8pp.

    Parameters
    ----------
    spectral_weight : float
        Weight of spectral gap risk in composite score. Default 0.3.
    embedder : optional
        sentence-transformers SentenceTransformer instance for semantic edges.
    delta_fit : float
        Reserved for future Fisher regularization (Wave 2+). Currently unused.
        Controls maximum per-parameter update magnitude during continual learning.
    return_density_matrix : bool
        If True, include the density matrix ndarray in the result. Default False.
    """

    def __init__(
        self,
        spectral_weight: float = 0.3,
        embedder: Optional[Any] = None,
        delta_fit: float = 1e-4,
        return_density_matrix: bool = False,
    ) -> None:
        self.spectral_weight = spectral_weight
        self.embedder = embedder
        self._delta_fit = delta_fit
        self._return_density_matrix = return_density_matrix
        self._m_correct: Optional[np.ndarray] = None
        self._m_incorrect: Optional[np.ndarray] = None
        self._fitted = False

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

    def fit(
        self,
        labeled_chains: List[Dict],
        embedder: Optional[Any] = None,
    ) -> "DensityMatrixScorer":
        """
        Fit measurement operators from labeled chains.

        Parameters
        ----------
        labeled_chains : List[Dict]
            Each dict: {"question": str, "steps": List[Dict],
                        "final_answer": str, "correct": bool}
        embedder : optional
            Override the scorer's embedder for this fit call only.

        Returns
        -------
        self
        """
        correct_chains = [c for c in labeled_chains if c.get("correct") is True]
        wrong_chains   = [c for c in labeled_chains if c.get("correct") is False]
        n_c, n_w = len(correct_chains), len(wrong_chains)
        if n_c < 10 or n_w < 10:
            raise ValueError(
                f"DensityMatrixScorer.fit() requires ≥10 correct and ≥10 wrong chains; "
                f"got {n_c} correct and {n_w} wrong."
            )

        emb = embedder or self.embedder
        correct_rhos, wrong_rhos = [], []
        for item in labeled_chains:
            steps = item.get("steps", [])
            if not steps:
                continue
            graph = build_pattern_graph(steps, embedder=emb)
            rho = _build_rho_from_steps(steps, graph)
            if "correct" not in item:
                raise ValueError(
                    f"Each chain dict must have a 'correct' key; got keys: {list(item.keys())}"
                )
            if item["correct"]:
                correct_rhos.append(rho)
            else:
                wrong_rhos.append(rho)
        self._m_correct = np.mean(correct_rhos, axis=0) if correct_rhos else np.eye(16) / 16
        self._m_incorrect = np.mean(wrong_rhos, axis=0) if wrong_rhos else np.eye(16) / 16
        self._fitted = True
        return self

    def score(
        self,
        question: str,
        steps: List[Dict],
        final_answer: str,
        embedder: Optional[Any] = None,
    ) -> DensityMatrixResult:
        """
        Score a reasoning chain.

        Parameters
        ----------
        question : str
        steps : List[Dict]
        final_answer : str
        embedder : optional
            Override the scorer's embedder for this call only.

        Returns
        -------
        DensityMatrixResult
        """
        if not self._fitted:
            raise RuntimeError(
                "DensityMatrixScorer must be fitted before scoring. "
                "Call scorer.fit(labeled_chains) first."
            )

        emb = embedder or self.embedder
        if not steps:
            return DensityMatrixResult(
                risk_score=0.5,
                sc_old_risk=0.5,
                graph_spectral_gap=0.0,
                mode="semantic" if emb is not None else "jaccard",
                density_matrix=None,
            )
        graph = build_pattern_graph(steps, embedder=emb)
        rho = _build_rho_from_steps(steps, graph)
        sg = graph.spectral_gap()
        risk_spectral = float(1.0 / (1.0 + sg))

        correct_w   = float(np.trace(self._m_correct   @ rho))
        incorrect_w = float(np.trace(self._m_incorrect @ rho))
        interference = correct_w - incorrect_w
        risk_interference = float(1.0 / (1.0 + np.exp(interference)))
        risk = (
            (1.0 - self.spectral_weight) * risk_interference
            + self.spectral_weight * risk_spectral
        )

        sc_old = _sc_old_score(question, steps, final_answer)
        return DensityMatrixResult(
            risk_score=float(np.clip(risk, 0.0, 1.0)),
            sc_old_risk=sc_old,
            graph_spectral_gap=sg,
            mode="semantic" if emb is not None else "jaccard",
            density_matrix=rho.copy() if self._return_density_matrix else None,
        )

    def score_batch(
        self,
        chains: List[Dict],
        embedder: Optional[Any] = None,
    ) -> List[DensityMatrixResult]:
        """Score a list of chain dicts. Each dict: {question, steps, final_answer}."""
        emb = embedder or self.embedder
        return [
            self.score(
                c.get("question", ""),
                c.get("steps", []),
                c.get("final_answer", c.get("answer", "")),
                embedder=emb,
            )
            for c in chains
        ]

    def features(self, steps: List[Dict], embedder: Optional[Any] = None) -> np.ndarray:
        """
        Extract upper-triangle of density matrix + spectral gap + entropy.
        Returns 138-dim vector for downstream ML (136 upper-tri + 2 scalars).
        """
        emb = embedder or self.embedder
        if not steps:
            return np.zeros(138)
        graph = build_pattern_graph(steps, embedder=emb)
        rho = _build_rho_from_steps(steps, graph)
        upper = rho[np.triu_indices(16)]  # 136 values
        extras = np.array([graph.spectral_gap() / 5.0, graph.pagerank_entropy()])
        return np.concatenate([upper, extras])
