//! Lightweight synchronous ensemble coordinator
//!
//! Wraps N [`ModelInferenceAdapter`] instances and aggregates their
//! predictions via confidence-weighted voting.  Models that are not
//! ready or that fail inference are gracefully skipped with warnings.

use std::collections::HashMap;

use rayon::prelude::*;

use crate::inference_adapter::{
    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction,
};
use crate::{MLError, MLResult};

/// Synchronous coordinator that aggregates predictions from multiple
/// [`ModelInferenceAdapter`] instances using confidence-weighted voting.
///
/// # Aggregation formula
///
/// ```text
/// weighted_direction = sum(direction_i * weight_i * confidence_i)
///                    / sum(weight_i * confidence_i)
/// ```
///
/// where `weight_i` defaults to 1.0 and can be overridden via
/// [`set_weight`](InferenceEnsemble::set_weight).
#[allow(missing_debug_implementations)]
pub struct InferenceEnsemble {
    adapters: Vec<Box<dyn ModelInferenceAdapter>>,
    weights: HashMap<String, f64>,
}

impl InferenceEnsemble {
    /// Create a new ensemble from a vector of adapters.
    ///
    /// All adapters start with a default weight of 1.0.
    pub fn new(adapters: Vec<Box<dyn ModelInferenceAdapter>>) -> Self {
        Self {
            adapters,
            weights: HashMap::new(),
        }
    }

    /// Set a custom weight for a model identified by name.
    ///
    /// Models without an explicit weight use 1.0.
    pub fn set_weight(&mut self, model_name: &str, weight: f64) {
        self.weights.insert(model_name.to_string(), weight);
    }

    /// Return the number of adapters whose `is_ready()` returns true.
    pub fn ready_count(&self) -> usize {
        self.adapters.iter().filter(|a| a.is_ready()).count()
    }

    /// Run inference across all ready adapters and aggregate via
    /// confidence-weighted voting.
    ///
    /// Uses [`predict_raw()`](ModelInferenceAdapter::predict_raw) to collect
    /// raw logits where available.  Models that return logits are
    /// aggregated via sigmoid -> weighted-sum.  Models without logits
    /// fall back to the original scalar-weighted-average path.  The two
    /// paths are merged by model-count-weighted average.
    ///
    /// Returns [`MLError::InferenceError`] if no models are ready.
    pub fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
        let ready_adapters: Vec<&Box<dyn ModelInferenceAdapter>> =
            self.adapters.iter().filter(|a| a.is_ready()).collect();

        if ready_adapters.is_empty() {
            return Err(MLError::InferenceError(
                "No models are ready for inference".to_owned(),
            ));
        }

        // Run predict_raw() in parallel across all ready adapters
        let raw_predictions: Vec<(String, RawPrediction)> = ready_adapters
            .par_iter()
            .filter_map(|adapter| {
                let model_name = adapter.model_name().to_string();
                match adapter.predict_raw(features) {
                    Ok(pred) => {
                        if !pred.direction_scalar.is_finite() || !pred.confidence.is_finite() {
                            tracing::warn!(
                                model = %model_name,
                                direction = %pred.direction_scalar,
                                confidence = %pred.confidence,
                                "Model returned NaN/Inf prediction, skipping (circuit breaker)"
                            );
                            None
                        } else {
                            Some((model_name, pred))
                        }
                    }
                    Err(e) => {
                        tracing::warn!(
                            model = %model_name,
                            error = %e,
                            "Model prediction failed, skipping"
                        );
                        None
                    }
                }
            })
            .collect();

        if raw_predictions.is_empty() {
            return Err(MLError::InferenceError(
                "All ready models failed during inference".to_owned(),
            ));
        }

        // Partition into logits-available vs CPU-scalar predictions
        let (logits_preds, cpu_preds): (Vec<_>, Vec<_>) = raw_predictions
            .into_iter()
            .partition(|(_, p)| p.logits.is_some());

