streaming_intent/tracker.py

"""
Streaming Intent Tracker with HMM belief updating and SPRT stopping rules.

Maintains rolling text window, debounce logic, and statistically-principled
decision rules using Sequential Probability Ratio Test (SPRT).

Key features:
- Log-space HMM belief filtering (numerical stability)
- SPRT-based escalation decisions (explicit error rate bounds)
- Near-miss tracking for operational monitoring
- Ablation support for comparing decision strategies
"""

import time
import math
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Callable, Any, Tuple
from collections import deque
from enum import Enum

import numpy as np

from .config import StreamingConfig
from .belief import BeliefUpdater, EmissionTransform
from .metrics import MetricsCounter


class DecisionMode(Enum):
    """Decision rule strategies for ablation studies."""
    K_CONSECUTIVE = "k_consecutive"  # Original: K steps above threshold
    SPRT = "sprt"                    # Sequential Probability Ratio Test
    HYBRID = "hybrid"                # SPRT + K-consecutive fallback


@dataclass
class Chunk:
    """Input chunk from STT system."""
    text: str
    is_final: bool = False
    timestamp_ms: Optional[int] = None


@dataclass
class Decision:
    """Output decision from tracker."""
    current_belief: Dict[str, float]
    top_intent: str
    top_intent_prob: float
    should_escalate: bool
    should_commit: bool
    committed_intent: Optional[str]
    escalation_reason: Optional[str] = None
    step_latency_ms: float = 0.0
    step_count: int = 0
    window_text: str = ""
    # SPRT diagnostics
    sprt_llr: Optional[float] = None  # Log-likelihood ratio
    sprt_upper_bound: Optional[float] = None  # Upper decision boundary
    sprt_lower_bound: Optional[float] = None  # Lower decision boundary


