Chapter 19: Human-AI Collaboration in Clinical Practice

"""
Production monitoring system for clinical AI with comprehensive equity evaluation.

This module implements a monitoring framework that tracks model performance
across multiple dimensions, detects performance degradation and fairness issues,
and generates alerts when intervention is required. The system is designed
specifically for healthcare applications where equity monitoring is essential.
"""

import logging
from typing import Dict, List, Optional, Tuple, Any, Set
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from pathlib import Path
from enum import Enum
import numpy as np
import pandas as pd
from scipy import stats
from sklearn.metrics import (
    roc_auc_score, average_precision_score, brier_score_loss,
    confusion_matrix, log_loss
)
from sklearn.calibration import calibration_curve
import json
import warnings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AlertSeverity(Enum):
    """Severity levels for monitoring alerts."""
    INFO = "info"
    WARNING = "warning"
    CRITICAL = "critical"
    EMERGENCY = "emergency"

class PerformanceTrend(Enum):
    """Trend direction for performance metrics."""
    IMPROVING = "improving"
    STABLE = "stable"
    DEGRADING = "degrading"
    INSUFFICIENT_DATA = "insufficient_data"

@dataclass
class PerformanceMetrics:
    """
    Container for model performance metrics.

    Attributes:
        n_samples: Number of predictions evaluated
        n_positive: Number of positive outcomes
        auc: Area under ROC curve
        average_precision: Average precision score
        brier_score: Brier calibration score
        sensitivity: True positive rate
        specificity: True negative rate
        ppv: Positive predictive value
        npv: Negative predictive value
        calibration_slope: Slope of calibration curve
        calibration_intercept: Intercept of calibration curve
        ece: Expected calibration error
        timestamp: When metrics were computed
    """
    n_samples: int
    n_positive: int
    auc: float
    average_precision: float
    brier_score: float
    sensitivity: float
    specificity: float
    ppv: float
    npv: float
    calibration_slope: float
    calibration_intercept: float
    ece: float
    timestamp: datetime

    def to_dict(self) -> Dict[str, Any]:
        """Convert metrics to dictionary for serialization."""
        return {
            'n_samples': self.n_samples,
            'n_positive': self.n_positive,
            'auc': float(self.auc) if not np.isnan(self.auc) else None,
            'average_precision': float(self.average_precision) if not np.isnan(self.average_precision) else None,
            'brier_score': float(self.brier_score),
            'sensitivity': float(self.sensitivity),
            'specificity': float(self.specificity),
            'ppv': float(self.ppv) if not np.isnan(self.ppv) else None,
            'npv': float(self.npv) if not np.isnan(self.npv) else None,
            'calibration_slope': float(self.calibration_slope),
            'calibration_intercept': float(self.calibration_intercept),
            'ece': float(self.ece),
            'timestamp': self.timestamp.isoformat()
        }

@dataclass
class FairnessMetrics:
    """
    Container for fairness metrics comparing groups.

    Attributes:
        reference_group: Name of reference group for comparison
        comparison_groups: Names of groups being compared to reference
        demographic_parity_ratio: Ratio of selection rates
        equal_opportunity_difference: Difference in true positive rates
        predictive_parity_difference: Difference in positive predictive values
        calibration_disparity: Maximum calibration difference across groups
        auc_disparity: Maximum AUC difference across groups
        timestamp: When metrics were computed
    """
    reference_group: str
    comparison_groups: List[str]
    demographic_parity_ratio: Dict[str, float]
    equal_opportunity_difference: Dict[str, float]
    predictive_parity_difference: Dict[str, float]
    calibration_disparity: float
    auc_disparity: float
    timestamp: datetime

    def to_dict(self) -> Dict[str, Any]:
        """Convert metrics to dictionary for serialization."""
        return {
            'reference_group': self.reference_group,
            'comparison_groups': self.comparison_groups,
            'demographic_parity_ratio': self.demographic_parity_ratio,
            'equal_opportunity_difference': self.equal_opportunity_difference,
            'predictive_parity_difference': self.predictive_parity_difference,
            'calibration_disparity': float(self.calibration_disparity),
            'auc_disparity': float(self.auc_disparity),
            'timestamp': self.timestamp.isoformat()
        }

@dataclass
class MonitoringAlert:
    """
    Alert generated by monitoring system.

    Attributes:
        alert_id: Unique identifier
        severity: Alert severity level
        alert_type: Type of issue detected
        affected_groups: Demographic groups affected
        metric_name: Performance metric that triggered alert
        threshold: Threshold that was exceeded
        observed_value: Actual observed value
        message: Human-readable alert description
        timestamp: When alert was generated
        resolution_status: Whether alert has been addressed
        resolution_notes: Notes on how alert was resolved
    """
    alert_id: str
    severity: AlertSeverity
    alert_type: str
    affected_groups: List[str]
    metric_name: str
    threshold: float
    observed_value: float
    message: str
    timestamp: datetime
    resolution_status: str = "unresolved"
    resolution_notes: str = ""

    def to_dict(self) -> Dict[str, Any]:
        """Convert alert to dictionary for serialization."""
        return {
            'alert_id': self.alert_id,
            'severity': self.severity.value,
            'alert_type': self.alert_type,
            'affected_groups': self.affected_groups,
            'metric_name': self.metric_name,
            'threshold': float(self.threshold),
            'observed_value': float(self.observed_value),
            'message': self.message,
            'timestamp': self.timestamp.isoformat(),
            'resolution_status': self.resolution_status,
            'resolution_notes': self.resolution_notes
        }

