foxhunt/services/trading_service/src/ab_testing_pipeline.rs

//! A/B Testing Pipeline for Model Deployment Decisions
//!
//! This module implements an automated A/B testing pipeline that:
//! 1. Creates A/B tests on model deployment
//! 2. Splits traffic 50/50 between control and treatment
//! 3. Collects metrics (Sharpe ratio, win rate, drawdown, PnL)
//! 4. Runs statistical tests (Welch's t-test, p < 0.05)
//! 5. Makes deployment decisions (rollout, revert, or continue)
//! 6. Integrates with ml/src/ensemble/ab_testing.rs
//!
//! ## Architecture
//!
//! ```text
//! New Model Deployed
//!         │
//!         ▼
//! Create A/B Test (control vs treatment)
//!         │
//!         ▼
//! Traffic Split (50/50 deterministic hash)
//!         │
//!         ▼
//! Collect Metrics (Sharpe, win rate, PnL, drawdown)
//!         │
//!         ▼
//! Statistical Testing (Welch's t-test, p < 0.05)
//!         │
//!         ▼
//! Deployment Decision:
//!   - RolloutTreatment (treatment significantly better)
//!   - RevertToControl (treatment significantly worse)
//!   - Neutral (no significant difference)
//!   - Inconclusive (insufficient samples)
//! ```

use anyhow::{anyhow, Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info};
use uuid::Uuid;

use ml::ensemble::ab_testing::{
    ABGroup, ABTestConfig as MLABTestConfig, ABTestRouter, GroupMetrics, StatisticalTestResult,
};

/// A/B Testing Pipeline Configuration
#[derive(Debug, Clone)]
pub struct ABTestingConfig {
    /// Test ID prefix (for namespacing)
    pub test_prefix: String,

    /// Minimum sample size per group before statistical testing
    pub min_sample_size: usize,

    /// Traffic split (0.5 = 50/50)
    pub traffic_split: f64,

    /// Statistical significance level (0.05 = p < 0.05)
    pub significance_level: f64,

    /// Maximum test duration in hours
    pub max_duration_hours: u64,
}

impl Default for ABTestingConfig {
    fn default() -> Self {
        Self {
            test_prefix: "ab_test".to_string(),
            min_sample_size: 150, // Realistic minimum for statistical significance
            traffic_split: 0.5,
            significance_level: 0.05,
            max_duration_hours: 168, // 1 week
        }
    }
}

/// A/B Test State
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ABTestState {
    /// Unique test ID
    pub test_id: String,

    /// Control model ID (baseline)
    pub control_model: String,

    /// Treatment model ID (new model)
    pub treatment_model: String,

    /// Symbol being tested
    pub symbol: String,

    /// Test status (running, completed, stopped)
    pub status: String,

    /// Start time
    pub start_time: DateTime<Utc>,

    /// End time (if completed)
    pub end_time: Option<DateTime<Utc>>,
}

/// Traffic Split Metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrafficSplitMetrics {
    /// Control group metrics
    pub control: ModelPerformanceMetrics,

    /// Treatment group metrics
    pub treatment: ModelPerformanceMetrics,
}

/// Model Performance Metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelPerformanceMetrics {
    /// Total predictions
    pub predictions: u64,

    /// Correct predictions
    pub correct_predictions: u64,

    /// Win rate
    pub win_rate: f64,

    /// Total PnL
    pub total_pnl: f64,

    /// Average PnL per prediction
    pub avg_pnl: f64,

    /// Sharpe ratio (annualized)
    pub sharpe_ratio: f64,

    /// Maximum drawdown
    pub max_drawdown: f64,

    /// Average latency (microseconds)
    pub avg_latency_us: f64,
}

impl Default for ModelPerformanceMetrics {
    fn default() -> Self {
        Self {
            predictions: 0,
            correct_predictions: 0,
            win_rate: 0.0,
            total_pnl: 0.0,
            avg_pnl: 0.0,
            sharpe_ratio: 0.0,
            max_drawdown: 0.0,
            avg_latency_us: 0.0,
        }
    }
}

