//! A/B Testing Framework for ML Model Deployments
//!
//! This module provides statistical A/B testing capabilities for comparing
//! model performance with traffic splitting and significance testing.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};

use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use uuid::Uuid;

use crate::{MLError, MLResult, Features, ModelPrediction, MLModel};

/// A/B test configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ABTestConfig {
    /// Test name/identifier
    pub test_name: String,
    /// Description of the test
    pub description: String,
    /// Control group percentage (0.0 to 1.0)
    pub control_percentage: f32,
    /// Treatment group percentage (0.0 to 1.0)
    pub treatment_percentage: f32,
    /// Minimum sample size per group
    pub min_sample_size: u32,
    /// Statistical significance threshold (e.g., 0.05 for 95% confidence)
    pub significance_threshold: f64,
    /// Test duration limit
    pub max_duration: Duration,
    /// Metrics to track for comparison
    pub tracked_metrics: Vec<String>,
    /// Early stopping criteria
    pub early_stopping: Option<EarlyStoppingConfig>,
    /// Traffic splitting strategy
    pub splitting_strategy: TrafficSplittingStrategy,
}

impl Default for ABTestConfig {
    fn default() -> Self {
        Self {
            test_name: "default_test".to_owned(),
            description: "Default A/B test configuration".to_owned(),
            control_percentage: 0.5,
            treatment_percentage: 0.5,
            min_sample_size: 1000,
            significance_threshold: 0.05,
            max_duration: Duration::from_secs(24 * 60 * 60),
            tracked_metrics: vec!["latency".to_owned(), "accuracy".to_owned(), "error_rate".to_owned()],
            early_stopping: Some(EarlyStoppingConfig::default()),
            splitting_strategy: TrafficSplittingStrategy::HashBased,
        }
    }
}

/// Early stopping configuration for A/B tests
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EarlyStoppingConfig {
    /// Check interval for early stopping
    pub check_interval: Duration,
    /// Minimum samples before early stopping is considered
    pub min_samples_for_early_stop: u32,
    /// Maximum degradation allowed before stopping (0.0 to 1.0)
    pub max_degradation_threshold: f64,
    /// Statistical power threshold for early stopping
    pub statistical_power_threshold: f64,
}

impl Default for EarlyStoppingConfig {
    fn default() -> Self {
        Self {
            check_interval: Duration::from_secs(5 * 60),
            min_samples_for_early_stop: 100,
            max_degradation_threshold: 0.1, // 10% degradation
            statistical_power_threshold: 0.8, // 80% power
        }
    }
}

/// Traffic splitting strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TrafficSplittingStrategy {
    /// Hash-based splitting using feature hash
    HashBased,
    /// Random splitting
    Random,
    /// Round-robin splitting
    RoundRobin,
    /// Weighted random splitting
    WeightedRandom,
    /// Thompson Sampling — Bayesian adaptive traffic allocation
    ThompsonSampling,
}

/// A/B test status
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ABTestStatus {
    /// Test is being set up
    Initializing,
    /// Test is running
    Running,
    /// Test is paused
    Paused,
    /// Test completed successfully
    Completed,
    /// Test was stopped early due to significance
    StoppedEarly,
    /// Test was stopped due to degradation
    StoppedDegradation,
    /// Test failed
    Failed,
}

/// A/B test experiment
#[derive(Debug, Clone)]
pub struct ABTestExperiment {
    /// Unique experiment ID
    pub experiment_id: Uuid,
    /// Test configuration
    pub config: ABTestConfig,
    /// Control model
    pub control_model: Arc<dyn MLModel>,
    /// Treatment model
    pub treatment_model: Arc<dyn MLModel>,
    /// Test status
    pub status: ABTestStatus,
    /// Start time
    pub start_time: SystemTime,
    /// End time (if completed)
    pub end_time: Option<SystemTime>,
    /// Traffic splitter
    pub traffic_splitter: Arc<TrafficSplitter>,
    /// Metrics collector
    pub metrics_collector: Arc<ABTestMetricsCollector>,
}

impl ABTestExperiment {
    /// Create new A/B test experiment
    pub fn new(
        config: ABTestConfig,
        control_model: Arc<dyn MLModel>,
        treatment_model: Arc<dyn MLModel>,
    ) -> Self {
        let experiment_id = Uuid::new_v4();
        let traffic_splitter = Arc::new(TrafficSplitter::new(
            config.control_percentage,
            config.treatment_percentage,
            config.splitting_strategy,
        ));
        let metrics_collector = Arc::new(ABTestMetricsCollector::new(
            experiment_id,
            config.tracked_metrics.clone(),
        ));

        Self {
            experiment_id,
            config,
            control_model,
            treatment_model,
            status: ABTestStatus::Initializing,
            start_time: SystemTime::now(),
            end_time: None,
            traffic_splitter,
            metrics_collector,
        }
    }

    /// Start the A/B test
    pub async fn start(&mut self) -> MLResult<()> {
        if self.status != ABTestStatus::Initializing {
            return Err(MLError::ValidationError {
                message: "Test can only be started from Initializing status".to_owned(),
            });
        }

        self.status = ABTestStatus::Running;
        self.start_time = SystemTime::now();

        tracing::info!(
            "Started A/B test {} with control model {} and treatment model {}",
            self.config.test_name,
            self.control_model.name(),
            self.treatment_model.name()
        );

        Ok(())
    }

