drift_detection.py

"""
Drift Detection Module

Statistical drift detection for monitoring metrics.
Implements rolling window analysis and threshold-based alerting.
"""

import uuid
from collections import deque
from dataclasses import dataclass, field
from datetime import datetime
from typing import Deque, Dict, List, Optional

from backend.logging.logger import get_logger

from .schemas import (
    AlertSeverity,
    AlertType,
    DriftDetectionResult,
    MonitoringConfig,
    RollingMetrics,
)


@dataclass
class MetricWindow:
    """Rolling window for a specific metric."""

    metric_name: str
    values: Deque[float] = field(default_factory=deque)
    timestamps: Deque[datetime] = field(default_factory=deque)
    window_size: int = 100
    max_stored_events: int = 10000  # Storage control: max events to retain

    def add(self, value: float, timestamp: datetime) -> None:
        """Add a new value to the window."""
        self.values.append(value)
        self.timestamps.append(timestamp)
        
        # Maintain window size for rolling calculations
        while len(self.values) > self.window_size:
            self.values.popleft()
            self.timestamps.popleft()
        
        # Storage control: Prune if exceeding max stored events
        self._prune_if_needed()

    def _prune_if_needed(self) -> None:
        """Prune oldest events if exceeding max storage limit."""
        if len(self.values) > self.max_stored_events:
            # Keep only the most recent max_stored_events
            excess = len(self.values) - self.max_stored_events
            for _ in range(excess):
                self.values.popleft()
                self.timestamps.popleft()
            
    def get_storage_stats(self) -> dict:
        """Get storage statistics for this window."""
        return {
            "metric_name": self.metric_name,
            "current_size": len(self.values),
            "max_stored_events": self.max_stored_events,
            "storage_usage_pct": (len(self.values) / self.max_stored_events) * 100,
        }
    
    def clear(self) -> None:
        """Clear all stored data."""
        self.values.clear()
        self.timestamps.clear()

    def compute_metrics(self) -> RollingMetrics:
        """Compute rolling statistics."""
        if not self.values:
            return RollingMetrics(
                metric_name=self.metric_name,
                current_value=0.0,
                window_size=self.window_size,
                sample_count=0,
                min_value=0.0,
                max_value=0.0,
                std_dev=0.0,
            )
        
        values_list = list(self.values)
        n = len(values_list)
        
        # Mean
        current_value = sum(values_list) / n
        
        # Min/Max
        min_value = min(values_list)
        max_value = max(values_list)
        
        # Standard deviation
        variance = sum((x - current_value) ** 2 for x in values_list) / n
        std_dev = variance ** 0.5
        
        return RollingMetrics(
            metric_name=self.metric_name,
            current_value=current_value,
            window_size=self.window_size,
            sample_count=n,
            min_value=min_value,
            max_value=max_value,
            std_dev=std_dev,
        )