/// Statistical Test Results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ABStatisticalTestResult {
    /// Sharpe ratio difference (treatment - control)
    pub sharpe_diff: f64,

    /// Sharpe ratio test result
    pub sharpe_test: StatisticalTestResult,

    /// Win rate difference (treatment - control)
    pub win_rate_diff: f64,

    /// Win rate test result
    pub win_rate_test: StatisticalTestResult,

    /// PnL difference (treatment - control)
    pub pnl_diff: f64,

    /// PnL test result
    pub pnl_test: StatisticalTestResult,
}

/// Deployment Decision
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DeploymentDecision {
    /// Roll out treatment to 100% (treatment significantly better)
    RolloutTreatment {
        reason: String,
        sharpe_improvement: f64,
        pnl_improvement: f64,
        p_value: f64,
    },

    /// Revert to control (treatment significantly worse)
    RevertToControl {
        reason: String,
        sharpe_degradation: f64,
        pnl_degradation: f64,
        p_value: f64,
    },

    /// No significant difference (use simpler model)
    Neutral {
        reason: String,
        sharpe_diff: f64,
        pnl_diff: f64,
    },

    /// Insufficient samples, continue testing
    Inconclusive {
        reason: String,
        control_samples: usize,
        treatment_samples: usize,
        required_samples: usize,
    },
}

/// A/B Testing Pipeline
pub struct ABTestingPipeline {
    /// Database connection pool
    db_pool: PgPool,

    /// Configuration
    config: ABTestingConfig,

    /// Active A/B test routers (keyed by test_id)
    active_tests: Arc<RwLock<HashMap<String, Arc<ABTestRouter>>>>,

    /// Traffic group assignments (keyed by test_id + user_id)
    traffic_assignments: Arc<RwLock<HashMap<String, String>>>,
}