    /// Perform prediction with A/B test traffic splitting
    pub async fn predict(&self, features: &Features) -> MLResult<ABTestPrediction> {
        if self.status != ABTestStatus::Running {
            return Err(MLError::ValidationError {
                message: "Test is not running".to_owned(),
            });
        }

        let start_time = Instant::now();
        
        // Determine which model to use
        let group = self.traffic_splitter.assign_group(features);
        
        let (model, _group_name) = match group {
            TestGroup::Control => (&self.control_model, "control"),
            TestGroup::Treatment => (&self.treatment_model, "treatment"),
        };

        // Make prediction
        let prediction = model.predict(features).await?;
        let latency = start_time.elapsed();

        // Record metrics
        self.metrics_collector.record_prediction(
            group,
            &prediction,
            latency,
        ).await?;

        Ok(ABTestPrediction {
            prediction,
            assigned_group: group,
            model_name: model.name().to_string(),
            latency,
            experiment_id: self.experiment_id,
        })
    }

    /// Get current test results
    pub async fn get_results(&self) -> ABTestResults {
        self.metrics_collector.get_results().await
    }

    /// Check if test should be stopped early
    pub async fn check_early_stopping(&mut self) -> MLResult<bool> {
        if let Some(ref early_config) = self.config.early_stopping {
            let results = self.get_results().await;
            
            // Check minimum samples
            if results.control_metrics.sample_count < early_config.min_samples_for_early_stop ||
               results.treatment_metrics.sample_count < early_config.min_samples_for_early_stop {
                return Ok(false);
            }

            // Check for statistical significance
            if let Some(significance) = self.calculate_statistical_significance(&results).await? {
                if significance.p_value < self.config.significance_threshold {
                    self.status = ABTestStatus::StoppedEarly;
                    self.end_time = Some(SystemTime::now());
                    
                    tracing::info!(
                        "A/B test {} stopped early due to statistical significance (p-value: {})",
                        self.config.test_name,
                        significance.p_value
                    );
                    
                    return Ok(true);
                }
            }

            // Check for degradation
            if self.check_degradation_threshold(&results, early_config.max_degradation_threshold).await? {
                self.status = ABTestStatus::StoppedDegradation;
                self.end_time = Some(SystemTime::now());
                
                tracing::warn!(
                    "A/B test {} stopped due to performance degradation",
                    self.config.test_name
                );
                
                return Ok(true);
            }
        }

        Ok(false)
    }

    /// Calculate statistical significance between groups
    async fn calculate_statistical_significance(&self, results: &ABTestResults) -> MLResult<Option<StatisticalSignificance>> {
        // For latency comparison (continuous metric)
        if let (Some(control_latency), Some(treatment_latency)) = 
           (results.control_metrics.avg_latency, results.treatment_metrics.avg_latency) {
            
            let t_stat = self.calculate_t_statistic(
                control_latency,
                treatment_latency,
                results.control_metrics.latency_std_dev.unwrap_or(0.0),
                results.treatment_metrics.latency_std_dev.unwrap_or(0.0),
                results.control_metrics.sample_count as f64,
                results.treatment_metrics.sample_count as f64,
            );

            let p_value = self.calculate_p_value(t_stat, 
                (results.control_metrics.sample_count + results.treatment_metrics.sample_count - 2) as f64);

            return Ok(Some(StatisticalSignificance {
                metric_name: "latency".to_owned(),
                t_statistic: t_stat,
                p_value,
                confidence_interval: self.calculate_confidence_interval(
                    control_latency,
                    treatment_latency,
                    results.control_metrics.latency_std_dev.unwrap_or(0.0),
                    results.treatment_metrics.latency_std_dev.unwrap_or(0.0),
                    results.control_metrics.sample_count as f64,
                    results.treatment_metrics.sample_count as f64,
                ),
                effect_size: (treatment_latency - control_latency) / control_latency,
            }));
        }

        Ok(None)
    }

    /// Calculate t-statistic for two-sample t-test
    fn calculate_t_statistic(&self, mean1: f64, mean2: f64, std1: f64, std2: f64, n1: f64, n2: f64) -> f64 {
        let pooled_variance = ((n1 - 1.0) * std1.powi(2) + (n2 - 1.0) * std2.powi(2)) / (n1 + n2 - 2.0);
        let standard_error = (pooled_variance * (1.0 / n1 + 1.0 / n2)).sqrt();
        
        if standard_error == 0.0 {
            0.0
        } else {
            (mean1 - mean2) / standard_error
        }
    }

    /// Calculate p-value from t-statistic (simplified approximation)
    fn calculate_p_value(&self, t_stat: f64, degrees_of_freedom: f64) -> f64 {
        // Simplified p-value calculation using normal approximation
        // In production, you'd use a proper statistical library
        let abs_t = t_stat.abs();
        
        if degrees_of_freedom > 30.0 {
            // Normal approximation for large samples
            2.0 * (1.0 - self.normal_cdf(abs_t))
        } else {
            // Conservative estimate for small samples
            if abs_t > 2.0 { 0.05 } else { 0.1 }
        }
    }

    /// Normal cumulative distribution function (approximation)
    fn normal_cdf(&self, x: f64) -> f64 {
        0.5 * (1.0 + erf(x / 2.0_f64.sqrt()))
    }

    /// Calculate confidence interval
    fn calculate_confidence_interval(&self, mean1: f64, mean2: f64, std1: f64, std2: f64, n1: f64, n2: f64) -> (f64, f64) {
        let diff = mean2 - mean1;
        let pooled_variance = ((n1 - 1.0) * std1.powi(2) + (n2 - 1.0) * std2.powi(2)) / (n1 + n2 - 2.0);
        let standard_error = (pooled_variance * (1.0 / n1 + 1.0 / n2)).sqrt();
        let margin_of_error = 1.96 * standard_error; // 95% confidence
        
        (diff - margin_of_error, diff + margin_of_error)
    }

