nexus_os_v2/unified_detector.py

"""
Unified Thermodynamic Hallucination Detector for NEXUS OS v2.1

Integrates four empirically-validated detection signals:
  1. EPR (Entropy Production Rate) — arXiv:2509.04492
  2. Spilled Energy — arXiv:2602.18671  
  3. CK-PLUG Confidence Gain — arXiv:2503.15888
  4. TWAVE Landau-Ginzburg — NEXUS OS novel framework

Plus novel composite signals:
  5. Energy-Entropy Product (EEP) — correlated instability
  6. Phase Transition Index (PTI) — divergence of order parameters
  7. Non-Equilibrium Work Index (NEWI) — Jarzynski cumulative measure

Architecture:
  Per-token:    Each detector produces a score ∈ [0,1]
  Per-sequence: Weighted fusion + ensemble disagreement
  Action:       None / Ground / Reflect / Halt / Switch Model

Detector Fusion Strategy:
  - Agreement mode: All detectors must agree for action (conservative)
  - Majority mode:  >50% detectors trigger for action (balanced)
  - Any mode:       Any detector triggers for action (sensitive)
  - Weighted mode:  Linear combination with learned weights
"""
import math
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass, field
from enum import Enum

from .epr_detector import EPRDetector, SequenceEPR
from .spilled_energy import SpilledEnergyDetector, SpilledEnergyReading
from .ckplug_retriever import CKPLUGCoupling
from .twave_tracker import TWAVETracker, TokenState, GenerationTrajectory


class FusionMode(Enum):
    AGREEMENT = "agreement"      # All detectors must agree
    MAJORITY = "majority"        # >50% detectors
    ANY = "any"                  # Any single detector
    WEIGHTED = "weighted"        # Linear combination


class Action(Enum):
    NONE = "none"                # Continue normally
    GROUND = "ground"            # Boost retrieval evidence
    REFLECT = "reflect"          # Backtrack and regenerate
    HALT = "halt"                # Stop generation
    SWITCH = "switch"            # Switch to different model


@dataclass
class DetectorReading:
    """Reading from a single detector at a token position."""
    detector_name: str
    score: float                 # [0,1] hallucination score
    confidence: float            # [0,1] detector confidence in this reading
    is_triggered: bool           # Does this detector trigger action?
    details: Dict[str, float] = field(default_factory=dict)


@dataclass
class TokenVerdict:
    """Unified verdict for a single token position."""
    position: int
    token_str: str
    readings: List[DetectorReading]
    fused_score: float           # Combined score [0,1]
    risk_level: str              # low/moderate/elevated/high/critical
    recommended_action: Action
    confidence: float            # Confidence in verdict [0,1]


@dataclass
class SequenceVerdict:
    """Unified verdict for an entire generation sequence."""
    token_verdicts: List[TokenVerdict]
    overall_risk: str
    overall_action: Action
    num_triggers: int
    trigger_positions: List[int]
    avg_fused_score: float
    max_fused_score: float
    detector_agreement: float    # How often detectors agree
    
    # Composite novel signals
    energy_entropy_product: float  # EEP = max(energy_score * entropy_score)
    phase_transition_index: float  # PTI = divergence of order parameters
    newi: float                    # Non-Equilibrium Work Index