class DriftDetector:
    """
    Drift detection engine for monitoring metrics.
    
    Implements:
    - Statistical drift detection (baseline vs live)
    - Confidence collapse detection
    - Threshold-based alerting
    
    Mathematical definitions:
    - Drift(H) = |mean(H_live) - mean(H_baseline)|
    - Alert if Drift(metric) > threshold
    """

    def __init__(
        self,
        config: Optional[MonitoringConfig] = None,
    ) -> None:
        """
        Initialize drift detector.
        
        Args:
            config: Monitoring configuration
        """
        self.logger = get_logger(__name__)
        self._config = config or MonitoringConfig()
        
        # Baseline windows (fixed reference)
        self._baseline_windows: Dict[str, MetricWindow] = {}
        
        # Live rolling windows
        self._live_windows: Dict[str, MetricWindow] = {
            "hallucination": MetricWindow("hallucination", window_size=self._config.window_size),
            "toxicity": MetricWindow("toxicity", window_size=self._config.window_size),
            "bias": MetricWindow("bias", window_size=self._config.window_size),
            "confidence": MetricWindow("confidence", window_size=self._config.window_size),
            "robustness": MetricWindow("robustness", window_size=self._config.window_size),
        }
        
        # Baseline values (established during initial period)
        self._baseline_values: Dict[str, float] = {}
        self._baseline_established: bool = False

    def update_baseline(self, baseline_values: Dict[str, float]) -> None:
        """
        Update baseline values for drift detection.
        
        Args:
            baseline_values: Dictionary of baseline metric values
        """
        self.logger.info("Updating baseline values", baseline_values=baseline_values)
        
        for metric_name, value in baseline_values.items():
            if metric_name in self._baseline_windows:
                # Add to baseline window
                self._baseline_windows[metric_name].add(value, datetime.utcnow())
            else:
                # Create new baseline window
                window = MetricWindow(metric_name, window_size=self._config.window_size)
                window.add(value, datetime.utcnow())
                self._baseline_windows[metric_name] = window
            
            self._baseline_values[metric_name] = value
        
        self._baseline_established = True
        self.logger.info("Baseline established", metrics=list(self._baseline_values.keys()))

    def record_metric(
        self,
        metric_name: str,
        value: float,
        timestamp: Optional[datetime] = None,
    ) -> None:
        """
        Record a new metric value.
        
        Args:
            metric_name: Name of the metric
            value: Metric value
            timestamp: Timestamp (defaults to now)
        """
        if timestamp is None:
            timestamp = datetime.utcnow()
        
        if metric_name in self._live_windows:
            self._live_windows[metric_name].add(value, timestamp)
            
            # Auto-establish baseline from first values if not set
            if not self._baseline_established and metric_name not in self._baseline_values:
                if self._live_windows[metric_name].compute_metrics().sample_count >= self._config.min_window_samples:
                    # Use initial rolling average as baseline
                    rolling = self._live_windows[metric_name].compute_metrics()
                    self._baseline_values[metric_name] = rolling.current_value
                    self.logger.info(f"Auto-established baseline for {metric_name}", baseline=rolling.current_value)
                    
                    # Check if all baselines are established
                    if all(m in self._baseline_values for m in self._live_windows.keys()):
                        self._baseline_established = True

    def detect_drift(self, metric_name: str) -> DriftDetectionResult:
        """
        Detect drift for a specific metric.
        
        Args:
            metric_name: Name of the metric to check
            
        Returns:
            Drift detection result
        """
        if metric_name not in self._live_windows:
            return DriftDetectionResult(
                metric_name=metric_name,
                baseline_value=0.0,
                live_value=0.0,
                drift_magnitude=0.0,
                threshold=0.0,
                is_drift_detected=False,
                severity=AlertSeverity.LOW,
            )
        
        # Get baseline value
        baseline_value = self._baseline_values.get(metric_name, 0.0)
        
        # Get live rolling metrics
        live_window = self._live_windows[metric_name]
        live_metrics = live_window.compute_metrics()
        
        if live_metrics.sample_count < self._config.min_window_samples:
            # Not enough samples yet
            return DriftDetectionResult(
                metric_name=metric_name,
                baseline_value=baseline_value,
                live_value=live_metrics.current_value,
                drift_magnitude=0.0,
                threshold=self._get_threshold(metric_name),
                is_drift_detected=False,
                severity=AlertSeverity.LOW,
            )
        
        # Calculate drift magnitude
        drift_magnitude = abs(live_metrics.current_value - baseline_value)
        
        # Get threshold for this metric
        threshold = self._get_threshold(metric_name)
        
        # Determine if drift exceeds threshold
        is_drift_detected = drift_magnitude > threshold
        
        # Calculate severity
        severity = self._calculate_severity(drift_magnitude, threshold)
        
        result = DriftDetectionResult(
            metric_name=metric_name,
            baseline_value=baseline_value,
            live_value=live_metrics.current_value,
            drift_magnitude=drift_magnitude,
            threshold=threshold,
            is_drift_detected=is_drift_detected,
            severity=severity,
        )
        
        if is_drift_detected:
            self.logger.warning(
                "Drift detected",
                metric_name=metric_name,
                baseline_value=baseline_value,
                live_value=live_metrics.current_value,
                drift_magnitude=drift_magnitude,
                threshold=threshold,
                severity=severity,
            )
        
        return result

    def detect_all_drift(self) -> Dict[str, DriftDetectionResult]:
        """
        Detect drift for all metrics.
        
        Returns:
            Dictionary of drift detection results by metric
        """
        results = {}
        
        for metric_name in self._live_windows.keys():
            results[metric_name] = self.detect_drift(metric_name)
        
        return results

    def get_rolling_metrics(self, metric_name: str) -> Optional[RollingMetrics]:
        """
        Get rolling metrics for a specific metric.
        
        Args:
            metric_name: Name of the metric
            
        Returns:
            Rolling metrics or None if not found
        """
        if metric_name in self._live_windows:
            return self._live_windows[metric_name].compute_metrics()
        return None

    def get_all_rolling_metrics(self) -> Dict[str, RollingMetrics]:
        """
        Get rolling metrics for all metrics.
        
        Returns:
            Dictionary of rolling metrics by metric name
        """
        return {
            name: window.compute_metrics()
            for name, window in self._live_windows.items()
        }

    def get_trend_data(self, metric_name: str, limit: int = 50) -> List[float]:
        """
        Get recent trend data for a metric.
        
        Args:
            metric_name: Name of the metric
            limit: Maximum number of values to return
            
        Returns:
            List of recent values
        """
        if metric_name not in self._live_windows:
            return []
        
        values = list(self._live_windows[metric_name].values)
        return values[-limit:] if len(values) > limit else values

    def _get_threshold(self, metric_name: str) -> float:
        """Get threshold for a specific metric."""
        threshold_map = {
            "hallucination": self._config.hallucination_threshold,
            "toxicity": self._config.toxicity_threshold,
            "bias": self._config.bias_threshold,
            "confidence": self._config.confidence_threshold,
            "robustness": self._config.robustness_threshold,
        }
        return threshold_map.get(metric_name, 0.1)

    def _calculate_severity(self, drift_magnitude: float, threshold: float) -> AlertSeverity:
        """Calculate alert severity based on drift magnitude vs threshold."""
        if threshold <= 0:
            return AlertSeverity.LOW
        
        ratio = drift_magnitude / threshold
        
        if ratio > 3.0:
            return AlertSeverity.CRITICAL
        elif ratio > 2.0:
            return AlertSeverity.HIGH
        elif ratio > 1.5:
            return AlertSeverity.MEDIUM
        else:
            return AlertSeverity.LOW

    def check_confidence_collapse(self) -> Optional[DriftDetectionResult]:
        """
        Check for confidence collapse.
        
        Returns:
            Drift detection result if collapse detected, None otherwise
        """
        # Confidence collapse is when live confidence drops below threshold
        confidence_rolling = self.get_rolling_metrics("confidence")
        
        if confidence_rolling is None or confidence_rolling.sample_count < self._config.min_window_samples:
            return None
        
        baseline_confidence = self._baseline_values.get("confidence", 0.5)
        threshold = self._config.confidence_threshold
        
        # Check if confidence has collapsed (dropped by more than threshold)
        collapse_magnitude = baseline_confidence - confidence_rolling.current_value
        
        if collapse_magnitude > threshold:
            severity = self._calculate_severity(collapse_magnitude, threshold)
            
            return DriftDetectionResult(
                metric_name="confidence",
                baseline_value=baseline_confidence,
                live_value=confidence_rolling.current_value,
                drift_magnitude=collapse_magnitude,
                threshold=threshold,
                is_drift_detected=True,
                severity=severity,
            )
        
        return None


# Global detector instance
_drift_detector: Optional[DriftDetector] = None


def get_drift_detector(config: Optional[MonitoringConfig] = None) -> DriftDetector:
    """
    Get the global drift detector instance.
    
    Args:
        config: Optional monitoring configuration
        
    Returns:
        DriftDetector singleton
    """
    global _drift_detector
    if _drift_detector is None:
        _drift_detector = DriftDetector(config=config)
    return _drift_detector


__all__ = [
    "DriftDetector",
    "MetricWindow",
    "DriftDetectionResult",
    "get_drift_detector",
]