    /// Check if treatment shows significant degradation
    async fn check_degradation_threshold(&self, results: &ABTestResults, threshold: f64) -> MLResult<bool> {
        // Check latency degradation
        if let (Some(control_latency), Some(treatment_latency)) = 
           (results.control_metrics.avg_latency, results.treatment_metrics.avg_latency) {
            let degradation = (treatment_latency - control_latency) / control_latency;
            if degradation > threshold {
                return Ok(true);
            }
        }

        // Check error rate degradation
        if let (Some(control_error), Some(treatment_error)) = 
           (results.control_metrics.error_rate, results.treatment_metrics.error_rate) {
            let degradation = (treatment_error - control_error) / control_error.max(0.001); // Avoid division by zero
            if degradation > threshold {
                return Ok(true);
            }
        }

        Ok(false)
    }

    /// Stop the test
    pub async fn stop(&mut self) -> MLResult<ABTestResults> {
        if self.status != ABTestStatus::Running {
            return Err(MLError::ValidationError {
                message: "Test is not running".to_owned(),
            });
        }

        self.status = ABTestStatus::Completed;
        self.end_time = Some(SystemTime::now());

        let results = self.get_results().await;
        
        tracing::info!(
            "A/B test {} completed with {} control samples and {} treatment samples",
            self.config.test_name,
            results.control_metrics.sample_count,
            results.treatment_metrics.sample_count
        );

        Ok(results)
    }
}

/// Test group assignment
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TestGroup {
    /// Control group (original model)
    Control,
    /// Treatment group (new model)
    Treatment,
}

/// Thompson Sampling state for Bayesian adaptive traffic splitting.
///
/// Maintains Beta distribution parameters for both champion (control) and
/// challenger (treatment) models. The champion starts with a strong prior
/// Beta(10,1) reflecting production confidence; the challenger starts with
/// an uninformative Beta(1,1).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThompsonSamplingState {
    /// Beta distribution alpha (successes + prior) for the champion model
    pub control_alpha: f64,
    /// Beta distribution beta (failures + prior) for the champion model
    pub control_beta: f64,
    /// Beta distribution alpha (successes + prior) for the challenger model
    pub treatment_alpha: f64,
    /// Beta distribution beta (failures + prior) for the challenger model
    pub treatment_beta: f64,
    /// Total trades observed for control
    pub control_trades: usize,
    /// Total trades observed for treatment
    pub treatment_trades: usize,
}

impl Default for ThompsonSamplingState {
    fn default() -> Self {
        Self::new()
    }
}

impl ThompsonSamplingState {
    /// Create a new Thompson Sampling state with default priors.
    ///
    /// Champion: Beta(10, 1) — strong prior reflecting production confidence.
    /// Challenger: Beta(1, 1) — uninformative prior.
    pub fn new() -> Self {
        Self {
            control_alpha: 10.0,
            control_beta: 1.0,
            treatment_alpha: 1.0,
            treatment_beta: 1.0,
            control_trades: 0,
            treatment_trades: 0,
        }
    }

    /// Sample from both Beta distributions and return the group whose
    /// sample is higher. This is the core Thompson Sampling decision:
    /// models that are believed to be better get more traffic, but
    /// exploration is preserved via randomness.
    ///
    /// Uses the Gamma-distribution trick: Beta(a,b) = X/(X+Y) where
    /// X ~ Gamma(a,1), Y ~ Gamma(b,1).
    pub fn sample_and_assign(&self) -> TestGroup {
        let control_sample = self.sample_beta(self.control_alpha, self.control_beta);
        let treatment_sample = self.sample_beta(self.treatment_alpha, self.treatment_beta);

        if control_sample >= treatment_sample {
            TestGroup::Control
        } else {
            TestGroup::Treatment
        }
    }

    /// Record the outcome of a trade for a given group.
    ///
    /// On success: alpha += 1 (evidence of reward).
    /// On failure: beta += 1 (evidence of non-reward).
    pub fn record_outcome(&mut self, group: TestGroup, success: bool) {
        match group {
            TestGroup::Control => {
                self.control_trades += 1;
                if success {
                    self.control_alpha += 1.0;
                } else {
                    self.control_beta += 1.0;
                }
            }
            TestGroup::Treatment => {
                self.treatment_trades += 1;
                if success {
                    self.treatment_alpha += 1.0;
                } else {
                    self.treatment_beta += 1.0;
                }
            }
        }
    }

    /// Sample from a Beta(alpha, beta) distribution using the Gamma trick.
    ///
    /// Beta(a,b) = X/(X+Y) where X ~ Gamma(a,1) and Y ~ Gamma(b,1).
    fn sample_beta(&self, alpha: f64, beta: f64) -> f64 {
        use rand::Rng;
        use rand_distr::{Distribution, Gamma};

        let mut rng = rand::thread_rng();

        // Alpha and beta must be > 0 for Gamma. Clamp to epsilon.
        let safe_alpha = alpha.max(1e-10);
        let safe_beta = beta.max(1e-10);

        let gamma_a = match Gamma::new(safe_alpha, 1.0) {
            Ok(d) => d.sample(&mut rng),
            Err(_) => rng.gen::<f64>(), // fallback: uniform [0,1)
        };
        let gamma_b = match Gamma::new(safe_beta, 1.0) {
            Ok(d) => d.sample(&mut rng),
            Err(_) => rng.gen::<f64>(),
        };

        let sum = gamma_a + gamma_b;
        if sum <= 0.0 {
            0.5 // degenerate case — return neutral
        } else {
            gamma_a / sum
        }
    }