        // Collect model names from both paths for the ensemble label
        let mut model_names: Vec<String> = Vec::new();
        let mut total_confidence_sum = 0.0_f64;
        let mut total_count: usize = 0;

        // --- Logits path: apply sigmoid, weighted-sum ---
        let logits_result = if !logits_preds.is_empty() {
            let result = self.aggregate_logits(&logits_preds);
            for (name, pred) in &logits_preds {
                model_names.push(name.clone());
                total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
            }
            total_count += logits_preds.len();
            result
        } else {
            None
        };

        // --- CPU path: existing scalar weighted average ---
        let cpu_result = if !cpu_preds.is_empty() {
            let result = self.aggregate_cpu(&cpu_preds);
            for (name, pred) in &cpu_preds {
                model_names.push(name.clone());
                total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
            }
            total_count += cpu_preds.len();
            result
        } else {
            None
        };

        // Merge logits and CPU directions by model-count-weighted average
        let direction = match (logits_result, cpu_result) {
            (Some((logits_dir, logits_n)), Some((cpu_dir, cpu_n))) => {
                let total_n = (logits_n + cpu_n) as f64;
                (logits_dir * logits_n as f64 + cpu_dir * cpu_n as f64) / total_n
            }
            (Some((dir, _)), None) => dir,
            (None, Some((dir, _))) => dir,
            (None, None) => 0.0,
        };

        let avg_confidence = if total_count > 0 {
            total_confidence_sum / total_count as f64
        } else {
            0.0
        };

        let ensemble_name = format!("ENSEMBLE({})", model_names.join("+"));