impl ABTestingPipeline {
    /// Create new A/B testing pipeline
    pub fn new(db_pool: PgPool, config: ABTestingConfig) -> Self {
        Self {
            db_pool,
            config,
            active_tests: Arc::new(RwLock::new(HashMap::new())),
            traffic_assignments: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Create A/B test on model deployment
    pub async fn create_ab_test(
        &self,
        control_model_id: &str,
        treatment_model_id: &str,
        symbol: &str,
    ) -> Result<ABTestState> {
        let test_id = format!("{}_{}", self.config.test_prefix, Uuid::new_v4());
        let start_time = Utc::now();

        info!(
            "Creating A/B test {} for symbol {} (control: {}, treatment: {})",
            test_id, symbol, control_model_id, treatment_model_id
        );

        // Create ML A/B test router
        let ml_config = MLABTestConfig {
            test_id: test_id.clone(),
            control_model: control_model_id.to_string(),
            treatment_model: treatment_model_id.to_string(),
            traffic_split: self.config.traffic_split,
            min_sample_size: self.config.min_sample_size,
            significance_level: self.config.significance_level,
            max_duration_hours: self.config.max_duration_hours,
            start_time: start_time.timestamp(),
        };

        let router = Arc::new(ABTestRouter::new(ml_config));

        // Store in active tests
        {
            let mut active_tests = self.active_tests.write().await;
            active_tests.insert(test_id.clone(), router.clone());
        }

        // Persist to database
        sqlx::query(
            r#"
            INSERT INTO ab_test_results (
                test_id, control_model, treatment_model, symbol,
                status, start_time, traffic_split, min_sample_size
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
            "#,
        )
        .bind(&test_id)
        .bind(control_model_id)
        .bind(treatment_model_id)
        .bind(symbol)
        .bind("running")
        .bind(start_time)
        .bind(self.config.traffic_split)
        .bind(self.config.min_sample_size as i32)
        .execute(&self.db_pool)
        .await
        .context("Failed to persist A/B test to database")?;

        Ok(ABTestState {
            test_id,
            control_model: control_model_id.to_string(),
            treatment_model: treatment_model_id.to_string(),
            symbol: symbol.to_string(),
            status: "running".to_string(),
            start_time,
            end_time: None,
        })
    }

    /// Assign traffic group (deterministic hash-based)
    pub async fn assign_traffic_group(&self, test_id: &str, user_id: &str) -> Result<String> {
        let assignment_key = format!("{}:{}", test_id, user_id);

        // Check if already assigned
        {
            let assignments = self.traffic_assignments.read().await;
            if let Some(group) = assignments.get(&assignment_key) {
                return Ok(group.clone());
            }
        }

        // Get router
        let router = {
            let active_tests = self.active_tests.read().await;
            active_tests
                .get(test_id)
                .ok_or_else(|| anyhow!("A/B test {} not found", test_id))?
                .clone()
        };

        // Assign group using ML router
        let group = router.get_or_assign_group(user_id).await;
        let group_str = match group {
            ABGroup::Control => "control".to_string(),
            ABGroup::Treatment => "treatment".to_string(),
        };

        // Cache assignment
        {
            let mut assignments = self.traffic_assignments.write().await;
            assignments.insert(assignment_key, group_str.clone());
        }

        debug!(
            "Assigned user {} to group {} for test {}",
            user_id, group_str, test_id
        );

        Ok(group_str)
    }

    /// Record prediction outcome
    pub async fn record_prediction_outcome(
        &self,
        test_id: &str,
        group: &str,
        correct: bool,
        pnl: f64,
        return_pct: f64,
        latency_us: u64,
    ) -> Result<()> {
        // Get router
        let router = {
            let active_tests = self.active_tests.read().await;
            active_tests
                .get(test_id)
                .ok_or_else(|| anyhow!("A/B test {} not found", test_id))?
                .clone()
        };

        // Convert group string to ABGroup
        let ab_group = match group {
            "control" => ABGroup::Control,
            "treatment" => ABGroup::Treatment,
            _ => return Err(anyhow!("Invalid group: {}", group)),
        };

        // Record outcome in ML router
        router
            .record_outcome(ab_group, correct, pnl, return_pct, latency_us)
            .await;

        debug!(
            "Recorded outcome for test {} group {}: correct={}, pnl={:.2}, return={:.4}",
            test_id, group, correct, pnl, return_pct
        );

        Ok(())
    }

    /// Get A/B test metrics
    pub async fn get_ab_test_metrics(&self, test_id: &str) -> Result<TrafficSplitMetrics> {
        // Get router
        let router = {
            let active_tests = self.active_tests.read().await;
            active_tests
                .get(test_id)
                .ok_or_else(|| anyhow!("A/B test {} not found", test_id))?
                .clone()
        };

        // Get ML results
        let ml_results = router
            .get_results()
            .await
            .map_err(|e| anyhow!("Failed to get ML results: {}", e))?;

        // Convert to our metrics format
        let control = Self::convert_group_metrics(&ml_results.control_group);
        let treatment = Self::convert_group_metrics(&ml_results.treatment_group);

        Ok(TrafficSplitMetrics { control, treatment })
    }

    /// Convert ML GroupMetrics to our ModelPerformanceMetrics
    fn convert_group_metrics(metrics: &GroupMetrics) -> ModelPerformanceMetrics {
        ModelPerformanceMetrics {
            predictions: metrics.predictions,
            correct_predictions: metrics.correct_predictions,
            win_rate: metrics.win_rate(),
            total_pnl: metrics.total_pnl,
            avg_pnl: metrics.avg_pnl(),
            sharpe_ratio: metrics.sharpe_ratio(),
            max_drawdown: Self::compute_max_drawdown(&metrics.pnl_samples),
            avg_latency_us: metrics.avg_latency_us,
        }
    }

    /// Compute maximum drawdown from cumulative PnL samples.
    ///
    /// Returns 0.0 when samples are empty. The result is a non-negative value
    /// representing the largest peak-to-trough decline in cumulative PnL.
    fn compute_max_drawdown(pnl_samples: &[f64]) -> f64 {
        if pnl_samples.is_empty() {
            return 0.0;
        }
        let mut cumulative = 0.0_f64;
        let mut peak = 0.0_f64;
        let mut max_dd = 0.0_f64;
        for &pnl in pnl_samples {
            cumulative += pnl;
            if cumulative > peak {
                peak = cumulative;
            }
            let drawdown = peak - cumulative;
            if drawdown > max_dd {
                max_dd = drawdown;
            }
        }
        max_dd
    }

    /// Run statistical tests
    pub async fn run_statistical_tests(&self, test_id: &str) -> Result<ABStatisticalTestResult> {
        // Get router
        let router = {
            let active_tests = self.active_tests.read().await;
            active_tests
                .get(test_id)
                .ok_or_else(|| anyhow!("A/B test {} not found", test_id))?
                .clone()
        };

        // Get ML results (includes statistical tests)
        let ml_results = router
            .get_results()
            .await
            .map_err(|e| anyhow!("Failed to get ML results: {}", e))?;

        Ok(ABStatisticalTestResult {
            sharpe_diff: ml_results.sharpe_diff,
            sharpe_test: ml_results.sharpe_test,
            win_rate_diff: ml_results.win_rate_diff,
            win_rate_test: ml_results.win_rate_test,
            pnl_diff: ml_results.pnl_diff,
            pnl_test: ml_results.pnl_test,
        })
    }

    /// Make deployment decision based on A/B test results
    pub async fn make_deployment_decision(&self, test_id: &str) -> Result<DeploymentDecision> {
        // Get metrics
        let metrics = self.get_ab_test_metrics(test_id).await?;

        // Check minimum sample size
        if metrics.control.predictions < self.config.min_sample_size as u64
            || metrics.treatment.predictions < self.config.min_sample_size as u64
        {
            return Ok(DeploymentDecision::Inconclusive {
                reason: format!(
                    "Insufficient samples. Need {} per group, got control={}, treatment={}",
                    self.config.min_sample_size,
                    metrics.control.predictions,
                    metrics.treatment.predictions
                ),
                control_samples: metrics.control.predictions as usize,
                treatment_samples: metrics.treatment.predictions as usize,
                required_samples: self.config.min_sample_size,
            });
        }

        // Run statistical tests
        let test_results = self.run_statistical_tests(test_id).await?;

        // Decision logic based on statistical significance
        let sharpe_diff = test_results.sharpe_diff;
        let pnl_diff = test_results.pnl_diff;
        let sharpe_significant = test_results.sharpe_test.is_significant;
        let pnl_significant = test_results.pnl_test.is_significant;
        let sharpe_p_value = test_results.sharpe_test.p_value;

        info!(
            "A/B test {} decision analysis: sharpe_diff={:.2}, pnl_diff={:.2}, sharpe_sig={}, pnl_sig={}, p_value={:.4}",
            test_id, sharpe_diff, pnl_diff, sharpe_significant, pnl_significant, sharpe_p_value
        );

        // Strong positive signal: both metrics significantly better
        if sharpe_significant && pnl_significant && sharpe_diff > 0.2 && pnl_diff > 0.0 {
            return Ok(DeploymentDecision::RolloutTreatment {
                reason: format!(
                    "Treatment significantly outperforms control: Sharpe +{:.2} (p={:.4}), PnL +${:.2}. Roll out to 100%.",
                    sharpe_diff, sharpe_p_value, pnl_diff
                ),
                sharpe_improvement: sharpe_diff,
                pnl_improvement: pnl_diff,
                p_value: sharpe_p_value,
            });
        }

        // Strong negative signal: both metrics significantly worse
        if sharpe_significant && pnl_significant && sharpe_diff < -0.2 && pnl_diff < 0.0 {
            return Ok(DeploymentDecision::RevertToControl {
                reason: format!(
                    "Treatment significantly underperforms control: Sharpe {:.2} (p={:.4}), PnL ${:.2}. Revert to control.",
                    sharpe_diff, sharpe_p_value, pnl_diff
                ),
                sharpe_degradation: sharpe_diff,
                pnl_degradation: pnl_diff,
                p_value: sharpe_p_value,
            });
        }

        // Moderate positive signal
        if sharpe_diff > 0.1 && pnl_diff > 0.0 {
            return Ok(DeploymentDecision::RolloutTreatment {
                reason: format!(
                    "Treatment shows improvement: Sharpe +{:.2}, PnL +${:.2}. Consider gradual rollout.",
                    sharpe_diff, pnl_diff
                ),
                sharpe_improvement: sharpe_diff,
                pnl_improvement: pnl_diff,
                p_value: sharpe_p_value,
            });
        }

        // Moderate negative signal
        if sharpe_diff < -0.1 && pnl_diff < 0.0 {
            return Ok(DeploymentDecision::RevertToControl {
                reason: format!(
                    "Treatment shows degradation: Sharpe {:.2}, PnL ${:.2}. Consider reverting.",
                    sharpe_diff, pnl_diff
                ),
                sharpe_degradation: sharpe_diff,
                pnl_degradation: pnl_diff,
                p_value: sharpe_p_value,
            });
        }

        // No meaningful difference
        Ok(DeploymentDecision::Neutral {
            reason: format!(
                "No meaningful difference detected (Sharpe diff: {:.2}, PnL diff: ${:.2}). Use simpler single-model for operational efficiency.",
                sharpe_diff, pnl_diff
            ),
            sharpe_diff,
            pnl_diff,
        })
    }

    /// Stop A/B test and persist results
    pub async fn stop_ab_test(&self, test_id: &str) -> Result<DeploymentDecision> {
        info!("Stopping A/B test {}", test_id);

        // Make final deployment decision
        let decision = self.make_deployment_decision(test_id).await?;

        // Get final metrics
        let metrics = self.get_ab_test_metrics(test_id).await?;

        // Update database
        let end_time = Utc::now();
        let status = match &decision {
            DeploymentDecision::RolloutTreatment { .. } => "completed_rollout",
            DeploymentDecision::RevertToControl { .. } => "completed_revert",
            DeploymentDecision::Neutral { .. } => "completed_neutral",
            DeploymentDecision::Inconclusive { .. } => "completed_inconclusive",
        };

        sqlx::query(
            r#"
            UPDATE ab_test_results
            SET status = $1,
                end_time = $2,
                control_predictions = $3,
                control_win_rate = $4,
                control_sharpe = $5,
                control_pnl = $6,
                treatment_predictions = $7,
                treatment_win_rate = $8,
                treatment_sharpe = $9,
                treatment_pnl = $10,
                decision = $11
            WHERE test_id = $12
            "#,
        )
        .bind(status)
        .bind(end_time)
        .bind(metrics.control.predictions as i64)
        .bind(metrics.control.win_rate)
        .bind(metrics.control.sharpe_ratio)
        .bind(metrics.control.total_pnl)
        .bind(metrics.treatment.predictions as i64)
        .bind(metrics.treatment.win_rate)
        .bind(metrics.treatment.sharpe_ratio)
        .bind(metrics.treatment.total_pnl)
        .bind(serde_json::to_string(&decision)?)
        .bind(test_id)
        .execute(&self.db_pool)
        .await
        .context("Failed to update A/B test results")?;

        // Remove from active tests
        {
            let mut active_tests = self.active_tests.write().await;
            active_tests.remove(test_id);
        }

        info!("A/B test {} stopped with decision: {:?}", test_id, decision);

        Ok(decision)
    }

    /// Get active A/B tests
    pub async fn get_active_tests(&self) -> Result<Vec<String>> {
        let active_tests = self.active_tests.read().await;
        Ok(active_tests.keys().cloned().collect())
    }

    /// Load A/B test state from database
    pub async fn load_ab_test(&self, test_id: &str) -> Result<ABTestState> {
        let result = sqlx::query_as::<_, ABTestStateRow>(
            r#"
            SELECT test_id, control_model, treatment_model, symbol,
                   status, start_time, end_time
            FROM ab_test_results
            WHERE test_id = $1
            "#,
        )
        .bind(test_id)
        .fetch_one(&self.db_pool)
        .await
        .context("Failed to load A/B test from database")?;

        Ok(ABTestState {
            test_id: result.test_id,
            control_model: result.control_model,
            treatment_model: result.treatment_model,
            symbol: result.symbol,
            status: result.status,
            start_time: result.start_time,
            end_time: result.end_time,
        })
    }
}

/// Database row for A/B test state
#[derive(Debug, sqlx::FromRow)]
struct ABTestStateRow {
    test_id: String,
    control_model: String,
    treatment_model: String,
    symbol: String,
    status: String,
    start_time: DateTime<Utc>,
    end_time: Option<DateTime<Utc>>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_config_defaults() {
        let config = ABTestingConfig::default();
        assert_eq!(config.min_sample_size, 150);
        assert_eq!(config.traffic_split, 0.5);
        assert_eq!(config.significance_level, 0.05);
    }

    #[test]
    fn test_model_performance_metrics_defaults() {
        let metrics = ModelPerformanceMetrics::default();
        assert_eq!(metrics.predictions, 0);
        assert_eq!(metrics.win_rate, 0.0);
    }
}