    /// Estimate P(treatment > control) via Monte Carlo sampling.
    ///
    /// Draws `n_samples` paired samples from both Beta distributions and
    /// returns the fraction where treatment sample exceeds control sample.
    pub fn estimate_probability_treatment_better(&self, n_samples: usize) -> f64 {
        if n_samples == 0 {
            return 0.0;
        }

        let mut treatment_wins = 0_usize;
        for _ in 0..n_samples {
            let ctrl = self.sample_beta(self.control_alpha, self.control_beta);
            let treat = self.sample_beta(self.treatment_alpha, self.treatment_beta);
            if treat > ctrl {
                treatment_wins += 1;
            }
        }

        treatment_wins as f64 / n_samples as f64
    }
}

/// Decision returned by promotion criteria evaluation.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PromotionDecision {
    /// Not enough evidence yet to promote
    NotReady {
        /// Explanation of why promotion is not ready
        reason: String,
    },
    /// All statistical criteria met; awaiting manual approval
    ReadyForApproval,
    /// All criteria met and manual approval not required — auto-promoted
    AutoPromoted,
}

/// Criteria for promoting a challenger model to champion.
///
/// All conditions must be satisfied simultaneously:
/// 1. Minimum trade count on the challenger
/// 2. Bayesian confidence threshold (P(challenger > champion) > threshold)
/// 3. Maximum drawdown ratio constraint
/// 4. Optional manual approval gate
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromotionCriteria {
    /// Minimum number of trades the challenger must have observed
    pub min_trades: usize,
    /// Required P(challenger > champion) for promotion (e.g. 0.95)
    pub bayesian_confidence: f64,
    /// Maximum allowed ratio of challenger DD to champion DD
    pub max_drawdown_ratio: f64,
    /// Whether a human must manually approve after criteria are met
    pub require_manual_approval: bool,
}

impl Default for PromotionCriteria {
    fn default() -> Self {
        Self {
            min_trades: 500,
            bayesian_confidence: 0.95,
            max_drawdown_ratio: 1.2,
            require_manual_approval: true,
        }
    }
}

impl PromotionCriteria {
    /// Evaluate whether the challenger should be promoted.
    ///
    /// Returns `NotReady` with a reason if any criterion fails,
    /// `ReadyForApproval` if manual approval is required, or
    /// `AutoPromoted` if all criteria pass and no approval needed.
    ///
    /// `n_mc_samples` controls Monte Carlo accuracy for Bayesian confidence
    /// (default: 10_000).
    pub fn should_promote(
        &self,
        state: &ThompsonSamplingState,
        challenger_dd: f64,
        champion_dd: f64,
    ) -> PromotionDecision {
        self.should_promote_with_samples(state, challenger_dd, champion_dd, 10_000)
    }

    /// Same as `should_promote` but with configurable MC sample count
    /// (useful for deterministic testing).
    pub fn should_promote_with_samples(
        &self,
        state: &ThompsonSamplingState,
        challenger_dd: f64,
        champion_dd: f64,
        n_mc_samples: usize,
    ) -> PromotionDecision {
        // 1. Check minimum trades
        if state.treatment_trades < self.min_trades {
            return PromotionDecision::NotReady {
                reason: format!(
                    "Insufficient trades: {} / {} required",
                    state.treatment_trades, self.min_trades
                ),
            };
        }

        // 2. Bayesian confidence: P(challenger > champion)
        let p_better = state.estimate_probability_treatment_better(n_mc_samples);
        if p_better < self.bayesian_confidence {
            return PromotionDecision::NotReady {
                reason: format!(
                    "Bayesian confidence too low: {:.4} < {:.4}",
                    p_better, self.bayesian_confidence
                ),
            };
        }

        // 3. Drawdown constraint: challenger DD must not exceed ratio * champion DD.
        //    When champion_dd == 0, dd_limit == 0 so any positive challenger_dd fails
        //    (correctly strict). Both == 0 passes (both perfect — valid).
        let dd_limit = champion_dd * self.max_drawdown_ratio;
        if challenger_dd > dd_limit {
            return PromotionDecision::NotReady {
                reason: format!(
                    "Challenger drawdown {:.4} exceeds {:.1}x champion drawdown {:.4} (limit {:.4})",
                    challenger_dd, self.max_drawdown_ratio, champion_dd, dd_limit
                ),
            };
        }

        // All criteria met
        if self.require_manual_approval {
            PromotionDecision::ReadyForApproval
        } else {
            PromotionDecision::AutoPromoted
        }
    }
}

/// Traffic splitter for A/B testing
#[derive(Debug)]
pub struct TrafficSplitter {
    /// Control group percentage
    control_percentage: f32,
    /// Treatment group percentage
    treatment_percentage: f32,
    /// Splitting strategy
    strategy: TrafficSplittingStrategy,
    /// Round-robin counter
    round_robin_counter: AtomicU64,
    /// Thompson Sampling state (only present when strategy == ThompsonSampling).
    /// Wrapped in Mutex for interior mutability through Arc, enabling
    /// `record_outcome` to update Beta distributions during inference.
    thompson_state: Option<std::sync::Mutex<ThompsonSamplingState>>,
}

impl TrafficSplitter {
    /// Create new traffic splitter
    pub fn new(
        control_percentage: f32,
        treatment_percentage: f32,
        strategy: TrafficSplittingStrategy,
    ) -> Self {
        let thompson_state =
            (strategy == TrafficSplittingStrategy::ThompsonSampling)
                .then(|| std::sync::Mutex::new(ThompsonSamplingState::new()));

        Self {
            control_percentage,
            treatment_percentage,
            strategy,
            round_robin_counter: AtomicU64::new(0),
            thompson_state,
        }
    }