class StratifiedPerformanceMonitor:
    """
    Comprehensive performance monitoring system with equity focus.

    This class implements stratified performance monitoring that tracks
    metrics across demographic subgroups, detects performance degradation
    and fairness issues, and generates appropriate alerts.
    """

    def __init__(
        self,
        model_name: str,
        stratification_variables: List[str],
        performance_thresholds: Dict[str, float],
        fairness_thresholds: Dict[str, float],
        min_sample_size: int = 30,
        monitoring_window_days: int = 7,
        storage_path: Optional[Path] = None
    ):
        """
        Initialize monitoring system.

        Args:
            model_name: Name of model being monitored
            stratification_variables: Demographic variables for stratification
            performance_thresholds: Minimum acceptable performance levels
            fairness_thresholds: Maximum acceptable fairness disparities
            min_sample_size: Minimum samples needed for reliable metrics
            monitoring_window_days: Days to aggregate for time series
            storage_path: Path for persisting monitoring data
        """
        self.model_name = model_name
        self.stratification_variables = stratification_variables
        self.performance_thresholds = performance_thresholds
        self.fairness_thresholds = fairness_thresholds
        self.min_sample_size = min_sample_size
        self.monitoring_window_days = monitoring_window_days
        self.storage_path = storage_path or Path("./monitoring_data")
        self.storage_path.mkdir(parents=True, exist_ok=True)

        # Monitoring state
        self.predictions_buffer: List[Dict] = []
        self.performance_history: Dict[str, List[PerformanceMetrics]] = {}
        self.fairness_history: List[FairnessMetrics] = []
        self.active_alerts: List[MonitoringAlert] = []
        self.alert_counter = 0

        logger.info(
            f"Initialized StratifiedPerformanceMonitor for {model_name} "
            f"with {len(stratification_variables)} stratification variables"
        )

    def record_prediction(
        self,
        prediction_id: str,
        features: Dict[str, Any],
        predicted_probability: float,
        predicted_class: int,
        demographic_attributes: Dict[str, Any],
        timestamp: Optional[datetime] = None
    ) -> None:
        """
        Record a model prediction for later evaluation.

        Args:
            prediction_id: Unique identifier for this prediction
            features: Feature values used for prediction
            predicted_probability: Model's predicted probability
            predicted_class: Model's predicted class
            demographic_attributes: Patient demographic information
            timestamp: Time of prediction (defaults to current time)
        """
        if timestamp is None:
            timestamp = datetime.now()

        prediction_record = {
            'prediction_id': prediction_id,
            'predicted_probability': predicted_probability,
            'predicted_class': predicted_class,
            'demographic_attributes': demographic_attributes,
            'timestamp': timestamp,
            'outcome_recorded': False,
            'true_outcome': None
        }

        self.predictions_buffer.append(prediction_record)

        # Persist to storage
        self._persist_prediction(prediction_record)

    def record_outcome(
        self,
        prediction_id: str,
        true_outcome: int,
        outcome_timestamp: Optional[datetime] = None
    ) -> None:
        """
        Record the true outcome for a previous prediction.

        Args:
            prediction_id: ID of the prediction to update
            true_outcome: Actual observed outcome
            outcome_timestamp: Time outcome was observed
        """
        # Find prediction in buffer
        for record in self.predictions_buffer:
            if record['prediction_id'] == prediction_id:
                record['true_outcome'] = true_outcome
                record['outcome_recorded'] = True
                if outcome_timestamp:
                    record['outcome_timestamp'] = outcome_timestamp
                break

        # Persist outcome
        self._persist_outcome(prediction_id, true_outcome, outcome_timestamp)

    def compute_metrics(
        self,
        y_true: np.ndarray,
        y_pred_proba: np.ndarray,
        y_pred_class: np.ndarray,
        timestamp: Optional[datetime] = None
    ) -> PerformanceMetrics:
        """
        Compute comprehensive performance metrics.

        Args:
            y_true: True binary labels
            y_pred_proba: Predicted probabilities
            y_pred_class: Predicted classes
            timestamp: Time metrics were computed

        Returns:
            PerformanceMetrics object
        """
        if timestamp is None:
            timestamp = datetime.now()

        n_samples = len(y_true)
        n_positive = int(np.sum(y_true))

        # Handle edge cases
        if n_samples < self.min_sample_size:
            logger.warning(
                f"Sample size ({n_samples}) below minimum ({self.min_sample_size})"
            )

        if n_positive == 0 or n_positive == n_samples:
            logger.warning("All samples have same class, some metrics will be undefined")

        # Compute discrimination metrics
        try:
            auc = roc_auc_score(y_true, y_pred_proba)
            average_precision = average_precision_score(y_true, y_pred_proba)
        except Exception as e:
            logger.error(f"Error computing discrimination metrics: {e}")
            auc = np.nan
            average_precision = np.nan

        # Compute calibration metrics
        brier_score = brier_score_loss(y_true, y_pred_proba)

        try:
            # Fit logistic calibration to get slope and intercept
            from sklearn.linear_model import LogisticRegression
            cal_model = LogisticRegression()
            # Use log-odds of predictions as feature
            log_odds = np.log(y_pred_proba / (1 - y_pred_proba + 1e-10))
            cal_model.fit(log_odds.reshape(-1, 1), y_true)
            calibration_slope = cal_model.coef_[0][0]
            calibration_intercept = cal_model.intercept_[0]
        except Exception as e:
            logger.warning(f"Could not compute calibration slope: {e}")
            calibration_slope = np.nan
            calibration_intercept = np.nan

        # Compute expected calibration error
        try:
            ece = self._compute_ece(y_true, y_pred_proba, n_bins=10)
        except Exception as e:
            logger.warning(f"Could not compute ECE: {e}")
            ece = np.nan

        # Compute confusion matrix metrics
        try:
            tn, fp, fn, tp = confusion_matrix(y_true, y_pred_class).ravel()
            sensitivity = tp / (tp + fn) if (tp + fn) > 0 else np.nan
            specificity = tn / (tn + fp) if (tn + fp) > 0 else np.nan
            ppv = tp / (tp + fp) if (tp + fp) > 0 else np.nan
            npv = tn / (tn + fn) if (tn + fn) > 0 else np.nan
        except Exception as e:
            logger.error(f"Error computing confusion matrix metrics: {e}")
            sensitivity = specificity = ppv = npv = np.nan

        return PerformanceMetrics(
            n_samples=n_samples,
            n_positive=n_positive,
            auc=auc,
            average_precision=average_precision,
            brier_score=brier_score,
            sensitivity=sensitivity,
            specificity=specificity,
            ppv=ppv,
            npv=npv,
            calibration_slope=calibration_slope,
            calibration_intercept=calibration_intercept,
            ece=ece,
            timestamp=timestamp
        )

    def _compute_ece(
        self,
        y_true: np.ndarray,
        y_pred_proba: np.ndarray,
        n_bins: int = 10
    ) -> float:
        """
        Compute expected calibration error.

        Args:
            y_true: True labels
            y_pred_proba: Predicted probabilities
            n_bins: Number of bins for calibration

        Returns:
            Expected calibration error
        """
        bin_boundaries = np.linspace(0, 1, n_bins + 1)
        bin_lowers = bin_boundaries[:-1]
        bin_uppers = bin_boundaries[1:]

        ece = 0.0
        for bin_lower, bin_upper in zip(bin_lowers, bin_uppers):
            in_bin = (y_pred_proba >= bin_lower) & (y_pred_proba < bin_upper)
            prop_in_bin = np.mean(in_bin)

            if prop_in_bin > 0:
                accuracy_in_bin = np.mean(y_true[in_bin])
                avg_confidence_in_bin = np.mean(y_pred_proba[in_bin])
                ece += np.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin

        return ece

    def evaluate_monitoring_window(
        self,
        window_end: Optional[datetime] = None,
        window_days: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        Evaluate performance over monitoring window with stratification.

        Args:
            window_end: End of evaluation window (defaults to now)
            window_days: Length of window (defaults to monitoring_window_days)

        Returns:
            Dictionary containing overall and stratified metrics
        """
        if window_end is None:
            window_end = datetime.now()
        if window_days is None:
            window_days = self.monitoring_window_days

        window_start = window_end - timedelta(days=window_days)

        # Filter to completed predictions in window
        window_predictions = [
            p for p in self.predictions_buffer
            if p['outcome_recorded'] and
               window_start <= p['timestamp'] <= window_end
        ]

        if len(window_predictions) == 0:
            logger.warning("No completed predictions in monitoring window")
            return {'error': 'No data available'}

        logger.info(
            f"Evaluating {len(window_predictions)} predictions from "
            f"{window_start.date()} to {window_end.date()}"
        )

        # Convert to arrays for metric computation
        y_true = np.array([p['true_outcome'] for p in window_predictions])
        y_pred_proba = np.array([p['predicted_probability'] for p in window_predictions])
        y_pred_class = np.array([p['predicted_class'] for p in window_predictions])

        # Compute overall metrics
        overall_metrics = self.compute_metrics(
            y_true, y_pred_proba, y_pred_class, timestamp=window_end
        )

        # Store in history
        if 'overall' not in self.performance_history:
            self.performance_history['overall'] = []
        self.performance_history['overall'].append(overall_metrics)

        # Compute stratified metrics
        stratified_metrics = {}
        for var in self.stratification_variables:
            # Get unique values for this stratification variable
            var_values = set()
            for p in window_predictions:
                if var in p['demographic_attributes']:
                    var_values.add(p['demographic_attributes'][var])

            stratified_metrics[var] = {}
            for value in var_values:
                # Filter to this subgroup
                subgroup_indices = [
                    i for i, p in enumerate(window_predictions)
                    if p['demographic_attributes'].get(var) == value
                ]

                if len(subgroup_indices) < self.min_sample_size:
                    logger.warning(
                        f"Insufficient samples for {var}={value}: "
                        f"{len(subgroup_indices)} < {self.min_sample_size}"
                    )
                    continue

                subgroup_y_true = y_true[subgroup_indices]
                subgroup_y_pred_proba = y_pred_proba[subgroup_indices]
                subgroup_y_pred_class = y_pred_class[subgroup_indices]

                group_metrics = self.compute_metrics(
                    subgroup_y_true,
                    subgroup_y_pred_proba,
                    subgroup_y_pred_class,
                    timestamp=window_end
                )

                group_key = f"{var}_{value}"
                stratified_metrics[var][value] = group_metrics

                # Store in history
                if group_key not in self.performance_history:
                    self.performance_history[group_key] = []
                self.performance_history[group_key].append(group_metrics)

        # Compute fairness metrics
        fairness_metrics = self._compute_fairness_metrics(
            window_predictions, stratified_metrics, window_end
        )
        if fairness_metrics:
            self.fairness_history.append(fairness_metrics)

        # Check for alerts
        alerts = self._check_for_alerts(
            overall_metrics, stratified_metrics, fairness_metrics
        )

        if alerts:
            logger.warning(f"Generated {len(alerts)} alerts")
            self.active_alerts.extend(alerts)

        return {
            'window_start': window_start,
            'window_end': window_end,
            'n_predictions': len(window_predictions),
            'overall_metrics': overall_metrics,
            'stratified_metrics': stratified_metrics,
            'fairness_metrics': fairness_metrics,
            'alerts': alerts
        }

    def _compute_fairness_metrics(
        self,
        predictions: List[Dict],
        stratified_metrics: Dict[str, Dict],
        timestamp: datetime
    ) -> Optional[FairnessMetrics]:
        """
        Compute fairness metrics comparing groups.

        Args:
            predictions: List of prediction records
            stratified_metrics: Metrics stratified by demographic variables
            timestamp: Time of evaluation

        Returns:
            FairnessMetrics object or None if insufficient data
        """
        # For simplicity, use first stratification variable as primary
        if not self.stratification_variables or not stratified_metrics:
            return None

        primary_var = self.stratification_variables[0]
        if primary_var not in stratified_metrics:
            return None

        groups = list(stratified_metrics[primary_var].keys())
        if len(groups) < 2:
            return None

        # Use most prevalent group as reference
        group_sizes = {}
        for group in groups:
            group_sizes[group] = stratified_metrics[primary_var][group].n_samples

        reference_group = max(group_sizes, key=group_sizes.get)
        comparison_groups = [g for g in groups if g != reference_group]

        ref_metrics = stratified_metrics[primary_var][reference_group]

        # Compute parity metrics
        demographic_parity = {}
        equal_opportunity_diff = {}
        predictive_parity_diff = {}

        auc_values = [ref_metrics.auc]
        calibration_slopes = [ref_metrics.calibration_slope]

        for group in comparison_groups:
            group_metrics = stratified_metrics[primary_var][group]

            # Demographic parity: ratio of positive prediction rates
            ref_positive_rate = np.sum([
                p['predicted_class'] for p in predictions
                if p['demographic_attributes'].get(primary_var) == reference_group
            ]) / group_sizes[reference_group]

            group_positive_rate = np.sum([
                p['predicted_class'] for p in predictions
                if p['demographic_attributes'].get(primary_var) == group
            ]) / group_sizes[group]

            demographic_parity[group] = group_positive_rate / (ref_positive_rate + 1e-10)

            # Equal opportunity: difference in true positive rates (sensitivity)
            equal_opportunity_diff[group] = group_metrics.sensitivity - ref_metrics.sensitivity

            # Predictive parity: difference in positive predictive values
            if not np.isnan(group_metrics.ppv) and not np.isnan(ref_metrics.ppv):
                predictive_parity_diff[group] = group_metrics.ppv - ref_metrics.ppv
            else:
                predictive_parity_diff[group] = np.nan

            # Collect for disparity calculation
            if not np.isnan(group_metrics.auc):
                auc_values.append(group_metrics.auc)
            if not np.isnan(group_metrics.calibration_slope):
                calibration_slopes.append(group_metrics.calibration_slope)

        # Compute maximum disparities
        auc_disparity = max(auc_values) - min(auc_values) if len(auc_values) > 1 else 0.0
        calibration_disparity = (
            max(calibration_slopes) - min(calibration_slopes)
            if len(calibration_slopes) > 1 else 0.0
        )

        return FairnessMetrics(
            reference_group=f"{primary_var}_{reference_group}",
            comparison_groups=[f"{primary_var}_{g}" for g in comparison_groups],
            demographic_parity_ratio=demographic_parity,
            equal_opportunity_difference=equal_opportunity_diff,
            predictive_parity_difference=predictive_parity_diff,
            calibration_disparity=calibration_disparity,
            auc_disparity=auc_disparity,
            timestamp=timestamp
        )

    def _check_for_alerts(
        self,
        overall_metrics: PerformanceMetrics,
        stratified_metrics: Dict[str, Dict],
        fairness_metrics: Optional[FairnessMetrics]
    ) -> List[MonitoringAlert]:
        """
        Check metrics against thresholds and generate alerts.

        Args:
            overall_metrics: Overall performance metrics
            stratified_metrics: Stratified performance metrics
            fairness_metrics: Fairness metrics across groups

        Returns:
            List of alerts generated
        """
        alerts = []

        # Check overall performance thresholds
        for metric_name, threshold in self.performance_thresholds.items():
            if hasattr(overall_metrics, metric_name):
                value = getattr(overall_metrics, metric_name)
                if not np.isnan(value) and value < threshold:
                    alert = self._create_alert(
                        severity=AlertSeverity.WARNING,
                        alert_type="overall_performance_degradation",
                        affected_groups=["overall"],
                        metric_name=metric_name,
                        threshold=threshold,
                        observed_value=value,
                        message=f"Overall {metric_name} ({value:.3f}) below threshold ({threshold:.3f})"
                    )
                    alerts.append(alert)

        # Check stratified performance
        for var, groups in stratified_metrics.items():
            for group, metrics in groups.items():
                for metric_name, threshold in self.performance_thresholds.items():
                    if hasattr(metrics, metric_name):
                        value = getattr(metrics, metric_name)
                        if not np.isnan(value) and value < threshold:
                            alert = self._create_alert(
                                severity=AlertSeverity.CRITICAL,
                                alert_type="subgroup_performance_degradation",
                                affected_groups=[f"{var}_{group}"],
                                metric_name=metric_name,
                                threshold=threshold,
                                observed_value=value,
                                message=(
                                    f"{metric_name} for {var}={group} ({value:.3f}) "
                                    f"below threshold ({threshold:.3f})"
                                )
                            )
                            alerts.append(alert)

        # Check fairness thresholds
        if fairness_metrics:
            # Check AUC disparity
            if 'auc_disparity' in self.fairness_thresholds:
                threshold = self.fairness_thresholds['auc_disparity']
                if fairness_metrics.auc_disparity > threshold:
                    alert = self._create_alert(
                        severity=AlertSeverity.CRITICAL,
                        alert_type="fairness_violation",
                        affected_groups=[fairness_metrics.reference_group] +
                                      fairness_metrics.comparison_groups,
                        metric_name="auc_disparity",
                        threshold=threshold,
                        observed_value=fairness_metrics.auc_disparity,
                        message=(
                            f"AUC disparity ({fairness_metrics.auc_disparity:.3f}) "
                            f"exceeds threshold ({threshold:.3f})"
                        )
                    )
                    alerts.append(alert)

            # Check equal opportunity differences
            if 'equal_opportunity' in self.fairness_thresholds:
                threshold = self.fairness_thresholds['equal_opportunity']
                for group, diff in fairness_metrics.equal_opportunity_difference.items():
                    if abs(diff) > threshold:
                        alert = self._create_alert(
                            severity=AlertSeverity.CRITICAL,
                            alert_type="fairness_violation",
                            affected_groups=[f"{self.stratification_variables[0]}_{group}"],
                            metric_name="equal_opportunity_difference",
                            threshold=threshold,
                            observed_value=diff,
                            message=(
                                f"Sensitivity difference for {group} ({diff:.3f}) "
                                f"exceeds threshold ({threshold:.3f})"
                            )
                        )
                        alerts.append(alert)

        return alerts

    def _create_alert(
        self,
        severity: AlertSeverity,
        alert_type: str,
        affected_groups: List[str],
        metric_name: str,
        threshold: float,
        observed_value: float,
        message: str
    ) -> MonitoringAlert:
        """Create a monitoring alert with unique ID."""
        self.alert_counter += 1
        alert_id = f"{self.model_name}_{datetime.now().strftime('%Y%m%d')}_{self.alert_counter:04d}"

        return MonitoringAlert(
            alert_id=alert_id,
            severity=severity,
            alert_type=alert_type,
            affected_groups=affected_groups,
            metric_name=metric_name,
            threshold=threshold,
            observed_value=observed_value,
            message=message,
            timestamp=datetime.now()
        )

    def _persist_prediction(self, prediction_record: Dict) -> None:
        """Persist prediction record to storage."""
        try:
            date_str = prediction_record['timestamp'].strftime('%Y%m%d')
            filepath = self.storage_path / f"predictions_{date_str}.jsonl"

            # Serialize record
            record_copy = prediction_record.copy()
            record_copy['timestamp'] = record_copy['timestamp'].isoformat()
            if 'outcome_timestamp' in record_copy:
                record_copy['outcome_timestamp'] = record_copy['outcome_timestamp'].isoformat()

            with open(filepath, 'a') as f:
                f.write(json.dumps(record_copy) + '\n')
        except Exception as e:
            logger.error(f"Error persisting prediction: {e}")

    def _persist_outcome(
        self,
        prediction_id: str,
        true_outcome: int,
        outcome_timestamp: Optional[datetime]
    ) -> None:
        """Persist outcome record to storage."""
        try:
            outcome_record = {
                'prediction_id': prediction_id,
                'true_outcome': true_outcome,
                'outcome_timestamp': outcome_timestamp.isoformat() if outcome_timestamp else None
            }

            filepath = self.storage_path / "outcomes.jsonl"
            with open(filepath, 'a') as f:
                f.write(json.dumps(outcome_record) + '\n')
        except Exception as e:
            logger.error(f"Error persisting outcome: {e}")

    def generate_monitoring_report(
        self,
        output_path: Optional[Path] = None
    ) -> str:
        """
        Generate comprehensive monitoring report.

        Args:
            output_path: Optional path to save report

        Returns:
            Report text
        """
        report_lines = []
        report_lines.append("=" * 80)
        report_lines.append(f"MONITORING REPORT: {self.model_name}")
        report_lines.append("=" * 80)
        report_lines.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        report_lines.append("")

        # Summary statistics
        report_lines.append("SUMMARY")
        report_lines.append("-" * 80)
        report_lines.append(f"Total predictions recorded: {len(self.predictions_buffer)}")
        completed = sum(1 for p in self.predictions_buffer if p['outcome_recorded'])
        report_lines.append(f"Completed predictions: {completed}")
        report_lines.append(f"Active alerts: {len([a for a in self.active_alerts if a.resolution_status == 'unresolved'])}")
        report_lines.append("")

        # Recent performance
        if 'overall' in self.performance_history and self.performance_history['overall']:
            report_lines.append("RECENT OVERALL PERFORMANCE")
            report_lines.append("-" * 80)
            recent_metrics = self.performance_history['overall'][-1]
            report_lines.append(f"Evaluation time: {recent_metrics.timestamp.strftime('%Y-%m-%d %H:%M:%S')}")
            report_lines.append(f"Samples: {recent_metrics.n_samples}")
            report_lines.append(f"AUC: {recent_metrics.auc:.3f}")
            report_lines.append(f"Sensitivity: {recent_metrics.sensitivity:.3f}")
            report_lines.append(f"Specificity: {recent_metrics.specificity:.3f}")
            report_lines.append(f"Expected Calibration Error: {recent_metrics.ece:.3f}")
            report_lines.append("")

        # Active alerts
        if self.active_alerts:
            unresolved = [a for a in self.active_alerts if a.resolution_status == 'unresolved']
            if unresolved:
                report_lines.append("ACTIVE ALERTS")
                report_lines.append("-" * 80)
                for alert in sorted(unresolved, key=lambda x: x.severity.value, reverse=True):
                    report_lines.append(f"\nAlert ID: {alert.alert_id}")
                    report_lines.append(f"Severity: {alert.severity.value.upper()}")
                    report_lines.append(f"Type: {alert.alert_type}")
                    report_lines.append(f"Affected groups: {', '.join(alert.affected_groups)}")
                    report_lines.append(f"Message: {alert.message}")
                    report_lines.append(f"Timestamp: {alert.timestamp.strftime('%Y-%m-%d %H:%M:%S')}")
                report_lines.append("")

        report_text = '\n'.join(report_lines)

        if output_path:
            with open(output_path, 'w') as f:
                f.write(report_text)
            logger.info(f"Monitoring report saved to {output_path}")

        return report_text

def example_monitoring_workflow():
    """
    Demonstrate monitoring workflow for a clinical risk prediction model.
    """
    print("Stratified Performance Monitoring Example")
    print("=" * 80)

    # Initialize monitoring system
    monitor = StratifiedPerformanceMonitor(
        model_name="readmission_risk_predictor_v1",
        stratification_variables=['race', 'insurance_type', 'language'],
        performance_thresholds={
            'auc': 0.70,
            'sensitivity': 0.75,
            'ppv': 0.30
        },
        fairness_thresholds={
            'auc_disparity': 0.05,
            'equal_opportunity': 0.10
        },
        min_sample_size=50,
        monitoring_window_days=7
    )

    # Simulate recording predictions over time
    np.random.seed(42)
    n_predictions = 1000

    print(f"\nSimulating {n_predictions} predictions...")

    for i in range(n_predictions):
        # Simulate demographic distribution
        race = np.random.choice(['White', 'Black', 'Hispanic', 'Asian'], p=[0.6, 0.2, 0.15, 0.05])
        insurance = np.random.choice(['Commercial', 'Medicare', 'Medicaid'], p=[0.5, 0.3, 0.2])
        language = np.random.choice(['English', 'Spanish', 'Other'], p=[0.8, 0.15, 0.05])

        # Simulate model predictions with bias
        # Model performs worse for certain groups
        base_risk = 0.3
        if race == 'Black':
            # Model miscalibrated for Black patients
            true_risk = base_risk + 0.05
            pred_risk = base_risk - 0.05  # Systematic underestimation
        elif insurance == 'Medicaid':
            true_risk = base_risk + 0.08
            pred_risk = base_risk  # Underestimates risk
        else:
            true_risk = base_risk
            pred_risk = base_risk

        pred_risk = np.clip(pred_risk + np.random.normal(0, 0.1), 0, 1)
        true_outcome = 1 if np.random.random() < true_risk else 0

        # Record prediction
        timestamp = datetime.now() - timedelta(days=np.random.randint(0, 7))
        monitor.record_prediction(
            prediction_id=f"pred_{i:05d}",
            features={'age': 65, 'comorbidity_count': 3},
            predicted_probability=pred_risk,
            predicted_class=1 if pred_risk > 0.5 else 0,
            demographic_attributes={
                'race': race,
                'insurance_type': insurance,
                'language': language
            },
            timestamp=timestamp
        )

        # Record outcome
        monitor.record_outcome(f"pred_{i:05d}", true_outcome)

    # Evaluate monitoring window
    print("\nEvaluating monitoring window...")
    results = monitor.evaluate_monitoring_window()

    if 'error' not in results:
        print(f"\nEvaluated {results['n_predictions']} predictions")
        print(f"From {results['window_start'].date()} to {results['window_end'].date()}")

        overall_metrics = results['overall_metrics']
        print(f"\nOverall Performance:")
        print(f"  AUC: {overall_metrics.auc:.3f}")
        print(f"  Sensitivity: {overall_metrics.sensitivity:.3f}")
        print(f"  Specificity: {overall_metrics.specificity:.3f}")
        print(f"  Expected Calibration Error: {overall_metrics.ece:.3f}")

        # Show stratified metrics for race
        if 'race' in results['stratified_metrics']:
            print(f"\nPerformance by Race:")
            for race, metrics in results['stratified_metrics']['race'].items():
                print(f"  {race}:")
                print(f"    N={metrics.n_samples}, AUC={metrics.auc:.3f}, "
                      f"Sensitivity={metrics.sensitivity:.3f}")

        # Show fairness metrics
        if results['fairness_metrics']:
            fm = results['fairness_metrics']
            print(f"\nFairness Evaluation:")
            print(f"  Reference group: {fm.reference_group}")
            print(f"  AUC disparity: {fm.auc_disparity:.3f}")
            print(f"  Calibration disparity: {fm.calibration_disparity:.3f}")

        # Show alerts
        if results['alerts']:
            print(f"\n{len(results['alerts'])} ALERTS GENERATED:")
            for alert in results['alerts']:
                print(f"  [{alert.severity.value.upper()}] {alert.message}")
                print(f"    Affected: {', '.join(alert.affected_groups)}")

    # Generate report
    print("\nGenerating monitoring report...")
    report = monitor.generate_monitoring_report()
    print("\n" + report)

if __name__ == "__main__":
    example_monitoring_workflow()

"""
Temporal monitoring methods for detecting performance drift over time.

This module implements sequential analysis and changepoint detection methods
for continuous monitoring of model performance with appropriate control of
false positive rates and early detection of concerning trends.
"""

import logging
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
import numpy as np
import pandas as pd
from scipy import stats
from datetime import datetime, timedelta
import matplotlib.pyplot as plt

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class SequentialTestResult:
    """
    Result from sequential hypothesis test.

    Attributes:
        decision: 'continue', 'reject_null', or 'accept_null'
        likelihood_ratio: Current likelihood ratio
        upper_boundary: Upper stopping boundary
        lower_boundary: Lower stopping boundary
        n_observations: Cumulative observations used
        timestamp: When test was conducted
    """
    decision: str
    likelihood_ratio: float
    upper_boundary: float
    lower_boundary: float
    n_observations: int
    timestamp: datetime

@dataclass
class ChangepointDetectionResult:
    """
    Result from changepoint detection analysis.

    Attributes:
        changepoints: List of detected changepoint indices
        changepoint_times: Times of detected changepoints
        segment_means: Mean metric values between changepoints
        posterior_probabilities: Probabilities of changepoints at each time
    """
    changepoints: List[int]
    changepoint_times: List[datetime]
    segment_means: List[float]
    posterior_probabilities: np.ndarray

class SequentialPerformanceTest:
    """
    Sequential probability ratio test for monitoring model performance.

    Implements sequential testing framework that enables continuous monitoring
    while controlling type I error rates through appropriate stopping boundaries.
    """

    def __init__(
        self,
        null_value: float,
        alternative_value: float,
        alpha: float = 0.05,
        beta: float = 0.20,
        test_statistic: str = 'auc'
    ):
        """
        Initialize sequential test.

        Args:
            null_value: Performance under null hypothesis (acceptable level)
            alternative_value: Performance under alternative (degraded level)
            alpha: Type I error rate (false positive rate)
            beta: Type II error rate (false negative rate)
            test_statistic: Which performance metric to monitor
        """
        self.null_value = null_value
        self.alternative_value = alternative_value
        self.alpha = alpha
        self.beta = beta
        self.test_statistic = test_statistic

        # Compute stopping boundaries using Wald's SPRT
        # A = (1 - beta) / alpha (upper boundary for rejecting null)
        # B = beta / (1 - alpha) (lower boundary for accepting null)
        self.upper_boundary = (1 - beta) / alpha
        self.lower_boundary = beta / (1 - alpha)

        # Track test history
        self.likelihood_ratios: List[float] = []
        self.observations: List[float] = []
        self.timestamps: List[datetime] = []

        logger.info(
            f"Initialized SequentialPerformanceTest for {test_statistic}: "
            f"H0={null_value}, H1={alternative_value}, α={alpha}, β={beta}"
        )

    def update(
        self,
        observed_performance: float,
        n_samples: int,
        timestamp: Optional[datetime] = None
    ) -> SequentialTestResult:
        """
        Update sequential test with new performance observation.

        Args:
            observed_performance: Observed value of test statistic
            n_samples: Sample size for this observation
            timestamp: Time of observation

        Returns:
            SequentialTestResult with current test status
        """
        if timestamp is None:
            timestamp = datetime.now()

        self.observations.append(observed_performance)
        self.timestamps.append(timestamp)

        # Compute likelihood ratio
        # For simplicity, assume normal approximation for AUC
        # L(data | H1) / L(data | H0)

        # Standard error for AUC (Hanley-McNeil approximation)
        se = self._compute_standard_error(n_samples)

        # Log likelihood under each hypothesis
        log_lik_alt = stats.norm.logpdf(
            observed_performance, loc=self.alternative_value, scale=se
        )
        log_lik_null = stats.norm.logpdf(
            observed_performance, loc=self.null_value, scale=se
        )

        # Cumulative log likelihood ratio
        if len(self.likelihood_ratios) == 0:
            cum_log_lr = log_lik_alt - log_lik_null
        else:
            cum_log_lr = self.likelihood_ratios[-1] + (log_lik_alt - log_lik_null)

        likelihood_ratio = np.exp(cum_log_lr)
        self.likelihood_ratios.append(cum_log_lr)

        # Make decision
        if likelihood_ratio >= self.upper_boundary:
            decision = "reject_null"
            logger.warning(
                f"Sequential test rejects null hypothesis: "
                f"Performance has degraded to {observed_performance:.3f}"
            )
        elif likelihood_ratio <= self.lower_boundary:
            decision = "accept_null"
            logger.info(
                f"Sequential test accepts null hypothesis: "
                f"Performance remains at acceptable level"
            )
        else:
            decision = "continue"

        return SequentialTestResult(
            decision=decision,
            likelihood_ratio=likelihood_ratio,
            upper_boundary=self.upper_boundary,
            lower_boundary=self.lower_boundary,
            n_observations=len(self.observations),
            timestamp=timestamp
        )

    def _compute_standard_error(self, n_samples: int) -> float:
        """
        Compute standard error for test statistic.

        Args:
            n_samples: Sample size

        Returns:
            Standard error estimate
        """
        # Hanley-McNeil approximation for AUC standard error
        # SE = sqrt((AUC * (1 - AUC)) / n)
        # Use average of null and alternative values
        avg_auc = (self.null_value + self.alternative_value) / 2
        se = np.sqrt((avg_auc * (1 - avg_auc)) / n_samples)
        return se

    def plot_sequential_test(self, output_path: Optional[str] = None):
        """
        Plot likelihood ratio trajectory with stopping boundaries.

        Args:
            output_path: Optional path to save plot
        """
        if len(self.likelihood_ratios) == 0:
            logger.warning("No observations to plot")
            return

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))

        # Plot likelihood ratio
        ax1.plot(range(len(self.likelihood_ratios)), self.likelihood_ratios,
                 'b-', linewidth=2, label='Cumulative Log LR')
        ax1.axhline(y=np.log(self.upper_boundary), color='r', linestyle='--',
                   label=f'Upper boundary (reject H0)')
        ax1.axhline(y=np.log(self.lower_boundary), color='g', linestyle='--',
                   label=f'Lower boundary (accept H0)')
        ax1.set_xlabel('Observation Number')
        ax1.set_ylabel('Cumulative Log Likelihood Ratio')
        ax1.set_title('Sequential Probability Ratio Test')
        ax1.legend()
        ax1.grid(True, alpha=0.3)

        # Plot observed performance
        ax2.plot(range(len(self.observations)), self.observations,
                'b-o', linewidth=2, markersize=4, label='Observed Performance')
        ax2.axhline(y=self.null_value, color='g', linestyle='--',
                   label=f'Null hypothesis ({self.null_value:.2f})')
        ax2.axhline(y=self.alternative_value, color='r', linestyle='--',
                   label=f'Alternative hypothesis ({self.alternative_value:.2f})')
        ax2.set_xlabel('Observation Number')
        ax2.set_ylabel(f'{self.test_statistic.upper()}')
        ax2.set_title('Observed Performance Over Time')
        ax2.legend()
        ax2.grid(True, alpha=0.3)

        plt.tight_layout()

        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            logger.info(f"Saved sequential test plot to {output_path}")
        else:
            plt.show()

class BayesianChangepointDetector:
    """
    Bayesian online changepoint detection for performance monitoring.

    Detects points in time where the distribution of performance metrics
    changes, indicating shifts in model behavior.
    """

    def __init__(
        self,
        hazard_rate: float = 0.01,
        prior_mean: float = 0.75,
        prior_variance: float = 0.01
    ):
        """
        Initialize changepoint detector.

        Args:
            hazard_rate: Prior probability of changepoint at each time
            prior_mean: Prior mean for performance metric
            prior_variance: Prior variance for performance metric
        """
        self.hazard_rate = hazard_rate
        self.prior_mean = prior_mean
        self.prior_variance = prior_variance

        # Initialize run length distribution
        # R[t] = P(run length at time t)
        self.run_length_dist = np.array([1.0])

        # Sufficient statistics for each run length
        self.mean_params = [prior_mean]
        self.variance_params = [prior_variance]
        self.observation_counts = [0]

        # Store results
        self.changepoint_probabilities: List[float] = []
        self.max_run_lengths: List[int] = []

        logger.info(
            f"Initialized BayesianChangepointDetector with hazard_rate={hazard_rate}"
        )

    def update(
        self,
        observation: float,
        observation_variance: float = 0.01
    ) -> float:
        """
        Update changepoint detection with new observation.

        Args:
            observation: New performance metric value
            observation_variance: Measurement uncertainty

        Returns:
            Probability of changepoint at this time
        """
        # Compute predictive probabilities for each run length
        n_run_lengths = len(self.run_length_dist)
        pred_probs = np.zeros(n_run_lengths)

        for r in range(n_run_lengths):
            # Predictive distribution is Student's t
            mean = self.mean_params[r]
            var = self.variance_params[r] + observation_variance
            pred_probs[r] = stats.norm.pdf(observation, loc=mean, scale=np.sqrt(var))

        # Compute growth probabilities (no changepoint)
        growth_probs = self.run_length_dist * pred_probs * (1 - self.hazard_rate)

        # Compute changepoint probability
        cp_prob = np.sum(self.run_length_dist * pred_probs * self.hazard_rate)

        # Update run length distribution
        # Shift and add changepoint
        new_run_length_dist = np.zeros(n_run_lengths + 1)
        new_run_length_dist[1:] = growth_probs
        new_run_length_dist[0] = cp_prob

        # Normalize
        new_run_length_dist = new_run_length_dist / np.sum(new_run_length_dist)
        self.run_length_dist = new_run_length_dist

        # Update sufficient statistics
        new_mean_params = [self.prior_mean]
        new_variance_params = [self.prior_variance]
        new_observation_counts = [0]

        for r in range(n_run_lengths):
            # Update statistics for continuing run lengths
            n = self.observation_counts[r] + 1
            old_mean = self.mean_params[r]
            new_mean = (old_mean * self.observation_counts[r] + observation) / n
            new_mean_params.append(new_mean)

            # Update variance (online algorithm)
            if n == 1:
                new_var = self.prior_variance
            else:
                old_var = self.variance_params[r]
                new_var = old_var + ((observation - old_mean) *
                                    (observation - new_mean) - old_var) / n
            new_variance_params.append(new_var)
            new_observation_counts.append(n)

        self.mean_params = new_mean_params
        self.variance_params = new_variance_params
        self.observation_counts = new_observation_counts

        # Store results
        self.changepoint_probabilities.append(cp_prob)
        self.max_run_lengths.append(np.argmax(self.run_length_dist))

        return cp_prob

    def detect_changepoints(
        self,
        observations: np.ndarray,
        timestamps: List[datetime],
        threshold: float = 0.5
    ) -> ChangepointDetectionResult:
        """
        Detect changepoints in sequence of observations.

        Args:
            observations: Array of performance metric values
            timestamps: Timestamps for observations
            threshold: Probability threshold for declaring changepoint

        Returns:
            ChangepointDetectionResult with detected changepoints
        """
        # Reset detector
        self.__init__(self.hazard_rate, self.prior_mean, self.prior_variance)

        # Process observations
        for obs in observations:
            self.update(obs)

        # Identify changepoints
        cp_probs = np.array(self.changepoint_probabilities)
        changepoint_indices = np.where(cp_probs > threshold)[0].tolist()
        changepoint_times = [timestamps[i] for i in changepoint_indices]

        # Compute segment means
        segment_means = []
        segment_starts = [0] + [i + 1 for i in changepoint_indices]
        segment_ends = changepoint_indices + [len(observations)]

        for start, end in zip(segment_starts, segment_ends):
            if start < end:
                segment_means.append(np.mean(observations[start:end]))

        logger.info(f"Detected {len(changepoint_indices)} changepoints")

        return ChangepointDetectionResult(
            changepoints=changepoint_indices,
            changepoint_times=changepoint_times,
            segment_means=segment_means,
            posterior_probabilities=cp_probs
        )

    def plot_changepoint_detection(
        self,
        observations: np.ndarray,
        timestamps: List[datetime],
        result: ChangepointDetectionResult,
        output_path: Optional[str] = None
    ):
        """
        Plot changepoint detection results.

        Args:
            observations: Observed performance metrics
            timestamps: Observation timestamps
            result: ChangepointDetectionResult from detection
            output_path: Optional path to save plot
        """
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10))

        # Plot observations with detected changepoints
        time_indices = range(len(observations))
        ax1.plot(time_indices, observations, 'b-o', linewidth=2,
                markersize=4, label='Observed Performance')

        # Mark changepoints
        for cp_idx in result.changepoints:
            ax1.axvline(x=cp_idx, color='r', linestyle='--', alpha=0.7)

        # Plot segment means
        segment_starts = [0] + [i + 1 for i in result.changepoints]
        segment_ends = result.changepoints + [len(observations)]

        for start, end, mean_val in zip(segment_starts, segment_ends, result.segment_means):
            ax1.hlines(y=mean_val, xmin=start, xmax=end-1,
                      colors='g', linewidth=3, alpha=0.5)

        ax1.set_xlabel('Observation Number')
        ax1.set_ylabel('Performance Metric')
        ax1.set_title('Performance with Detected Changepoints')
        ax1.legend()
        ax1.grid(True, alpha=0.3)

        # Plot changepoint probabilities
        ax2.plot(time_indices, result.posterior_probabilities,
                'r-', linewidth=2, label='Changepoint Probability')
        ax2.axhline(y=0.5, color='k', linestyle='--', alpha=0.5,
                   label='Detection Threshold')
        ax2.set_xlabel('Observation Number')
        ax2.set_ylabel('Changepoint Probability')
        ax2.set_title('Posterior Changepoint Probabilities')
        ax2.legend()
        ax2.grid(True, alpha=0.3)

        plt.tight_layout()

        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            logger.info(f"Saved changepoint detection plot to {output_path}")
        else:
            plt.show()

def example_temporal_monitoring():
    """
    Demonstrate sequential testing and changepoint detection.
    """
    print("Temporal Monitoring Example")
    print("=" * 80)

    # Simulate performance over time with degradation
    np.random.seed(42)
    n_weeks = 52

    # Generate performance that degrades after week 30
    performance = []
    timestamps = []
    base_date = datetime(2024, 1, 1)

    for week in range(n_weeks):
        if week < 30:
            # Stable performance
            perf = np.random.normal(0.78, 0.02)
        else:
            # Gradual degradation
            degradation = (week - 30) * 0.005
            perf = np.random.normal(0.78 - degradation, 0.02)

        performance.append(np.clip(perf, 0, 1))
        timestamps.append(base_date + timedelta(weeks=week))

    performance = np.array(performance)

    print(f"Simulated {n_weeks} weeks of performance data")
    print(f"Performance degrades starting at week 30")
    print(f"Initial AUC: {performance[:30].mean():.3f}")
    print(f"Final AUC: {performance[40:].mean():.3f}")

    # Sequential testing
    print("\n" + "=" * 80)
    print("Sequential Probability Ratio Test")
    print("=" * 80)

    sprt = SequentialPerformanceTest(
        null_value=0.75,  # Acceptable performance
        alternative_value=0.70,  # Degraded performance
        alpha=0.05,
        beta=0.20,
        test_statistic='auc'
    )

    for week, (perf, timestamp) in enumerate(zip(performance, timestamps)):
        result = sprt.update(
            observed_performance=perf,
            n_samples=1000,  # Simulate 1000 predictions per week
            timestamp=timestamp
        )

        if week % 10 == 0:
            print(f"\nWeek {week}:")
            print(f"  Observed AUC: {perf:.3f}")
            print(f"  Likelihood Ratio: {result.likelihood_ratio:.3f}")
            print(f"  Decision: {result.decision}")

        if result.decision != "continue":
            print(f"\nSequential test terminated at week {week}")
            print(f"  Final decision: {result.decision}")
            break

    # Plot sequential test
    sprt.plot_sequential_test()

    # Changepoint detection
    print("\n" + "=" * 80)
    print("Bayesian Changepoint Detection")
    print("=" * 80)

    detector = BayesianChangepointDetector(
        hazard_rate=0.02,  # Prior probability of changepoint each week
        prior_mean=0.78,
        prior_variance=0.01
    )

    result = detector.detect_changepoints(
        observations=performance,
        timestamps=timestamps,
        threshold=0.5
    )

    print(f"\nDetected {len(result.changepoints)} changepoints:")
    for i, (cp_idx, cp_time) in enumerate(zip(result.changepoints, result.changepoint_times)):
        print(f"  Changepoint {i+1}: Week {cp_idx} ({cp_time.date()})")

    print(f"\nSegment means:")
    for i, mean_val in enumerate(result.segment_means):
        print(f"  Segment {i+1}: AUC = {mean_val:.3f}")

    # Plot changepoint detection
    detector.plot_changepoint_detection(performance, timestamps, result)

    print("\nKey observations:")
    print("1. Sequential test detected degradation before reaching threshold")
    print("2. Changepoint detection identified approximate time of drift")
    print("3. Both methods provide early warning of performance issues")
    print("4. For equity monitoring, apply separately to each demographic subgroup")

if __name__ == "__main__":
    example_temporal_monitoring()

"""
Data drift detection methods for monitoring distribution shifts.

This module implements statistical tests and methods for detecting covariate
shift, label shift, and concept drift in production clinical AI systems,
with equity-focused analysis across demographic subgroups.
"""

import logging
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass
from datetime import datetime
import numpy as np
import pandas as pd
from scipy import stats
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
from sklearn.metrics import roc_auc_score
import warnings

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class DriftDetectionResult:
    """
    Result from drift detection analysis.

    Attributes:
        drift_detected: Whether significant drift was found
        drift_type: Type of drift detected (covariate, label, concept)
        drift_magnitude: Quantitative measure of drift severity
        affected_features: Features showing significant drift
        test_statistics: Dictionary of test statistics and p-values
        timestamp: When drift detection was performed
        recommendation: Suggested action based on drift severity
    """
    drift_detected: bool
    drift_type: str
    drift_magnitude: float
    affected_features: List[str]
    test_statistics: Dict[str, Any]
    timestamp: datetime
    recommendation: str

    def to_dict(self) -> Dict[str, Any]:
        """Convert result to dictionary."""
        return {
            'drift_detected': self.drift_detected,
            'drift_type': self.drift_type,
            'drift_magnitude': float(self.drift_magnitude),
            'affected_features': self.affected_features,
            'test_statistics': self.test_statistics,
            'timestamp': self.timestamp.isoformat(),
            'recommendation': self.recommendation
        }

class CovariateShiftDetector:
    """
    Detects covariate shift using multivariate and univariate methods.

    Implements classifier-based drift detection and univariate statistical
    tests to identify changes in feature distributions between reference
    and current data.
    """

    def __init__(
        self,
        feature_names: List[str],
        categorical_features: Optional[List[str]] = None,
        significance_level: float = 0.05,
        min_samples: int = 100
    ):
        """
        Initialize covariate shift detector.

        Args:
            feature_names: Names of features to monitor
            categorical_features: Names of categorical features
            significance_level: Significance level for statistical tests
            min_samples: Minimum samples needed for reliable detection
        """
        self.feature_names = feature_names
        self.categorical_features = categorical_features or []
        self.significance_level = significance_level
        self.min_samples = min_samples

        # Classifier for multivariate drift detection
        self.drift_classifier = RandomForestClassifier(
            n_estimators=100,
            max_depth=5,
            random_state=42
        )

        logger.info(
            f"Initialized CovariateShiftDetector for {len(feature_names)} features"
        )

    def detect_drift(
        self,
        reference_data: pd.DataFrame,
        current_data: pd.DataFrame
    ) -> DriftDetectionResult:
        """
        Detect covariate shift between reference and current data.

        Args:
            reference_data: Historical reference data
            current_data: Current production data

        Returns:
            DriftDetectionResult with detected drift information
        """
        timestamp = datetime.now()

        # Validate inputs
        if len(reference_data) < self.min_samples or len(current_data) < self.min_samples:
            logger.warning(
                f"Insufficient samples for drift detection: "
                f"reference={len(reference_data)}, current={len(current_data)}"
            )
            return DriftDetectionResult(
                drift_detected=False,
                drift_type="unknown",
                drift_magnitude=0.0,
                affected_features=[],
                test_statistics={'error': 'insufficient_samples'},
                timestamp=timestamp,
                recommendation="Collect more data before detecting drift"
            )

        # Multivariate drift detection using classifier
        multivariate_result = self._classifier_based_detection(
            reference_data, current_data
        )

        # Univariate drift detection for each feature
        univariate_results = self._univariate_tests(
            reference_data, current_data
        )

        # Identify features with significant drift
        affected_features = [
            feat for feat, result in univariate_results.items()
            if result['significant']
        ]

        # Determine overall drift
        drift_detected = (
            multivariate_result['auc'] > 0.6 or  # Classifier can distinguish
            len(affected_features) > len(self.feature_names) * 0.1  # >10% features drifted
        )

        drift_magnitude = max(
            multivariate_result['auc'] - 0.5,  # Distance from random
            len(affected_features) / len(self.feature_names)  # Proportion affected
        )

        # Generate recommendation
        if drift_magnitude < 0.1:
            recommendation = "Monitor but no action needed"
        elif drift_magnitude < 0.2:
            recommendation = "Investigate affected features, consider recalibration"
        elif drift_magnitude < 0.3:
            recommendation = "Significant drift detected, plan model update"
        else:
            recommendation = "Severe drift detected, consider retraining or pausing deployment"

        test_statistics = {
            'multivariate': multivariate_result,
            'univariate': univariate_results,
            'n_affected_features': len(affected_features),
            'proportion_affected': len(affected_features) / len(self.feature_names)
        }

        if drift_detected:
            logger.warning(
                f"Covariate shift detected: {len(affected_features)} features affected, "
                f"magnitude={drift_magnitude:.3f}"
            )

        return DriftDetectionResult(
            drift_detected=drift_detected,
            drift_type="covariate_shift",
            drift_magnitude=drift_magnitude,
            affected_features=affected_features,
            test_statistics=test_statistics,
            timestamp=timestamp,
            recommendation=recommendation
        )

    def _classifier_based_detection(
        self,
        reference_data: pd.DataFrame,
        current_data: pd.DataFrame
    ) -> Dict[str, float]:
        """
        Use classifier to detect multivariate distribution shift.

        Args:
            reference_data: Reference data
            current_data: Current data

        Returns:
            Dictionary with detection results
        """
        # Create labels (0 for reference, 1 for current)
        n_ref = len(reference_data)
        n_cur = len(current_data)

        X = pd.concat([reference_data[self.feature_names],
                      current_data[self.feature_names]], axis=0)
        y = np.concatenate([np.zeros(n_ref), np.ones(n_cur)])

        # Handle missing values
        X = X.fillna(X.median())

        # Train classifier to distinguish reference from current
        self.drift_classifier.fit(X, y)
        y_pred_proba = self.drift_classifier.predict_proba(X)[:, 1]

        # Compute AUC as measure of drift
        # AUC near 0.5 means no drift, higher values indicate drift
        auc = roc_auc_score(y, y_pred_proba)

        # Get feature importances to identify main contributors to drift
        feature_importances = dict(zip(
            self.feature_names,
            self.drift_classifier.feature_importances_
        ))

        return {
            'auc': auc,
            'feature_importances': feature_importances,
            'drift_score': abs(auc - 0.5) * 2  # Scale to [0, 1]
        }

    def _univariate_tests(
        self,
        reference_data: pd.DataFrame,
        current_data: pd.DataFrame
    ) -> Dict[str, Dict]:
        """
        Perform univariate statistical tests for each feature.

        Args:
            reference_data: Reference data
            current_data: Current data

        Returns:
            Dictionary mapping feature names to test results
        """
        results = {}

        for feature in self.feature_names:
            ref_vals = reference_data[feature].dropna()
            cur_vals = current_data[feature].dropna()

            if len(ref_vals) < 10 or len(cur_vals) < 10:
                results[feature] = {
                    'test': 'insufficient_data',
                    'statistic': np.nan,
                    'pvalue': np.nan,
                    'significant': False
                }
                continue

            if feature in self.categorical_features:
                # Chi-square test for categorical variables
                ref_counts = ref_vals.value_counts()
                cur_counts = cur_vals.value_counts()

                # Align categories
                all_categories = set(ref_counts.index) | set(cur_counts.index)
                ref_aligned = [ref_counts.get(cat, 0) for cat in all_categories]
                cur_aligned = [cur_counts.get(cat, 0) for cat in all_categories]

                try:
                    statistic, pvalue = stats.chisquare(cur_aligned, ref_aligned)
                    test_name = 'chi_square'
                except Exception as e:
                    logger.warning(f"Chi-square test failed for {feature}: {e}")
                    statistic, pvalue = np.nan, np.nan
                    test_name = 'failed'

            else:
                # Kolmogorov-Smirnov test for continuous variables
                try:
                    statistic, pvalue = stats.ks_2samp(ref_vals, cur_vals)
                    test_name = 'kolmogorov_smirnov'
                except Exception as e:
                    logger.warning(f"KS test failed for {feature}: {e}")
                    statistic, pvalue = np.nan, np.nan
                    test_name = 'failed'

            # Apply Bonferroni correction for multiple testing
            corrected_alpha = self.significance_level / len(self.feature_names)
            significant = pvalue < corrected_alpha if not np.isnan(pvalue) else False

            results[feature] = {
                'test': test_name,
                'statistic': float(statistic) if not np.isnan(statistic) else None,
                'pvalue': float(pvalue) if not np.isnan(pvalue) else None,
                'significant': significant,
                'effect_size': self._compute_effect_size(ref_vals, cur_vals, feature)
            }

        return results

    def _compute_effect_size(
        self,
        ref_vals: pd.Series,
        cur_vals: pd.Series,
        feature: str
    ) -> float:
        """
        Compute effect size for distributional difference.

        Args:
            ref_vals: Reference feature values
            cur_vals: Current feature values
            feature: Feature name

        Returns:
            Effect size measure
        """
        if feature in self.categorical_features:
            # Cramér's V for categorical variables
            ref_props = ref_vals.value_counts(normalize=True)
            cur_props = cur_vals.value_counts(normalize=True)
            all_cats = set(ref_props.index) | set(cur_props.index)
            diff = sum(abs(ref_props.get(cat, 0) - cur_props.get(cat, 0))
                      for cat in all_cats)
            return diff / 2  # Total variation distance
        else:
            # Cohen's d for continuous variables
            ref_mean, ref_std = ref_vals.mean(), ref_vals.std()
            cur_mean, cur_std = cur_vals.mean(), cur_vals.std()
            pooled_std = np.sqrt((ref_std**2 + cur_std**2) / 2)
            if pooled_std > 0:
                return abs(cur_mean - ref_mean) / pooled_std
            else:
                return 0.0

class LabelShiftDetector:
    """
    Detects label shift by monitoring outcome prevalence over time.

    Tracks changes in outcome distribution while accounting for sampling
    variability and demographic stratification.
    """

    def __init__(
        self,
        baseline_prevalence: float,
        significance_level: float = 0.05,
        min_samples: int = 100
    ):
        """
        Initialize label shift detector.

        Args:
            baseline_prevalence: Expected outcome prevalence
            significance_level: Significance level for tests
            min_samples: Minimum samples for reliable detection
        """
        self.baseline_prevalence = baseline_prevalence
        self.significance_level = significance_level
        self.min_samples = min_samples

        logger.info(
            f"Initialized LabelShiftDetector with baseline={baseline_prevalence:.3f}"
        )

    def detect_shift(
        self,
        current_outcomes: np.ndarray,
        demographic_groups: Optional[np.ndarray] = None
    ) -> DriftDetectionResult:
        """
        Detect label shift in current outcomes.

        Args:
            current_outcomes: Binary outcomes from current data
            demographic_groups: Optional group labels for stratified analysis

        Returns:
            DriftDetectionResult with shift information
        """
        timestamp = datetime.now()

        if len(current_outcomes) < self.min_samples:
            return DriftDetectionResult(
                drift_detected=False,
                drift_type="unknown",
                drift_magnitude=0.0,
                affected_features=[],
                test_statistics={'error': 'insufficient_samples'},
                timestamp=timestamp,
                recommendation="Collect more data"
            )

        # Test overall prevalence change
        current_prevalence = np.mean(current_outcomes)

        # Binomial test for significant change
        n_positive = int(np.sum(current_outcomes))
        n_total = len(current_outcomes)

        # Two-sided test
        pvalue = 2 * min(
            stats.binom_test(n_positive, n_total, self.baseline_prevalence, alternative='less'),
            stats.binom_test(n_positive, n_total, self.baseline_prevalence, alternative='greater')
        )

        shift_detected = pvalue < self.significance_level
        shift_magnitude = abs(current_prevalence - self.baseline_prevalence)

        test_statistics = {
            'baseline_prevalence': self.baseline_prevalence,
            'current_prevalence': float(current_prevalence),
            'prevalence_change': float(shift_magnitude),
            'pvalue': float(pvalue),
            'n_samples': n_total,
            'n_positive': n_positive
        }

        # Stratified analysis if demographic groups provided
        if demographic_groups is not None:
            stratified_results = {}
            affected_groups = []

            for group in np.unique(demographic_groups):
                group_mask = demographic_groups == group
                group_outcomes = current_outcomes[group_mask]

                if len(group_outcomes) >= 30:  # Minimum for group analysis
                    group_prevalence = np.mean(group_outcomes)
                    group_n_pos = int(np.sum(group_outcomes))
                    group_pval = 2 * min(
                        stats.binom_test(group_n_pos, len(group_outcomes),
                                       self.baseline_prevalence, alternative='less'),
                        stats.binom_test(group_n_pos, len(group_outcomes),
                                       self.baseline_prevalence, alternative='greater')
                    )

                    stratified_results[str(group)] = {
                        'prevalence': float(group_prevalence),
                        'change': float(abs(group_prevalence - self.baseline_prevalence)),
                        'pvalue': float(group_pval),
                        'n': len(group_outcomes)
                    }

                    if group_pval < self.significance_level:
                        affected_groups.append(str(group))

            test_statistics['stratified'] = stratified_results
            test_statistics['affected_groups'] = affected_groups

        # Generate recommendation
        if shift_magnitude < 0.05:
            recommendation = "Minor shift, continue monitoring"
        elif shift_magnitude < 0.10:
            recommendation = "Moderate shift detected, consider recalibration"
        else:
            recommendation = "Significant label shift, recalibration required"

        if shift_detected:
            logger.warning(
                f"Label shift detected: prevalence changed from {self.baseline_prevalence:.3f} "
                f"to {current_prevalence:.3f} (p={pvalue:.4f})"
            )

        return DriftDetectionResult(
            drift_detected=shift_detected,
            drift_type="label_shift",
            drift_magnitude=shift_magnitude,
            affected_features=['outcome_prevalence'],
            test_statistics=test_statistics,
            timestamp=timestamp,
            recommendation=recommendation
        )

def example_drift_detection():
    """
    Demonstrate drift detection methods.
    """
    print("Data Drift Detection Example")
    print("=" * 80)

    np.random.seed(42)

    # Generate reference dataset
    n_ref = 1000
    n_features = 10

    reference_features = np.random.randn(n_ref, n_features)
    reference_df = pd.DataFrame(
        reference_features,
        columns=[f'feature_{i}' for i in range(n_features)]
    )
    reference_df['categorical_feat'] = np.random.choice(['A', 'B', 'C'], n_ref)
    reference_outcomes = (reference_features[:, 0] +
                         reference_features[:, 1] +
                         np.random.randn(n_ref) * 0.5 > 0).astype(int)

    print(f"Generated reference dataset: {n_ref} samples, {n_features} features")
    print(f"Reference outcome prevalence: {reference_outcomes.mean():.3f}")

    # Generate current dataset with covariate shift
    n_cur = 800
    # Add systematic shift to some features
    current_features = np.random.randn(n_cur, n_features)
    current_features[:, 0] += 0.5  # Shift feature 0
    current_features[:, 2] *= 1.5  # Scale feature 2

    current_df = pd.DataFrame(
        current_features,
        columns=[f'feature_{i}' for i in range(n_features)]
    )
    # Change categorical distribution
    current_df['categorical_feat'] = np.random.choice(
        ['A', 'B', 'C'], n_cur, p=[0.5, 0.3, 0.2]
    )
    current_outcomes = (current_features[:, 0] +
                       current_features[:, 1] +
                       np.random.randn(n_cur) * 0.5 > 0.3).astype(int)  # Label shift

    print(f"\nGenerated current dataset: {n_cur} samples")
    print(f"Current outcome prevalence: {current_outcomes.mean():.3f}")

    # Detect covariate shift
    print("\n" + "=" * 80)
    print("Covariate Shift Detection")
    print("=" * 80)

    covariate_detector = CovariateShiftDetector(
        feature_names=[f'feature_{i}' for i in range(n_features)] + ['categorical_feat'],
        categorical_features=['categorical_feat']
    )

    covariate_result = covariate_detector.detect_drift(reference_df, current_df)

    print(f"\nDrift detected: {covariate_result.drift_detected}")
    print(f"Drift magnitude: {covariate_result.drift_magnitude:.3f}")
    print(f"Affected features: {', '.join(covariate_result.affected_features)}")
    print(f"Recommendation: {covariate_result.recommendation}")

    print("\nMultivariate drift detection:")
    multivariate = covariate_result.test_statistics['multivariate']
    print(f"  Classifier AUC: {multivariate['auc']:.3f}")
    print(f"  Drift score: {multivariate['drift_score']:.3f}")

    print("\nTop features contributing to drift:")
    importances = multivariate['feature_importances']
    sorted_features = sorted(importances.items(), key=lambda x: x[1], reverse=True)[:5]
    for feat, imp in sorted_features:
        print(f"  {feat}: {imp:.4f}")

    # Detect label shift
    print("\n" + "=" * 80)
    print("Label Shift Detection")
    print("=" * 80)

    label_detector = LabelShiftDetector(
        baseline_prevalence=reference_outcomes.mean()
    )

    # Add demographic groups for stratified analysis
    current_demographics = np.random.choice(['Group_A', 'Group_B'], n_cur)

    label_result = label_detector.detect_shift(
        current_outcomes,
        demographic_groups=current_demographics
    )

    print(f"\nShift detected: {label_result.drift_detected}")
    print(f"Prevalence change: {label_result.drift_magnitude:.3f}")
    print(f"P-value: {label_result.test_statistics['pvalue']:.4f}")
    print(f"Recommendation: {label_result.recommendation}")

    if 'stratified' in label_result.test_statistics:
        print("\nStratified analysis:")
        for group, stats in label_result.test_statistics['stratified'].items():
            print(f"  {group}: prevalence={stats['prevalence']:.3f}, "
                  f"change={stats['change']:.3f}, p={stats['pvalue']:.4f}")

    print("\nKey observations:")
    print("1. Covariate shift detected in multiple features")
    print("2. Label shift indicates changing outcome prevalence")
    print("3. Stratified analysis reveals group-specific patterns")
    print("4. Both types of drift require different mitigation strategies")
    print("5. For equity, monitor drift separately for each demographic subgroup")

if __name__ == "__main__":
    example_drift_detection()