class StreamingIntentTracker:
    """
    Streaming intent classifier with HMM-style belief updating and SPRT decisions.
    
    Features:
    - Rolling window of recent text (max_tokens configurable)
    - Debounce logic to prevent thrashing on partial updates
    - SPRT-based decision rules with explicit error rate bounds
    - Near-miss tracking for operational monitoring
    - Ablation support (use_hmm, decision_mode, emission_transform)
    
    SPRT (Sequential Probability Ratio Test):
        Tests H0: P(ESCALATION) = p0 vs H1: P(ESCALATION) = p1
        - Upper boundary A = log((1-beta)/alpha) → decide H1 (escalate)
        - Lower boundary B = log(beta/(1-alpha)) → decide H0 (don't escalate)
        - Continue if B < LLR < A
    
    Usage:
        tracker = StreamingIntentTracker(model_fn=my_inference)
        
        for chunk in stt_stream:
            decision = tracker.update(chunk)
            if decision.should_escalate:
                handle_escalation()
            elif decision.should_commit:
                handle_intent(decision.committed_intent)
    """
    
    def __init__(
        self,
        config: Optional[StreamingConfig] = None,
        model_fn: Optional[Callable[[str], Dict[str, float]]] = None,
        config_path: Optional[str] = None,
        use_hmm: bool = True,
        decision_mode: DecisionMode = DecisionMode.HYBRID,
        emission_transform: EmissionTransform = EmissionTransform.POWER,
        # SPRT parameter overrides (None = use config values)
        sprt_alpha: Optional[float] = None,
        sprt_beta: Optional[float] = None,
        sprt_p0: Optional[float] = None,
        sprt_p1: Optional[float] = None,
    ):
        """
        Initialize streaming intent tracker.
        
        Args:
            config: StreamingConfig instance (loads default if None)
            model_fn: Function that takes text and returns intent probabilities.
            config_path: Path to config YAML (used if config is None)
            use_hmm: If False, bypass HMM filtering (ablation mode).
            decision_mode: K_CONSECUTIVE, SPRT, or HYBRID.
            emission_transform: How to transform neural outputs.
            sprt_alpha: Target false escalation rate (overrides config).
            sprt_beta: Target missed escalation rate (overrides config).
            sprt_p0: Null hypothesis escalation probability (overrides config).
            sprt_p1: Alternative hypothesis escalation probability (overrides config).
        """
        if config is None:
            self.config = StreamingConfig.from_yaml(config_path)
        else:
            self.config = config
        
        self.model_fn = model_fn
        self.use_hmm = use_hmm
        self.decision_mode = decision_mode
        
        # Initialize belief updater with ablation options
        self.belief_updater = BeliefUpdater(
            self.config,
            use_hmm=use_hmm,
            emission_transform=emission_transform,
        )
        self.metrics = MetricsCounter()
        
        # SPRT parameters (use config values, allow overrides)
        self.sprt_alpha = sprt_alpha if sprt_alpha is not None else self.config.sprt_alpha
        self.sprt_beta = sprt_beta if sprt_beta is not None else self.config.sprt_beta
        self.sprt_p0 = sprt_p0 if sprt_p0 is not None else self.config.sprt_p0
        self.sprt_p1 = sprt_p1 if sprt_p1 is not None else self.config.sprt_p1
        
        # Compute SPRT boundaries (Wald's approximation)
        # A = log((1-beta)/alpha) - upper boundary → escalate
        # B = log(beta/(1-alpha)) - lower boundary → don't escalate
        self.sprt_upper_bound = math.log((1 - self.sprt_beta) / self.sprt_alpha)
        self.sprt_lower_bound = math.log(self.sprt_beta / (1 - self.sprt_alpha))
        
        # State
        self._text_buffer: List[str] = []
        self._window_text: str = ""
        self._top_intent_history: deque = deque(maxlen=self.config.K + 5)
        self._escalation_history: deque = deque(maxlen=self.config.K + 5)
        self._last_update_ms: int = 0
        self._last_text_hash: int = 0
        self._committed: bool = False
        self._committed_intent: Optional[str] = None
        self._escalated: bool = False
        
        # SPRT state
        self._sprt_llr: float = 0.0  # Cumulative log-likelihood ratio
        
        # Near-miss tracking
        self._peak_escalation_prob: float = 0.0
        self._near_miss_threshold: float = self.config.theta_med
    
    def reset(self) -> None:
        """Reset tracker state for new conversation."""
        # Track near-miss BEFORE reset (only if we actually processed data)
        had_data = self.belief_updater.step_count > 0
        if had_data:
            if self._peak_escalation_prob >= self._near_miss_threshold and not self._escalated:
                self.metrics.record_near_miss(self._peak_escalation_prob)
        
        # Now reset all state
        self.belief_updater.reset()
        self._text_buffer = []
        self._window_text = ""
        self._top_intent_history.clear()
        self._escalation_history.clear()
        self._last_update_ms = 0
        self._last_text_hash = 0
        self._committed = False
        self._committed_intent = None
        self._escalated = False
        
        # Reset SPRT state
        self._sprt_llr = 0.0
        self._peak_escalation_prob = 0.0
        
        self.metrics.start_session()
    
    def update(self, chunk: Chunk) -> Decision:
        """
        Process a chunk and return decision.
        
        Args:
            chunk: Input chunk with text, is_final flag, and optional timestamp.
            
        Returns:
            Decision with current belief, escalation status, and commitment status.
        """
        start_time = time.perf_counter()
        
        # Get current timestamp
        current_ms = chunk.timestamp_ms or int(time.time() * 1000)
        
        # Check debounce conditions
        if not self._should_update(chunk, current_ms):
            # Return current state without update
            top_intent, top_prob = self.belief_updater.get_top_intent()
            return Decision(
                current_belief=self.belief_updater.get_belief_dict(),
                top_intent=top_intent,
                top_intent_prob=top_prob,
                should_escalate=self._escalated,
                should_commit=self._committed,
                committed_intent=self._committed_intent,
                step_latency_ms=0.0,
                step_count=self.belief_updater.step_count,
                window_text=self._window_text,
            )
        
        # Update text buffer and window
        self._update_text_buffer(chunk)
        
        # Run model inference
        if self.model_fn is None:
            raise ValueError("model_fn must be set before calling update()")
        
        drivehealthbert_probs = self.model_fn(self._window_text)
        
        # Update belief state
        self.belief_updater.update(drivehealthbert_probs)
        
        # Track history
        top_intent, top_prob = self.belief_updater.get_top_intent()
        self._top_intent_history.append(top_intent)
        
        escalation_prob = self.belief_updater.get_intent_prob("ESCALATION")
        self._escalation_history.append(escalation_prob)
        
        # Update timestamps
        self._last_update_ms = current_ms
        self._last_text_hash = hash(self._window_text)
        
        # Make decisions
        should_escalate, escalation_reason = self._check_escalation()
        should_commit, committed_intent = self._check_commitment(top_intent, top_prob)
        
        # Update metrics
        self.metrics.record_step()
        if should_escalate and not self._escalated:
            self.metrics.record_escalation()
            self._escalated = True
        if should_commit and not self._committed:
            self.metrics.record_commitment(committed_intent)
            self._committed = True
            self._committed_intent = committed_intent
        
        # Check gray zone
        if not should_escalate and not should_commit:
            self.metrics.record_gray_zone()
        
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        
        return Decision(
            current_belief=self.belief_updater.get_belief_dict(),
            top_intent=top_intent,
            top_intent_prob=top_prob,
            should_escalate=should_escalate or self._escalated,
            should_commit=should_commit or self._committed,
            committed_intent=self._committed_intent if (should_commit or self._committed) else None,
            escalation_reason=escalation_reason,
            step_latency_ms=elapsed_ms,
            step_count=self.belief_updater.step_count,
            window_text=self._window_text,
            # SPRT diagnostics
            sprt_llr=self._sprt_llr,
            sprt_upper_bound=self.sprt_upper_bound,
            sprt_lower_bound=self.sprt_lower_bound,
        )
    
    def _should_update(self, chunk: Chunk, current_ms: int) -> bool:
        """Check if we should process this chunk (debounce logic)."""
        # Always update on final segments
        if chunk.is_final:
            return True
        
        # Check time-based debounce
        elapsed = current_ms - self._last_update_ms
        if elapsed < self.config.debounce_ms:
            return False
        
        # Check if text changed significantly
        new_text = self._compute_window_text(chunk.text)
        text_change = abs(len(new_text) - len(self._window_text))
        if text_change < self.config.min_change_chars:
            # Also check hash for content changes
            if hash(new_text) == self._last_text_hash:
                return False
        
        return True
    
    def _update_text_buffer(self, chunk: Chunk) -> None:
        """Update text buffer and compute rolling window."""
        if chunk.is_final:
            # Final segment - append to buffer
            self._text_buffer.append(chunk.text)
        else:
            # Partial - replace last entry if buffer not empty, otherwise append
            if self._text_buffer:
                self._text_buffer[-1] = chunk.text
            else:
                self._text_buffer.append(chunk.text)
        
        self._window_text = self._compute_window_text()
    
    def _compute_window_text(self, pending_text: str = "") -> str:
        """Compute rolling window text (last N tokens)."""
        # Combine buffer with pending text
        full_text = " ".join(self._text_buffer)
        if pending_text:
            full_text = full_text + " " + pending_text if full_text else pending_text
        
        # Simple token approximation: split on whitespace
        tokens = full_text.split()
        
        # Keep last max_tokens
        max_tokens = min(self.config.max_tokens, self.config.max_tokens_limit)
        if len(tokens) > max_tokens:
            tokens = tokens[-max_tokens:]
        
        return " ".join(tokens)
    
    def _update_sprt(self, escalation_prob: float) -> None:
        """
        Update SPRT log-likelihood ratio.
        
        LLR_n = LLR_{n-1} + log(P(x|H1) / P(x|H0))
        
        For continuous probability observations, we use a Bernoulli likelihood
        where the observation is treated as a soft indicator:
            LLR contribution = p * log(p1/p0) + (1-p) * log((1-p1)/(1-p0))
        
        This is equivalent to expected log-likelihood ratio under the
        observed probability distribution.
        """
        eps = 1e-10
        
        # Bernoulli log-likelihood ratio with soft observation
        # This correctly handles the full [0,1] range of escalation_prob
        log_ratio_escalate = math.log((self.sprt_p1 + eps) / (self.sprt_p0 + eps))
        log_ratio_no_escalate = math.log((1 - self.sprt_p1 + eps) / (1 - self.sprt_p0 + eps))
        
        # Expected LLR contribution = p * log(p1/p0) + (1-p) * log((1-p1)/(1-p0))
        llr_contribution = (
            escalation_prob * log_ratio_escalate +
            (1 - escalation_prob) * log_ratio_no_escalate
        )
        
        self._sprt_llr += llr_contribution
    
    def _check_escalation_sprt(self) -> Tuple[bool, Optional[str]]:
        """
        Check escalation using SPRT (Sequential Probability Ratio Test).
        
        Decision boundaries (Wald's approximation):
        - Upper: A = log((1-beta)/alpha) → decide H1 (escalate)
        - Lower: B = log(beta/(1-alpha)) → decide H0 (don't escalate)
        
        Returns:
            (should_escalate, reason)
        """
        if self._sprt_llr >= self.sprt_upper_bound:
            return True, f"sprt_upper_bound_crossed_llr_{self._sprt_llr:.3f}"
        
        # Note: we don't use lower bound to "commit" to non-escalation
        # because safety requires we keep checking until conversation ends
        return False, None
    
    def _check_escalation_k_consecutive(self, escalation_prob: float) -> Tuple[bool, Optional[str]]:
        """
        Check escalation using K-consecutive rule (original method).
        
        Rules:
        1. Immediate if belief[ESCALATION] >= theta_hi
        2. Escalate if belief[ESCALATION] >= theta_med for K consecutive steps
        
        Returns:
            (should_escalate, reason)
        """
        # Rule 1: Immediate escalation
        if escalation_prob >= self.config.theta_hi:
            return True, f"immediate_high_prob_{escalation_prob:.3f}"
        
        # Rule 2: Consecutive steps above theta_med
        if len(self._escalation_history) >= self.config.K:
            recent = list(self._escalation_history)[-self.config.K:]
            if all(p >= self.config.theta_med for p in recent):
                return True, f"consecutive_{self.config.K}_steps_above_theta_med"
        
        return False, None
    
    def _check_escalation(self) -> Tuple[bool, Optional[str]]:
        """
        Check escalation decision rules based on decision_mode.
        
        Modes:
        - K_CONSECUTIVE: Original threshold-based rules
        - SPRT: Sequential Probability Ratio Test only
        - HYBRID: SPRT with K-consecutive as fast-path fallback
        
        Returns:
            (should_escalate, reason)
        """
        if self._escalated:
            return True, "previously_escalated"
        
        escalation_prob = self.belief_updater.get_intent_prob("ESCALATION")
        
        # Track peak for near-miss detection
        self._peak_escalation_prob = max(self._peak_escalation_prob, escalation_prob)
        
        # Update SPRT state
        self._update_sprt(escalation_prob)
        
        if self.decision_mode == DecisionMode.K_CONSECUTIVE:
            return self._check_escalation_k_consecutive(escalation_prob)
        
        elif self.decision_mode == DecisionMode.SPRT:
            return self._check_escalation_sprt()
        
        elif self.decision_mode == DecisionMode.HYBRID:
            # Fast-path: immediate high probability (K-consecutive rule 1)
            if escalation_prob >= self.config.theta_hi:
                return True, f"hybrid_immediate_high_prob_{escalation_prob:.3f}"
            
            # SPRT for statistical rigor
            sprt_result, sprt_reason = self._check_escalation_sprt()
            if sprt_result:
                return True, f"hybrid_{sprt_reason}"
            
            # Fallback: K-consecutive as safety net
            k_result, k_reason = self._check_escalation_k_consecutive(escalation_prob)
            if k_result:
                return True, f"hybrid_{k_reason}"
            
            return False, None
        
        # Default fallback
        return self._check_escalation_k_consecutive(escalation_prob)
    
    def _check_commitment(self, top_intent: str, top_prob: float) -> Tuple[bool, Optional[str]]:
        """
        Check non-escalation intent commitment rules.
        
        Rules:
        - Commit only if top_prob >= theta_lock AND stable top intent for K steps
        
        Returns:
            (should_commit, committed_intent)
        """
        if self._committed:
            return True, self._committed_intent
        
        # Don't commit to ESCALATION through this path
        if top_intent == "ESCALATION":
            return False, None
        
        # Check probability threshold
        if top_prob < self.config.theta_lock:
            return False, None
        
        # Check stability: same top intent for K steps
        if len(self._top_intent_history) < self.config.K:
            return False, None
        
        recent = list(self._top_intent_history)[-self.config.K:]
        if all(intent == top_intent for intent in recent):
            return True, top_intent
        
        return False, None
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get current metrics."""
        return self.metrics.get_summary()
    
    def set_model_fn(self, model_fn: Callable[[str], Dict[str, float]]) -> None:
        """Set or update the model inference function."""
        self.model_fn = model_fn
    
    @property
    def window_text(self) -> str:
        """Current rolling window text."""
        return self._window_text
    
    @property
    def belief(self) -> Dict[str, float]:
        """Current belief state."""
        return self.belief_updater.get_belief_dict()
    
    @property
    def is_committed(self) -> bool:
        """Whether a final intent has been committed."""
        return self._committed
    
    @property
    def is_escalated(self) -> bool:
        """Whether escalation has been triggered."""
        return self._escalated