    /// Access the Thompson Sampling state lock (if active).
    ///
    /// Callers should use `record_outcome` to update the state after each trade.
    pub fn thompson_state(&self) -> Option<&std::sync::Mutex<ThompsonSamplingState>> {
        self.thompson_state.as_ref()
    }

    /// Record an outcome for the given test group in the Thompson Sampling state.
    ///
    /// No-op if the strategy is not ThompsonSampling.
    pub fn record_outcome(&self, group: TestGroup, success: bool) {
        if let Some(ref mutex) = self.thompson_state {
            if let Ok(mut state) = mutex.lock() {
                state.record_outcome(group, success);
            }
        }
    }

    /// Assign test group for given features
    pub fn assign_group(&self, features: &Features) -> TestGroup {
        match self.strategy {
            TrafficSplittingStrategy::HashBased => self.hash_based_assignment(features),
            TrafficSplittingStrategy::Random => self.random_assignment(),
            TrafficSplittingStrategy::RoundRobin => self.round_robin_assignment(),
            TrafficSplittingStrategy::WeightedRandom => self.weighted_random_assignment(),
            TrafficSplittingStrategy::ThompsonSampling => {
                if let Some(ref mutex) = self.thompson_state {
                    if let Ok(state) = mutex.lock() {
                        state.sample_and_assign()
                    } else {
                        self.random_assignment()
                    }
                } else {
                    // Fallback to random if state is somehow missing
                    self.random_assignment()
                }
            }
        }
    }

    /// Hash-based assignment using feature hash
    fn hash_based_assignment(&self, features: &Features) -> TestGroup {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        
        // Hash feature values
        for value in &features.values {
            value.to_bits().hash(&mut hasher);
        }
        
        // Hash timestamp for additional randomness
        features.timestamp.hash(&mut hasher);

        let hash_value = hasher.finish();
        let normalized = (hash_value % 1000) as f32 / 1000.0;

        if normalized < self.control_percentage {
            TestGroup::Control
        } else {
            TestGroup::Treatment
        }
    }

    /// Random assignment
    fn random_assignment(&self) -> TestGroup {
        let random_value: f32 = rand::random();
        if random_value < self.control_percentage {
            TestGroup::Control
        } else {
            TestGroup::Treatment
        }
    }

    /// Round-robin assignment
    fn round_robin_assignment(&self) -> TestGroup {
        let counter = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);
        let normalized = (counter % 1000) as f32 / 1000.0;

        if normalized < self.control_percentage {
            TestGroup::Control
        } else {
            TestGroup::Treatment
        }
    }

    /// Weighted random assignment
    fn weighted_random_assignment(&self) -> TestGroup {
        // Similar to random but with more sophisticated weighting
        self.random_assignment()
    }
}

/// A/B test prediction result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ABTestPrediction {
    /// Underlying model prediction
    pub prediction: ModelPrediction,
    /// Assigned test group
    pub assigned_group: TestGroup,
    /// Model name used
    pub model_name: String,
    /// Prediction latency
    pub latency: Duration,
    /// Experiment ID
    pub experiment_id: Uuid,
}

/// Metrics collector for A/B tests
#[derive(Debug)]
pub struct ABTestMetricsCollector {
    /// Experiment ID
    experiment_id: Uuid,
    /// Tracked metrics
    tracked_metrics: Vec<String>,
    /// Control group metrics
    control_metrics: Arc<Mutex<GroupMetrics>>,
    /// Treatment group metrics
    treatment_metrics: Arc<Mutex<GroupMetrics>>,
}

/// Metrics for a test group
#[derive(Debug, Clone, Default)]
struct GroupMetrics {
    /// Total predictions
    sample_count: u32,
    /// Total errors
    error_count: u32,
    /// Latency measurements
    latencies: Vec<f64>,
    /// Accuracy scores
    accuracy_scores: Vec<f64>,
    /// Custom metrics
    custom_metrics: HashMap<String, Vec<f64>>,
}

impl ABTestMetricsCollector {
    /// Create new metrics collector
    pub fn new(experiment_id: Uuid, tracked_metrics: Vec<String>) -> Self {
        Self {
            experiment_id,
            tracked_metrics,
            control_metrics: Arc::new(Mutex::new(GroupMetrics::default())),
            treatment_metrics: Arc::new(Mutex::new(GroupMetrics::default())),
        }
    }

    /// Record prediction metrics
    pub async fn record_prediction(
        &self,
        group: TestGroup,
        prediction: &ModelPrediction,
        latency: Duration,
    ) -> MLResult<()> {
        let metrics = match group {
            TestGroup::Control => &self.control_metrics,
            TestGroup::Treatment => &self.treatment_metrics,
        };

        let mut group_metrics = metrics.lock().await;
        group_metrics.sample_count += 1;
        group_metrics.latencies.push(latency.as_micros() as f64);
        group_metrics.accuracy_scores.push(prediction.confidence);

        Ok(())
    }

    /// Record error
    pub async fn record_error(&self, group: TestGroup) -> MLResult<()> {
        let metrics = match group {
            TestGroup::Control => &self.control_metrics,
            TestGroup::Treatment => &self.treatment_metrics,
        };

        let mut group_metrics = metrics.lock().await;
        group_metrics.error_count += 1;

        Ok(())
    }

    /// Get test results
    pub async fn get_results(&self) -> ABTestResults {
        let control_metrics = self.control_metrics.lock().await;
        let treatment_metrics = self.treatment_metrics.lock().await;

        ABTestResults {
            experiment_id: self.experiment_id,
            control_metrics: GroupMetricsSummary::from_group_metrics(&control_metrics),
            treatment_metrics: GroupMetricsSummary::from_group_metrics(&treatment_metrics),
            test_duration: SystemTime::now(),
        }
    }
}