        Ok(EnsemblePrediction {
            model_name: ensemble_name,
            direction,
            confidence: avg_confidence,
            metadata: PredictionMeta::default(),
        })
    }

    /// Aggregate predictions with logits: apply sigmoid to each logit,
    /// compute confidence-weighted sum, remap to [-1,1] direction space.
    fn aggregate_logits(
        &self,
        preds: &[(String, RawPrediction)],
    ) -> Option<(f64, usize)> {
        if preds.is_empty() {
            return None;
        }

        let mut weighted_direction_sum = 0.0_f64;
        let mut weight_confidence_sum = 0.0_f64;

        for (name, pred) in preds {
            let logits = match pred.logits.as_ref() {
                Some(l) if !l.is_empty() => l,
                _ => continue,
            };

            // Apply sigmoid to logits and compute mean
            let sigmoid_sum: f64 = logits
                .iter()
                .map(|&x| 1.0 / (1.0 + (-f64::from(x)).exp()))
                .sum();
            let sigmoid_mean = sigmoid_sum / logits.len() as f64;

            // Remap [0,1] sigmoid output to [-1,1] direction space
            let direction = sigmoid_mean * 2.0 - 1.0;

            let confidence = pred.confidence.clamp(0.0, 1.0);
            let w = self.weights.get(name).copied().unwrap_or(1.0);
            let wc = w * confidence;
            weighted_direction_sum += direction * wc;
            weight_confidence_sum += wc;
        }

        let direction = if weight_confidence_sum.abs() < f64::EPSILON {
            0.0
        } else {
            weighted_direction_sum / weight_confidence_sum
        };

        Some((direction, preds.len()))
    }

    /// Batched inference across all ready adapters: processes N feature vectors
    /// through each model with a single GPU upload -> forward -> download per model,
    /// then aggregates predictions per sample via confidence-weighted voting.
    ///
    /// Returns one [`EnsemblePrediction`] per input feature vector.
    pub fn predict_batch(&self, batch: &[FeatureVector]) -> MLResult<Vec<EnsemblePrediction>> {
        if batch.is_empty() {
            return Ok(Vec::new());
        }

        let ready_adapters: Vec<&Box<dyn ModelInferenceAdapter>> =
            self.adapters.iter().filter(|a| a.is_ready()).collect();

        if ready_adapters.is_empty() {
            return Err(MLError::InferenceError(
                "No models are ready for inference".to_owned(),
            ));
        }

        // Run predict_batch() in parallel across all ready adapters.
        // Each adapter does a single batched GPU roundtrip internally.
        let all_results: Vec<(String, Vec<EnsemblePrediction>)> = ready_adapters
            .par_iter()
            .filter_map(|adapter| {
                let name = adapter.model_name().to_string();
                match adapter.predict_batch(batch) {
                    Ok(preds) => Some((name, preds)),
                    Err(e) => {
                        tracing::warn!(
                            model = %name,
                            error = %e,
                            "Batched prediction failed, skipping model"
                        );
                        None
                    }
                }
            })
            .collect();

        if all_results.is_empty() {
            return Err(MLError::InferenceError(
                "All ready models failed during batched inference".to_owned(),
            ));
        }

        // Aggregate per-sample: for each index, collect predictions from all models
        let n = batch.len();
        let mut ensemble_preds = Vec::with_capacity(n);

        for i in 0..n {
            let mut weighted_dir_sum = 0.0_f64;
            let mut weight_conf_sum = 0.0_f64;
            let mut conf_sum = 0.0_f64;
            let mut model_names = Vec::new();
            let mut count: usize = 0;

            for (name, preds) in &all_results {
                if let Some(pred) = preds.get(i) {
                    if !pred.direction.is_finite() || !pred.confidence.is_finite() {
                        continue;
                    }
                    let conf = pred.confidence.clamp(0.0, 1.0);
                    let w = self.weights.get(name).copied().unwrap_or(1.0);
                    let wc = w * conf;
                    weighted_dir_sum += pred.direction * wc;
                    weight_conf_sum += wc;
                    conf_sum += conf;
                    model_names.push(name.as_str());
                    count += 1;
                }
            }

            let direction = if weight_conf_sum.abs() < f64::EPSILON {
                0.0
            } else {
                weighted_dir_sum / weight_conf_sum
            };

            let avg_confidence = if count > 0 {
                conf_sum / count as f64
            } else {
                0.0
            };

            ensemble_preds.push(EnsemblePrediction {
                model_name: format!("ENSEMBLE({})", model_names.join("+")),
                direction,
                confidence: avg_confidence,
                metadata: PredictionMeta::default(),
            });
        }

        Ok(ensemble_preds)
    }

    /// CPU-side scalar weighted average (original algorithm).
    fn aggregate_cpu(
        &self,
        preds: &[(String, RawPrediction)],
    ) -> Option<(f64, usize)> {
        if preds.is_empty() {
            return None;
        }

        let mut weighted_direction_sum = 0.0_f64;
        let mut weight_confidence_sum = 0.0_f64;

        for (name, pred) in preds {
            let confidence = pred.confidence.clamp(0.0, 1.0);
            let w = self.weights.get(name).copied().unwrap_or(1.0);
            let wc = w * confidence;
            weighted_direction_sum += pred.direction_scalar * wc;
            weight_confidence_sum += wc;
        }

        let direction = if weight_confidence_sum.abs() < f64::EPSILON {
            0.0
        } else {
            weighted_direction_sum / weight_confidence_sum
        };

        Some((direction, preds.len()))
    }
}

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

    /// A simple test adapter with configurable direction, confidence, and readiness.
    struct DummyAdapter {
        name: String,
        direction: f64,
        confidence: f64,
        ready: bool,
    }

    impl ModelInferenceAdapter for DummyAdapter {
        fn model_name(&self) -> &str {
            &self.name
        }

        fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
            Ok(EnsemblePrediction {
                model_name: self.name.clone(),
                direction: self.direction,
                confidence: self.confidence,
                metadata: PredictionMeta::default(),
            })
        }

        fn is_ready(&self) -> bool {
            self.ready
        }
    }

    fn make_features() -> FeatureVector {
        FeatureVector {
            values: vec![0.0; 51],
            timestamp: 1_700_000_000_000_000,
        }
    }

    #[test]
    fn test_ensemble_equal_weight_aggregation() {
        // Model A: bullish (dir=1.0, conf=0.8)
        // Model B: bearish (dir=-1.0, conf=0.6)
        // weighted_direction = (1.0*1.0*0.8 + (-1.0)*1.0*0.6) / (1.0*0.8 + 1.0*0.6)
        //                    = (0.8 - 0.6) / 1.4
        //                    = 0.2 / 1.4
        //                    ~ 0.1429
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "A".to_owned(),
                direction: 1.0,
                confidence: 0.8,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "B".to_owned(),
                direction: -1.0,
                confidence: 0.6,
                ready: true,
            }),
        ];

        let ensemble = InferenceEnsemble::new(adapters);
        let features = make_features();
        let pred = ensemble.predict(&features).expect("predict should succeed");

        // Net direction should be positive (bullish model has higher confidence)
        assert!(
            pred.direction > 0.0,
            "direction should be positive, got {}",
            pred.direction
        );
        // But not too strong since the bearish model partially cancels
        assert!(
            pred.direction < 0.5,
            "direction should be < 0.5, got {}",
            pred.direction
        );
        // Confidence = average = (0.8 + 0.6) / 2 = 0.7
        assert!(
            (pred.confidence - 0.7).abs() < 1e-9,
            "confidence should be 0.7, got {}",
            pred.confidence
        );
    }

    #[test]
    fn test_ensemble_skips_unready_models() {
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "Ready".to_owned(),
                direction: 1.0,
                confidence: 0.9,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "NotReady".to_owned(),
                direction: -1.0,
                confidence: 0.9,
                ready: false,
            }),
        ];

        let ensemble = InferenceEnsemble::new(adapters);
        assert_eq!(ensemble.ready_count(), 1);

        let features = make_features();
        let pred = ensemble.predict(&features).expect("predict should succeed");

        // Only the ready (bullish) model contributes
        assert!(
            pred.direction > 0.5,
            "direction should be > 0.5 with only bullish model, got {}",
            pred.direction
        );
        // Model name should only include the ready model
        assert!(
            pred.model_name.contains("Ready"),
            "model_name should contain 'Ready', got {}",
            pred.model_name
        );
        assert!(
            !pred.model_name.contains("NotReady"),
            "model_name should NOT contain 'NotReady', got {}",
            pred.model_name
        );
    }

    #[test]
    fn test_ensemble_custom_weights() {
        // Model A: bullish, weight 0.75
        // Model B: bearish, weight 0.25
        // weighted_direction = (1.0*0.75*0.8 + (-1.0)*0.25*0.8) / (0.75*0.8 + 0.25*0.8)
        //                    = (0.6 - 0.2) / (0.6 + 0.2)
        //                    = 0.4 / 0.8
        //                    = 0.5
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "Heavy".to_owned(),
                direction: 1.0,
                confidence: 0.8,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "Light".to_owned(),
                direction: -1.0,
                confidence: 0.8,
                ready: true,
            }),
        ];

        let mut ensemble = InferenceEnsemble::new(adapters);
        ensemble.set_weight("Heavy", 0.75);
        ensemble.set_weight("Light", 0.25);

        let features = make_features();
        let pred = ensemble.predict(&features).expect("predict should succeed");

        // The heavier-weighted bullish model should dominate
        assert!(
            pred.direction > 0.0,
            "direction should be positive (heavy bullish dominates), got {}",
            pred.direction
        );
        // With equal confidences, the 3:1 weight ratio strongly favors bullish
        // Expected: 0.5, but any positive value > 0.3 confirms dominance
        assert!(
            pred.direction > 0.3,
            "direction should be > 0.3 showing heavy model dominance, got {}",
            pred.direction
        );
    }

    #[test]
    fn test_ensemble_filters_nan_predictions() {
        // NaN model should be skipped; the valid model's prediction stands alone.
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "Valid".to_owned(),
                direction: 0.8,
                confidence: 0.7,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "NaN_Model".to_owned(),
                direction: f64::NAN,
                confidence: 0.9,
                ready: true,
            }),
        ];

        let ensemble = InferenceEnsemble::new(adapters);
        let features = make_features();
        let pred = ensemble.predict(&features).expect("predict should succeed");

        // Only the Valid model should contribute
        assert!(
            pred.model_name.contains("Valid"),
            "model_name should contain 'Valid', got {}",
            pred.model_name
        );
        assert!(
            !pred.model_name.contains("NaN_Model"),
            "model_name should NOT contain 'NaN_Model', got {}",
            pred.model_name
        );
        // direction should come entirely from the Valid model
        assert!(
            (pred.direction - 0.8).abs() < 1e-9,
            "direction should be 0.8, got {}",
            pred.direction
        );
        // confidence should be from the single valid model
        assert!(
            (pred.confidence - 0.7).abs() < 1e-9,
            "confidence should be 0.7, got {}",
            pred.confidence
        );
    }

    #[test]
    fn test_predict_batch_aggregates_correctly() {
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "Bull".to_owned(),
                direction: 1.0,
                confidence: 0.8,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "Bear".to_owned(),
                direction: -1.0,
                confidence: 0.6,
                ready: true,
            }),
        ];

        let ensemble = InferenceEnsemble::new(adapters);
        let batch = vec![make_features(), make_features(), make_features()];
        let preds = ensemble.predict_batch(&batch).expect("batch should succeed");

        assert_eq!(preds.len(), 3, "should produce one prediction per input");
        for pred in &preds {
            // Same adapters, same features -> same result for each sample
            assert!(pred.direction > 0.0, "net direction should be bullish");
            assert!(
                pred.model_name.contains("Bull"),
                "should include Bull model"
            );
            assert!(
                pred.model_name.contains("Bear"),
                "should include Bear model"
            );
        }
    }

    #[test]
    fn test_predict_batch_empty_input() {
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![Box::new(DummyAdapter {
            name: "A".to_owned(),
            direction: 0.5,
            confidence: 0.9,
            ready: true,
        })];
        let ensemble = InferenceEnsemble::new(adapters);
        let preds = ensemble.predict_batch(&[]).expect("empty batch should succeed");
        assert!(preds.is_empty(), "empty input should produce empty output");
    }

    #[test]
    fn test_predict_batch_skips_nan() {
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "Valid".to_owned(),
                direction: 0.5,
                confidence: 0.8,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "NaN".to_owned(),
                direction: f64::NAN,
                confidence: 0.9,
                ready: true,
            }),
        ];

        let ensemble = InferenceEnsemble::new(adapters);
        let batch = vec![make_features()];
        let preds = ensemble.predict_batch(&batch).expect("should succeed");

        assert_eq!(preds.len(), 1);
        assert!(
            (preds[0].direction - 0.5).abs() < 1e-9,
            "should use only valid model, got {}",
            preds[0].direction
        );
    }

    #[test]
    fn test_ensemble_filters_extreme_confidence() {
        // Model with confidence 5.0 should be clamped to 1.0
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "Normal".to_owned(),
                direction: 1.0,
                confidence: 0.6,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "Extreme".to_owned(),
                direction: 1.0,
                confidence: 5.0,
                ready: true,
            }),
        ];

        let ensemble = InferenceEnsemble::new(adapters);
        let features = make_features();
        let pred = ensemble.predict(&features).expect("predict should succeed");

        // Both models should contribute (extreme confidence is clamped, not skipped)
        assert!(
            pred.model_name.contains("Normal"),
            "model_name should contain 'Normal', got {}",
            pred.model_name
        );
        assert!(
            pred.model_name.contains("Extreme"),
            "model_name should contain 'Extreme', got {}",
            pred.model_name
        );
        // Confidence should be clamped: avg of 0.6 and 1.0 = 0.8
        assert!(
            (pred.confidence - 0.8).abs() < 1e-9,
            "confidence should be 0.8 (avg of 0.6 and clamped 1.0), got {}",
            pred.confidence
        );
    }
}