class UnifiedThermodynamicDetector:
    """
    Production unified hallucination detector.
    
    Combines multiple thermodynamic signals with configurable fusion.
    """
    
    # Default thresholds (calibrated on validation set)
    RISK_THRESHOLDS = {
        "low": 0.0,
        "moderate": 0.25,
        "elevated": 0.40,
        "high": 0.60,
        "critical": 0.80,
    }
    
    # Action mapping from risk level
    ACTION_MAP = {
        "low": Action.NONE,
        "moderate": Action.NONE,
        "elevated": Action.GROUND,
        "high": Action.REFLECT,
        "critical": Action.HALT,
    }
    
    def __init__(
        self,
        fusion_mode: FusionMode = FusionMode.WEIGHTED,
        weights: Optional[Dict[str, float]] = None,
        enable_epr: bool = True,
        enable_spilled: bool = True,
        enable_ckplug: bool = True,
        enable_twave: bool = True,
    ):
        self.fusion_mode = fusion_mode
        self.enable_epr = enable_epr
        self.enable_spilled = enable_spilled
        self.enable_ckplug = enable_ckplug
        self.enable_twave = enable_twave
        
        # Initialize sub-detectors
        self.epr = EPRDetector() if enable_epr else None
        self.spilled = SpilledEnergyDetector() if enable_spilled else None
        self.ckplug = CKPLUGCoupling(mu_0=0.5) if enable_ckplug else None
        self.twave = TWAVETracker(T_c=1.0, mu_0=0.5, kappa=0.1) if enable_twave else None
        
        # Detector weights for weighted fusion
        default_weights = {
            "epr": 0.25,
            "spilled": 0.25,
            "ckplug": 0.25,
            "twave": 0.25,
        }
        self.weights = weights or default_weights
        
        # Normalize weights
        total = sum(self.weights.values())
        if total > 0:
            self.weights = {k: v / total for k, v in self.weights.items()}
    
    def _fuse_scores(self, readings: List[DetectorReading]) -> Tuple[float, str, float]:
        """
        Fuse detector scores into unified score, risk, and confidence.
        
        Returns: (fused_score, risk_level, confidence)
        """
        if not readings:
            return 0.0, "low", 1.0
        
        scores = [r.score for r in readings]
        confidences = [r.confidence for r in readings]
        triggered = [r for r in readings if r.is_triggered]
        
        if self.fusion_mode == FusionMode.AGREEMENT:
            # All must agree
            if len(triggered) == len(readings) and len(readings) > 0:
                fused = max(scores)
            else:
                fused = max(scores) * 0.5  # Dampen if not all agree
        
        elif self.fusion_mode == FusionMode.MAJORITY:
            # >50% must trigger
            if len(triggered) > len(readings) / 2:
                fused = sum(scores) / len(scores)
            else:
                fused = max(scores) * 0.3
        
        elif self.fusion_mode == FusionMode.ANY:
            # Any single trigger
            if triggered:
                fused = max(r.score for r in triggered)
            else:
                fused = sum(scores) / len(scores)
        
        else:  # WEIGHTED
            # Weighted linear combination
            fused = 0.0
            weight_sum = 0.0
            for r in readings:
                w = self.weights.get(r.detector_name, 0.25)
                fused += w * r.score
                weight_sum += w
            
            if weight_sum > 0:
                fused /= weight_sum
        
        # Confidence: inverse of detector disagreement
        if len(scores) > 1:
            score_variance = sum((s - sum(scores)/len(scores))**2 for s in scores) / len(scores)
            agreement = 1.0 - min(1.0, score_variance * 4.0)  # Scale: var=0.25 → agreement=0
        else:
            agreement = 1.0
        
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0.5
        overall_confidence = 0.5 * agreement + 0.5 * avg_confidence
        
        # Risk level from fused score
        risk = "low"
        for level, threshold in sorted(self.RISK_THRESHOLDS.items(), key=lambda x: x[1], reverse=True):
            if fused >= threshold:
                risk = level
                break
        
        return fused, risk, overall_confidence
    
    def evaluate_token(
        self,
        position: int,
        token_str: str,
        # EPR inputs
        topk_probs: Optional[List[float]] = None,
        topk_logprobs: Optional[List[float]] = None,
        token_id: int = 0,
        # Spilled energy inputs
        full_logits: Optional[List[float]] = None,
        sampled_token_id: int = 0,
        # CK-PLUG inputs
        p_query: Optional[List[float]] = None,
        p_rag: Optional[List[float]] = None,
        # TWAVE inputs
        probs_distribution: Optional[List[float]] = None,
        log_prob_policy: float = 0.0,
        log_prob_ref: float = 0.0,
        visual_attention: float = 1.0,
        prev_psi: float = 0.0,
    ) -> TokenVerdict:
        """
        Evaluate a single token with all enabled detectors.
        
        This is the core per-token evaluation function.
        """
        readings = []
        
        # EPR detector
        if self.epr and (topk_probs is not None or topk_logprobs is not None):
            epr_reading = self.epr.compute_token_entropy(
                token_id=token_id,
                topk_probs=topk_probs,
                topk_logprobs=topk_logprobs,
                token_str=token_str,
            )
            readings.append(DetectorReading(
                detector_name="epr",
                score=epr_reading.entropy_normalized,  # [0,1]
                confidence=epr_reading.confidence,
                is_triggered=epr_reading.is_anomaly,
                details={"entropy": epr_reading.entropy, "normalized": epr_reading.entropy_normalized},
            ))
        
        # Spilled energy detector
        if self.spilled and full_logits is not None:
            energy_reading = self.spilled.compute_spilled_energy(
                logits=full_logits,
                sampled_token_id=sampled_token_id,
                token_str=token_str,
            )
            # Normalize delta_E to [0,1] score
            normalized_delta = min(1.0, max(0.0, energy_reading.delta_E / 5.0))
            readings.append(DetectorReading(
                detector_name="spilled",
                score=normalized_delta,
                confidence=energy_reading.confidence,
                is_triggered=energy_reading.is_anomaly,
                details={"delta_E": energy_reading.delta_E, "E_ell": energy_reading.E_ell},
            ))
        
        # CK-PLUG detector
        if self.ckplug and p_query is not None and p_rag is not None:
            import numpy as np
            p_q = np.array(p_query) / sum(p_query)
            p_r = np.array(p_rag) / sum(p_rag)
            CG, H_para, H_cont = CKPLUGCoupling.confidence_gain(p_q, p_r)
            mu_ret = self.ckplug.compute_chemical_potential(p_q, p_r)
            
            # Negative CG = conflict = hallucination risk
            # Map CG to [0,1] score: CG=-0.5 → 1.0, CG=+0.5 → 0.0
            score = max(0.0, min(1.0, -CG * 2.0 + 0.5))
            readings.append(DetectorReading(
                detector_name="ckplug",
                score=score,
                confidence=0.7,  # CK-PLUG has moderate confidence
                is_triggered=CG < self.ckplug.epsilon,
                details={"CG": CG, "H_para": H_para, "H_cont": H_cont, "mu_ret": mu_ret},
            ))
        
        # TWAVE detector
        if self.twave and probs_distribution is not None:
            import numpy as np
            probs = np.array(probs_distribution) / sum(probs_distribution)
            CG = 0.0  # Neutral for standalone TWAVE
            
            state = self.twave.update_state(
                position=position,
                probs=probs,
                log_prob_policy=log_prob_policy,
                log_prob_ref=log_prob_ref,
                CG=CG,
                visual_attention=visual_attention,
                prev_psi=prev_psi,
            )
            
            stability = self.twave.evaluate_stability(state)
            # Map risk level to score
            risk_map = {
                "low": 0.0,
                "elevated": 0.3,
                "high": 0.6,
                "critical": 1.0,
            }
            score = risk_map.get(stability["hallucination_risk"], 0.0)
            
            readings.append(DetectorReading(
                detector_name="twave",
                score=score,
                confidence=0.6,  # TWAVE needs calibration
                is_triggered=stability["action"] != "continue",
                details={
                    "T_eff": state.temperature_eff,
                    "psi": state.psi,
                    "E_exc": state.E_excitation,
                    "action": stability["action"],
                },
            ))
        
        # Fuse
        fused, risk, confidence = self._fuse_scores(readings)
        action = self.ACTION_MAP.get(risk, Action.NONE)
        
        return TokenVerdict(
            position=position,
            token_str=token_str,
            readings=readings,
            fused_score=fused,
            risk_level=risk,
            recommended_action=action,
            confidence=confidence,
        )
    
    def evaluate_sequence(self, token_verdicts: List[TokenVerdict]) -> SequenceVerdict:
        """Evaluate an entire sequence and compute composite signals."""
        if not token_verdicts:
            return SequenceVerdict(
                token_verdicts=[],
                overall_risk="low",
                overall_action=Action.NONE,
                num_triggers=0,
                trigger_positions=[],
                avg_fused_score=0.0,
                max_fused_score=0.0,
                detector_agreement=0.0,
                energy_entropy_product=0.0,
                phase_transition_index=0.0,
                newi=0.0,
            )
        
        scores = [v.fused_score for v in token_verdicts]
        avg_score = sum(scores) / len(scores)
        max_score = max(scores)
        
        # Trigger positions
        trigger_positions = [v.position for v in token_verdicts if v.recommended_action != Action.NONE]
        
        # Detector agreement
        if token_verdicts and token_verdicts[0].readings:
            num_detectors = len(token_verdicts[0].readings)
            agreements = 0
            total = 0
            for v in token_verdicts:
                triggered = [r.detector_name for r in v.readings if r.is_triggered]
                if len(triggered) == 0 or len(triggered) == num_detectors:
                    agreements += 1
                total += 1
            agreement_rate = agreements / total if total > 0 else 0.0
        else:
            agreement_rate = 0.0
        
        # Composite novel signals
        # EEP: Energy-Entropy Product (correlated instability)
        epr_scores = []
        spilled_scores = []
        for v in token_verdicts:
            for r in v.readings:
                if r.detector_name == "epr":
                    epr_scores.append(r.score)
                elif r.detector_name == "spilled":
                    spilled_scores.append(r.score)
        
        if epr_scores and spilled_scores:
            eep = max(a * b for a, b in zip(epr_scores, spilled_scores))
        else:
            eep = 0.0
        
        # PTI: Phase Transition Index — divergence of order parameters
        # Approximated by entropy variance spikes
        if len(scores) > 3:
            mean_score = sum(scores) / len(scores)
            variance = sum((s - mean_score) ** 2 for s in scores) / len(scores)
            pti = min(1.0, variance * 4.0)
        else:
            pti = 0.0
        
        # NEWI: Non-Equilibrium Work Index
        # Approximated by cumulative fused score divergence from equilibrium
        newi = sum(abs(s - avg_score) for s in scores) / max(1, len(scores))
        
        # Overall risk: max of average and peak
        overall_score = max(avg_score, max_score * 0.7)
        overall_risk = "low"
        for level, threshold in sorted(self.RISK_THRESHOLDS.items(), key=lambda x: x[1], reverse=True):
            if overall_score >= threshold:
                overall_risk = level
                break
        
        overall_action = self.ACTION_MAP.get(overall_risk, Action.NONE)
        
        return SequenceVerdict(
            token_verdicts=token_verdicts,
            overall_risk=overall_risk,
            overall_action=overall_action,
            num_triggers=len(trigger_positions),
            trigger_positions=trigger_positions,
            avg_fused_score=avg_score,
            max_fused_score=max_score,
            detector_agreement=agreement_rate,
            energy_entropy_product=eep,
            phase_transition_index=pti,
            newi=newi,
        )
    
    def reset(self):
        """Reset all sub-detectors for new sequence."""
        if self.epr:
            self.epr.reset()
        if self.spilled:
            self.spilled.reset()
    
    def calibrate_weights(
        self,
        validation_sequences: List[Tuple[List[Dict], bool]],
    ) -> Dict[str, float]:
        """
        Calibrate detector weights on validation data.
        
        Args:
            validation_sequences: List of (token_readings_list, is_hallucination) tuples
        
        Returns:
            Optimized weight dictionary
        """
        # Simple grid search over weight combinations
        best_weights = self.weights.copy()
        best_f1 = 0.0
        
        # Grid: 0.1-0.4 for each detector
        import itertools
        
        for w_epr, w_spilled, w_ckplug, w_twave in itertools.product(
            [0.1, 0.2, 0.3, 0.4],
            [0.1, 0.2, 0.3, 0.4],
            [0.1, 0.2, 0.3, 0.4],
            [0.1, 0.2, 0.3, 0.4],
        ):
            total = w_epr + w_spilled + w_ckplug + w_twave
            if total == 0:
                continue
            
            weights = {
                "epr": w_epr / total,
                "spilled": w_spilled / total,
                "ckplug": w_ckplug / total,
                "twave": w_twave / total,
            }
            
            # Evaluate on validation set
            tp, fp, tn, fn = 0, 0, 0, 0
            for token_readings, is_hallucination in validation_sequences:
                # Simulate with these weights
                self.weights = weights
                readings = [
                    DetectorReading(
                        detector_name=r["detector"],
                        score=r["score"],
                        confidence=r["confidence"],
                        is_triggered=r["score"] > 0.5,
                    )
                    for r in token_readings
                ]
                fused, _, _ = self._fuse_scores(readings)
                predicted = fused > 0.5
                
                if predicted and is_hallucination:
                    tp += 1
                elif predicted and not is_hallucination:
                    fp += 1
                elif not predicted and not is_hallucination:
                    tn += 1
                else:
                    fn += 1
            
            precision = tp / (tp + fp) if (tp + fp) > 0 else 0
            recall = tp / (tp + fn) if (tp + fn) > 0 else 0
            f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
            
            if f1 > best_f1:
                best_f1 = f1
                best_weights = weights.copy()
        
        self.weights = best_weights
        return best_weights