/// Summary of group metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GroupMetricsSummary {
    /// Sample count
    pub sample_count: u32,
    /// Error count
    pub error_count: u32,
    /// Error rate
    pub error_rate: Option<f64>,
    /// Average latency
    pub avg_latency: Option<f64>,
    /// Latency standard deviation
    pub latency_std_dev: Option<f64>,
    /// Average accuracy
    pub avg_accuracy: Option<f64>,
    /// Accuracy standard deviation
    pub accuracy_std_dev: Option<f64>,
    /// Custom metrics
    pub custom_metrics: HashMap<String, f64>,
}

impl GroupMetricsSummary {
    /// Create summary from group metrics
    fn from_group_metrics(metrics: &GroupMetrics) -> Self {
        let error_rate = (metrics.sample_count > 0)
            .then(|| metrics.error_count as f64 / metrics.sample_count as f64);

        let (avg_latency, latency_std_dev) = if !metrics.latencies.is_empty() {
            let avg = metrics.latencies.iter().sum::<f64>() / metrics.latencies.len() as f64;
            let variance = metrics.latencies.iter()
                .map(|x| (x - avg).powi(2))
                .sum::<f64>() / metrics.latencies.len() as f64;
            (Some(avg), Some(variance.sqrt()))
        } else {
            (None, None)
        };

        let (avg_accuracy, accuracy_std_dev) = if !metrics.accuracy_scores.is_empty() {
            let avg = metrics.accuracy_scores.iter().sum::<f64>() / metrics.accuracy_scores.len() as f64;
            let variance = metrics.accuracy_scores.iter()
                .map(|x| (x - avg).powi(2))
                .sum::<f64>() / metrics.accuracy_scores.len() as f64;
            (Some(avg), Some(variance.sqrt()))
        } else {
            (None, None)
        };

        Self {
            sample_count: metrics.sample_count,
            error_count: metrics.error_count,
            error_rate,
            avg_latency,
            latency_std_dev,
            avg_accuracy,
            accuracy_std_dev,
            custom_metrics: HashMap::new(), // Would be populated from metrics.custom_metrics
        }
    }
}

/// A/B test results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ABTestResults {
    /// Experiment ID
    pub experiment_id: Uuid,
    /// Control group results
    pub control_metrics: GroupMetricsSummary,
    /// Treatment group results
    pub treatment_metrics: GroupMetricsSummary,
    /// Test duration
    pub test_duration: SystemTime,
}

/// Statistical significance result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StatisticalSignificance {
    /// Metric name
    pub metric_name: String,
    /// T-statistic
    pub t_statistic: f64,
    /// P-value
    pub p_value: f64,
    /// Confidence interval (lower, upper)
    pub confidence_interval: (f64, f64),
    /// Effect size
    pub effect_size: f64,
}

/// Error function approximation for normal CDF
fn erf(x: f64) -> f64 {
    // Abramowitz and Stegun approximation
    let a1 = 0.254829592;
    let a2 = -0.284496736;
    let a3 = 1.421413741;
    let a4 = -1.453152027;
    let a5 = 1.061405429;
    let p = 0.3275911;

    let sign = if x < 0.0 { -1.0 } else { 1.0 };
    let x = x.abs();

    let t = 1.0 / (1.0 + p * x);
    let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp();

    sign * y
}

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

    #[test]
    fn test_traffic_splitter_creation() {
        let splitter = TrafficSplitter::new(0.5, 0.5, TrafficSplittingStrategy::Random);
        
        // Test that assignments are roughly balanced over many iterations
        let mut control_count = 0;
        let mut treatment_count = 0;
        
        for _ in 0..1000 {
            let features = Features::new(vec![1.0, 2.0], vec!["f1".to_owned(), "f2".to_owned()]);
            match splitter.assign_group(&features) {
                TestGroup::Control => control_count += 1,
                TestGroup::Treatment => treatment_count += 1,
            }
        }
        
        // Should be roughly balanced (within 20% of expected)
        let total = control_count + treatment_count;
        let control_ratio = control_count as f64 / total as f64;
        assert!(control_ratio > 0.3 && control_ratio < 0.7);
    }

    #[test]
    fn test_group_metrics_summary() {
        let mut metrics = GroupMetrics::default();
        metrics.sample_count = 100;
        metrics.error_count = 5;
        metrics.latencies = vec![100.0, 200.0, 150.0, 175.0, 125.0];
        metrics.accuracy_scores = vec![0.8, 0.85, 0.9, 0.75, 0.88];

        let summary = GroupMetricsSummary::from_group_metrics(&metrics);
        
        assert_eq!(summary.sample_count, 100);
        assert_eq!(summary.error_count, 5);
        assert_eq!(summary.error_rate, Some(0.05));
        assert!(summary.avg_latency.is_some());
        assert!(summary.avg_accuracy.is_some());
    }

    #[tokio::test]
    async fn test_ab_test_experiment_creation() {
        let control_model = Arc::from(model_factory::create_dqn_wrapper().unwrap());
        let treatment_model = Arc::from(model_factory::create_dqn_wrapper().unwrap());
        
        let config = ABTestConfig::default();
        let experiment = ABTestExperiment::new(config, control_model, treatment_model);
        
        assert_eq!(experiment.status, ABTestStatus::Initializing);
        assert!(experiment.end_time.is_none());
    }

    #[test]
    fn test_statistical_calculations() {
        let experiment = ABTestExperiment::new(
            ABTestConfig::default(),
            Arc::from(model_factory::create_dqn_wrapper().unwrap()),
            Arc::from(model_factory::create_dqn_wrapper().unwrap()),
        );

        // Test t-statistic calculation
        let t_stat = experiment.calculate_t_statistic(100.0, 110.0, 10.0, 12.0, 50.0, 50.0);
        assert!(t_stat.is_finite());

        // Test confidence interval calculation
        let (lower, upper) = experiment.calculate_confidence_interval(100.0, 110.0, 10.0, 12.0, 50.0, 50.0);
        assert!(lower < upper);
    }

    // ---------------------------------------------------------------
    // Thompson Sampling tests
    // ---------------------------------------------------------------

    #[test]
    fn test_thompson_sampling_state_new() {
        let state = ThompsonSamplingState::new();
        // Champion: strong prior Beta(10, 1)
        assert!((state.control_alpha - 10.0).abs() < f64::EPSILON);
        assert!((state.control_beta - 1.0).abs() < f64::EPSILON);
        // Challenger: uninformative prior Beta(1, 1)
        assert!((state.treatment_alpha - 1.0).abs() < f64::EPSILON);
        assert!((state.treatment_beta - 1.0).abs() < f64::EPSILON);
        assert_eq!(state.control_trades, 0);
        assert_eq!(state.treatment_trades, 0);
    }

    #[test]
    fn test_thompson_sampling_default_matches_new() {
        let from_new = ThompsonSamplingState::new();
        let from_default = ThompsonSamplingState::default();
        assert!((from_new.control_alpha - from_default.control_alpha).abs() < f64::EPSILON);
        assert!((from_new.treatment_alpha - from_default.treatment_alpha).abs() < f64::EPSILON);
    }

    #[test]
    fn test_thompson_sampling_record_outcome() {
        let mut state = ThompsonSamplingState::new();
        let orig_ctrl_alpha = state.control_alpha;
        let orig_ctrl_beta = state.control_beta;
        let orig_treat_alpha = state.treatment_alpha;
        let orig_treat_beta = state.treatment_beta;

        // Control success: alpha should increase by 1
        state.record_outcome(TestGroup::Control, true);
        assert!((state.control_alpha - (orig_ctrl_alpha + 1.0)).abs() < f64::EPSILON);
        assert!((state.control_beta - orig_ctrl_beta).abs() < f64::EPSILON);
        assert_eq!(state.control_trades, 1);

        // Control failure: beta should increase by 1
        state.record_outcome(TestGroup::Control, false);
        assert!((state.control_beta - (orig_ctrl_beta + 1.0)).abs() < f64::EPSILON);
        assert_eq!(state.control_trades, 2);

        // Treatment success
        state.record_outcome(TestGroup::Treatment, true);
        assert!((state.treatment_alpha - (orig_treat_alpha + 1.0)).abs() < f64::EPSILON);
        assert!((state.treatment_beta - orig_treat_beta).abs() < f64::EPSILON);
        assert_eq!(state.treatment_trades, 1);

        // Treatment failure
        state.record_outcome(TestGroup::Treatment, false);
        assert!((state.treatment_beta - (orig_treat_beta + 1.0)).abs() < f64::EPSILON);
        assert_eq!(state.treatment_trades, 2);
    }

    #[test]
    fn test_thompson_sampling_returns_valid_group() {
        let state = ThompsonSamplingState::new();
        // Run many samples — each must return a valid group
        for _ in 0..500 {
            let group = state.sample_and_assign();
            assert!(group == TestGroup::Control || group == TestGroup::Treatment);
        }
    }

    #[test]
    fn test_thompson_sampling_strong_prior() {
        // Champion with overwhelming evidence Beta(100, 1) vs Challenger Beta(1, 1)
        let state = ThompsonSamplingState {
            control_alpha: 100.0,
            control_beta: 1.0,
            treatment_alpha: 1.0,
            treatment_beta: 1.0,
            control_trades: 0,
            treatment_trades: 0,
        };

        let mut control_count = 0;
        let iterations = 1000;
        for _ in 0..iterations {
            if state.sample_and_assign() == TestGroup::Control {
                control_count += 1;
            }
        }

        // With Beta(100,1) vs Beta(1,1), champion should win >80% of the time
        let control_ratio = control_count as f64 / iterations as f64;
        assert!(
            control_ratio > 0.80,
            "Champion with Beta(100,1) should dominate, but ratio was {control_ratio:.3}"
        );
    }

    #[test]
    fn test_thompson_sampling_estimate_probability() {
        // Treatment is much better: Beta(50, 5) vs Control Beta(5, 50)
        let state = ThompsonSamplingState {
            control_alpha: 5.0,
            control_beta: 50.0,
            treatment_alpha: 50.0,
            treatment_beta: 5.0,
            control_trades: 55,
            treatment_trades: 55,
        };

        let p = state.estimate_probability_treatment_better(5000);
        // Treatment should be better with very high probability
        assert!(
            p > 0.95,
            "P(treatment > control) should be > 0.95, got {p:.4}"
        );
    }

    #[test]
    fn test_promotion_criteria_defaults() {
        let criteria = PromotionCriteria::default();
        assert_eq!(criteria.min_trades, 500);
        assert!((criteria.bayesian_confidence - 0.95).abs() < f64::EPSILON);
        assert!((criteria.max_drawdown_ratio - 1.2).abs() < f64::EPSILON);
        assert!(criteria.require_manual_approval);
    }

    #[test]
    fn test_promotion_not_ready_insufficient_trades() {
        let criteria = PromotionCriteria::default();
        let state = ThompsonSamplingState {
            control_alpha: 1.0,
            control_beta: 1.0,
            treatment_alpha: 100.0,
            treatment_beta: 1.0,
            control_trades: 1000,
            treatment_trades: 499, // below min_trades (500)
        };

        let decision = criteria.should_promote(&state, 0.05, 0.10);
        match decision {
            PromotionDecision::NotReady { reason } => {
                assert!(
                    reason.contains("Insufficient trades"),
                    "Expected insufficient trades reason, got: {reason}"
                );
            }
            other => panic!("Expected NotReady, got {other:?}"),
        }
    }

    #[test]
    fn test_promotion_not_ready_low_confidence() {
        let criteria = PromotionCriteria {
            min_trades: 10,
            bayesian_confidence: 0.95,
            max_drawdown_ratio: 2.0,
            require_manual_approval: false,
        };
        // Treatment is much worse: Beta(1, 50) — will fail confidence check
        let state = ThompsonSamplingState {
            control_alpha: 50.0,
            control_beta: 1.0,
            treatment_alpha: 1.0,
            treatment_beta: 50.0,
            control_trades: 51,
            treatment_trades: 51,
        };

        let decision = criteria.should_promote_with_samples(&state, 0.01, 0.10, 5000);
        match decision {
            PromotionDecision::NotReady { reason } => {
                assert!(
                    reason.contains("confidence"),
                    "Expected confidence reason, got: {reason}"
                );
            }
            other => panic!("Expected NotReady due to low confidence, got {other:?}"),
        }
    }

    #[test]
    fn test_promotion_drawdown_check() {
        let criteria = PromotionCriteria {
            min_trades: 10,
            bayesian_confidence: 0.01, // very low threshold so confidence always passes
            max_drawdown_ratio: 1.2,
            require_manual_approval: false,
        };
        // Treatment has good stats
        let state = ThompsonSamplingState {
            control_alpha: 10.0,
            control_beta: 10.0,
            treatment_alpha: 50.0,
            treatment_beta: 5.0,
            control_trades: 20,
            treatment_trades: 55,
        };

        // Challenger DD (0.25) > 1.2 * Champion DD (0.10) = 0.12 => should fail
        let decision = criteria.should_promote_with_samples(&state, 0.25, 0.10, 5000);
        match decision {
            PromotionDecision::NotReady { reason } => {
                assert!(
                    reason.contains("drawdown"),
                    "Expected drawdown reason, got: {reason}"
                );
            }
            other => panic!("Expected NotReady due to drawdown, got {other:?}"),
        }

        // Challenger DD (0.11) <= 1.2 * 0.10 = 0.12 => should pass
        let decision = criteria.should_promote_with_samples(&state, 0.11, 0.10, 5000);
        assert_eq!(decision, PromotionDecision::AutoPromoted);
    }

    #[test]
    fn test_promotion_ready_for_approval() {
        let criteria = PromotionCriteria {
            min_trades: 10,
            bayesian_confidence: 0.01,
            max_drawdown_ratio: 2.0,
            require_manual_approval: true,
        };
        let state = ThompsonSamplingState {
            control_alpha: 10.0,
            control_beta: 10.0,
            treatment_alpha: 50.0,
            treatment_beta: 5.0,
            control_trades: 20,
            treatment_trades: 55,
        };

        let decision = criteria.should_promote_with_samples(&state, 0.05, 0.10, 5000);
        assert_eq!(decision, PromotionDecision::ReadyForApproval);
    }

    #[test]
    fn test_promotion_auto_promoted() {
        let criteria = PromotionCriteria {
            min_trades: 10,
            bayesian_confidence: 0.01,
            max_drawdown_ratio: 2.0,
            require_manual_approval: false,
        };
        let state = ThompsonSamplingState {
            control_alpha: 10.0,
            control_beta: 10.0,
            treatment_alpha: 50.0,
            treatment_beta: 5.0,
            control_trades: 20,
            treatment_trades: 55,
        };

        let decision = criteria.should_promote_with_samples(&state, 0.05, 0.10, 5000);
        assert_eq!(decision, PromotionDecision::AutoPromoted);
    }

    #[test]
    fn test_traffic_splitter_thompson() {
        let splitter = TrafficSplitter::new(0.5, 0.5, TrafficSplittingStrategy::ThompsonSampling);

        // Thompson state should be initialized
        assert!(splitter.thompson_state().is_some());

        let mut control_count = 0;
        let mut treatment_count = 0;
        for _ in 0..500 {
            let features = Features::new(vec![1.0, 2.0], vec!["f1".to_owned(), "f2".to_owned()]);
            match splitter.assign_group(&features) {
                TestGroup::Control => control_count += 1,
                TestGroup::Treatment => treatment_count += 1,
            }
        }

        // With default priors Beta(10,1) vs Beta(1,1), control should dominate
        // but both groups should get some traffic (exploration)
        let total = control_count + treatment_count;
        assert_eq!(total, 500);
        // Control should get the majority with Beta(10,1)
        assert!(
            control_count > treatment_count,
            "Control ({control_count}) should exceed Treatment ({treatment_count}) with strong prior"
        );
        // Treatment should still get some traffic (Thompson Sampling explores)
        assert!(
            treatment_count > 0,
            "Treatment should get some exploration traffic"
        );
    }

    #[test]
    fn test_traffic_splitter_non_thompson_has_no_state() {
        let splitter = TrafficSplitter::new(0.5, 0.5, TrafficSplittingStrategy::Random);
        assert!(splitter.thompson_state().is_none());

        let splitter = TrafficSplitter::new(0.5, 0.5, TrafficSplittingStrategy::HashBased);
        assert!(splitter.thompson_state().is_none());
    }
}
