//! Stream-aware ensemble with per-model CUDA device handles.
//!
//! Each model runs on its own CUDA stream for true GPU-level parallelism.
//! Falls back to standard rayon parallelism on CPU.
//!
//! # Design (Option C — per-model CudaDevice handles)
//!
//! On CUDA: each model adapter gets its own `CudaDevice::new(0)` handle so the
//! CUDA driver assigns a unique default stream. `predict()` runs all models via
//! rayon with per-device streams, syncs all, then aggregates on the primary device.
//!
//! On CPU: all adapters share the same device, rayon provides parallelism (same
//! behaviour as [`InferenceEnsemble`](super::inference_ensemble::InferenceEnsemble)).

use rayon::prelude::*;

use ml_core::device::MlDevice;
#[cfg(feature = "cuda")]
use ml_core::cuda_autograd::{ActivationKernels, GpuTensor};

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

/// Ensemble that runs models on separate CUDA streams for true
/// GPU-level parallelism. Falls back to rayon on CPU.
///
/// Uses [`CudaStreamPool`] for stream management and synchronisation.
#[allow(missing_debug_implementations)]
pub struct StreamAwareEnsemble {
    /// Model adapters (each may have its own device handle)
    adapters: Vec<Box<dyn ModelInferenceAdapter>>,
    /// Weights per model (normalised, indexed by adapter position)
    weights: Vec<f64>,
    /// Stream pool for CUDA synchronisation (no-op on CPU)
    stream_pool: CudaStreamPool,
}

impl StreamAwareEnsemble {
    /// Create a new stream-aware ensemble.
    ///
    /// `weights` are normalised so they sum to 1.0. If the provided weights
    /// slice is shorter than `adapters`, missing entries default to equal
    /// share. If `weights` is empty, all models get equal weight.
    pub fn new(
        adapters: Vec<Box<dyn ModelInferenceAdapter>>,
        weights: Vec<f64>,
        device: &MlDevice,
    ) -> MLResult<Self> {
        let n = adapters.len();
        let stream_pool = CudaStreamPool::new(device, n)?;

        let raw_weights = if weights.is_empty() || n == 0 {
            vec![1.0_f64; n]
        } else {
            let mut w = weights;
            w.resize(n, 1.0);
            w
        };

        let weight_sum: f64 = raw_weights.iter().sum();
        let normalized = if weight_sum > 0.0 {
            raw_weights.iter().map(|w| w / weight_sum).collect()
        } else {
            vec![1.0 / n.max(1) as f64; n]
        };

        Ok(Self {
            adapters,
            weights: normalized,
            stream_pool,
        })
    }

    /// Run ensemble prediction with stream-level parallelism.
    ///
    /// 1. Runs all ready adapters in parallel via rayon (each on its own
    ///    CUDA stream if GPU, or rayon thread if CPU).
    /// 2. Synchronises all CUDA streams (no-op on CPU).
    /// 3. Aggregates predictions via confidence-weighted voting with
    ///    GPU-logit and CPU-scalar paths (same algorithm as
    ///    [`InferenceEnsemble`]).
    pub fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
        // Parallel inference via rayon — each adapter on its own stream if CUDA
        let predictions: Vec<(usize, String, RawPrediction)> = self
            .adapters
            .par_iter()
            .enumerate()
            .filter(|(_, a)| a.is_ready())
            .filter_map(|(i, adapter)| {
                let 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 = %name,
                                direction = %pred.direction_scalar,
                                confidence = %pred.confidence,
                                "StreamAwareEnsemble: model returned NaN/Inf, skipping"
                            );
                            None
                        } else {
                            Some((i, name, pred))
                        }
                    }
                    Err(e) => {
                        tracing::warn!(
                            model = %name,
                            error = %e,
                            "StreamAwareEnsemble: model prediction failed, skipping"
                        );
                        None
                    }
                }
            })
            .collect();

        if predictions.is_empty() {
            return Ok(EnsemblePrediction {
                model_name: "STREAM_ENSEMBLE(empty)".to_owned(),
                direction: 0.0,
                confidence: 0.0,
                metadata: PredictionMeta::default(),
            });
        }

        // Synchronise all CUDA streams before aggregation (no-op on CPU)
        self.stream_pool.sync_all()?;

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

        let mut model_names: Vec<String> = Vec::new();
        let mut total_confidence_sum = 0.0_f64;
        let mut total_count: usize = 0;

        // --- Logits path: sigmoid + weighted-sum on host f32 vectors ---
        let logit_result = if !logit_preds.is_empty() {
            match self.aggregate_logits(&logit_preds) {
                Ok((direction, count)) => {
                    for (_, name, pred) in &logit_preds {
                        model_names.push(name.clone());
                        total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
                    }
                    total_count += logit_preds.len();
                    Some((direction, count))
                }
                Err(e) => {
                    tracing::warn!(
                        error = %e,
                        "Logits aggregation failed, falling back to CPU for {} models",
                        logit_preds.len()
                    );
                    let fallback = self.aggregate_cpu(&logit_preds);
                    for (_, name, pred) in &logit_preds {
                        model_names.push(name.clone());
                        total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
                    }
                    total_count += logit_preds.len();
                    fallback
                }
            }
        } else {
            None
        };

        // --- CPU path: 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 (logit_result, cpu_result) {
            (Some((logit_dir, logit_n)), Some((cpu_dir, cpu_n))) => {
                let total_n = (logit_n + cpu_n) as f64;
                (logit_dir * logit_n as f64 + cpu_dir * cpu_n as f64) / total_n
            }
            (Some((dir, _)), None) | (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!("STREAM_ENSEMBLE({})", model_names.join("+"));

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

    /// Number of adapters in the ensemble.
    pub fn adapter_count(&self) -> usize {
        self.adapters.len()
    }

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

    /// Whether the underlying stream pool is on a CUDA device.
    pub fn is_cuda(&self) -> bool {
        self.stream_pool.is_cuda()
    }

    // -----------------------------------------------------------------------
    // Private aggregation helpers
    // -----------------------------------------------------------------------

    /// Logits aggregation on GPU: concatenate all models' logits into a single
    /// GPU tensor, apply sigmoid via `ActivationKernels` CUDA kernel, compute
    /// per-model mean and weighted sum via a fused CUDA kernel, then read back
    /// only the final scalar.
    ///
    /// Falls back to a CPU path when CUDA is not available.
    fn aggregate_logits(
        &self,
        preds: &[(usize, String, RawPrediction)],
    ) -> Result<(f64, usize), MLError> {
        #[cfg(feature = "cuda")]
        {
            if let Some(result) = self.aggregate_logits_gpu(preds)? {
                return Ok(result);
            }
        }

        // CPU fallback (non-CUDA builds or when stream pool has no CUDA streams)
        self.aggregate_logits_cpu(preds)
    }

    /// GPU path: sigmoid + weighted-mean-reduce entirely on device.
    ///
    /// 1. Concatenate all models' logit vectors into one flat `GpuTensor`.
    /// 2. Apply sigmoid via `ActivationKernels::sigmoid_fwd()` (one kernel launch).
    /// 3. Launch a fused reduction kernel that computes per-model means and the
    ///    final confidence-weighted sum in a single pass.
    /// 4. Read back only the single f32 result.
    ///
    /// Returns `Ok(None)` if no CUDA stream is available (caller should fall back
    /// to CPU).
    #[cfg(feature = "cuda")]
    #[allow(unsafe_code)] // CUDA kernel launches require unsafe FFI.
    fn aggregate_logits_gpu(
        &self,
        preds: &[(usize, String, RawPrediction)],
    ) -> Result<Option<(f64, usize)>, MLError> {
        use std::sync::Arc;
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        // Need a CUDA stream from the pool
        let stream = match self.stream_pool.get_stream(0) {
            Some(s) => Arc::clone(s),
            None => return Ok(None),
        };

        // Build per-model metadata: (weight*confidence, logit_offset, logit_len)
        let mut all_logits: Vec<f32> = Vec::new();
        let mut segment_offsets: Vec<i32> = Vec::new();
        let mut segment_lengths: Vec<i32> = Vec::new();
        let mut model_weights: Vec<f32> = Vec::new();

        let mut weight_sum = 0.0_f32;

        for (i, (idx, _, pred)) in preds.iter().enumerate() {
            let logits = match pred.logits.as_ref() {
                Some(l) if !l.is_empty() => l,
                _ => continue,
            };
            let conf = pred.confidence.clamp(0.0, 1.0) as f32;
            let w = self.weights.get(*idx).copied().unwrap_or(1.0) as f32;
            let wc = w * conf;

            segment_offsets.push(all_logits.len() as i32);
            segment_lengths.push(logits.len() as i32);
            model_weights.push(wc);
            weight_sum += wc;

            all_logits.extend_from_slice(logits);
            let _ = i; // suppress unused
        }

        if all_logits.is_empty() || weight_sum.abs() < f32::EPSILON {
            return Ok(Some((0.0, preds.len())));
        }

        let n_models = segment_offsets.len();

        // Normalize weights
        for w in &mut model_weights {
            *w /= weight_sum;
        }

        // 1) Upload concatenated logits to GPU
        let logit_tensor = GpuTensor::from_host(
            &all_logits,
            vec![all_logits.len()],
            &stream,
        )?;

        // 2) Apply sigmoid on GPU via ActivationKernels
        let act_kernels = ActivationKernels::new(&stream)?;
        let (sigmoid_tensor, _saved) = act_kernels.sigmoid_fwd(&logit_tensor, &stream)?;

        // 3) Upload segment metadata and weights to GPU, then launch a
        //    fused per-model-mean + weighted-sum reduction kernel.
        let mut d_offsets = stream.alloc_zeros::<i32>(n_models).map_err(|e| {
            MLError::ModelError(format!("alloc offsets: {e}"))
        })?;
        stream.memcpy_htod(&segment_offsets, &mut d_offsets).map_err(|e| {
            MLError::ModelError(format!("htod offsets: {e}"))
        })?;

        let mut d_lengths = stream.alloc_zeros::<i32>(n_models).map_err(|e| {
            MLError::ModelError(format!("alloc lengths: {e}"))
        })?;
        stream.memcpy_htod(&segment_lengths, &mut d_lengths).map_err(|e| {
            MLError::ModelError(format!("htod lengths: {e}"))
        })?;

        let mut d_weights = stream.alloc_zeros::<f32>(n_models).map_err(|e| {
            MLError::ModelError(format!("alloc weights: {e}"))
        })?;
        stream.memcpy_htod(&model_weights, &mut d_weights).map_err(|e| {
            MLError::ModelError(format!("htod weights: {e}"))
        })?;

        // Output: single f32 (weighted sum of per-model sigmoid means)
        let d_output = stream.alloc_zeros::<f32>(1).map_err(|e| {
            MLError::ModelError(format!("alloc output: {e}"))
        })?;

        let context = stream.context();
        let ptx = ml_core::cuda_compile::compile_ptx_for_device(
            ENSEMBLE_REDUCE_CUDA_SRC,
            context,
        ).map_err(|e| MLError::ModelError(format!("ensemble reduce compile: {e}")))?;

        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("ensemble reduce module load: {e}"))
        })?;
        let kernel = module.load_function("weighted_sigmoid_mean_reduce").map_err(|e| {
            MLError::ModelError(format!("ensemble reduce kernel load: {e}"))
        })?;

        let n_models_i32 = n_models as i32;
        let cfg = LaunchConfig {
            grid_dim: (1, 1, 1),
            block_dim: (n_models.min(256) as u32, 1, 1),
            shared_mem_bytes: 0,
        };

        // SAFETY: kernel arguments match the CUDA kernel signature
        // (sigmoid_vals, offsets, lengths, weights, output, n_models).
        // All buffers are GPU-allocated with correct sizes above.
        unsafe {
            stream
                .launch_builder(&kernel)
                .arg(sigmoid_tensor.data())
                .arg(&d_offsets)
                .arg(&d_lengths)
                .arg(&d_weights)
                .arg(&d_output)
                .arg(&n_models_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("ensemble reduce launch: {e}")))?;
        }

        // 4) Read back the single scalar result
        let mut result_host = [0.0_f32];
        stream.memcpy_dtoh(&d_output, &mut result_host).map_err(|e| { // gpu-exit: final prediction scalar exits system
            MLError::ModelError(format!("ensemble reduce dtoh: {e}"))
        })?;

        // Sigmoid output is [0,1]; remap to [-1,1] direction space
        let direction = (f64::from(result_host[0]) * 2.0) - 1.0;

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

    /// CPU fallback for logits aggregation (non-CUDA builds).
    fn aggregate_logits_cpu(
        &self,
        preds: &[(usize, String, RawPrediction)],
    ) -> Result<(f64, usize), MLError> {
        let weights_f32: Vec<f32> = preds
            .iter()
            .map(|(idx, _, p)| {
                let conf = p.confidence.clamp(0.0, 1.0) as f32;
                let w = self.weights.get(*idx).copied().unwrap_or(1.0) as f32;
                w * conf
            })
            .collect();

        let weight_sum: f32 = weights_f32.iter().sum();
        if weight_sum.abs() < f32::EPSILON {
            return Ok((0.0, preds.len()));
        }

        let mut weighted_sum = 0.0_f64;

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

            // CPU sigmoid + mean (fallback only)
            let sum_sigmoid: f32 = logits
                .iter()
                .map(|&x| 1.0_f32 / (1.0_f32 + (-x).exp()))
                .sum();
            let mean_sigmoid = sum_sigmoid / logits.len() as f32;

            let normalized_w = weights_f32.get(i).copied().unwrap_or(0.0) / weight_sum;
            weighted_sum += f64::from(mean_sigmoid) * f64::from(normalized_w);
        }

        let direction = (weighted_sum * 2.0) - 1.0;
        Ok((direction, preds.len()))
    }

    /// CPU-side scalar weighted average.
    fn aggregate_cpu(
        &self,
        preds: &[(usize, 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 (idx, _, pred) in preds {
            let confidence = pred.confidence.clamp(0.0, 1.0);
            let w = self.weights.get(*idx).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()))
    }
}

// ── CUDA source for ensemble weighted-sigmoid-mean reduction ─────────────
//
// One thread per model. Each thread loops over its logit segment (already
// post-sigmoid), computes the mean, multiplies by the normalized model weight,
// and atomically adds to a single output scalar.
//
// This keeps the entire aggregation on GPU after sigmoid_fwd() — only one f32
// is read back to CPU.
#[cfg(feature = "cuda")]
const ENSEMBLE_REDUCE_CUDA_SRC: &str = r#"
extern "C" __global__
void weighted_sigmoid_mean_reduce(
    const float* __restrict__ sigmoid_vals,
    const int*   __restrict__ offsets,
    const int*   __restrict__ lengths,
    const float* __restrict__ weights,
    float*       __restrict__ output,
    int n_models)
{
    int model = blockIdx.x * blockDim.x + threadIdx.x;
    if (model >= n_models) return;

    int off = offsets[model];
    int len = lengths[model];
    if (len <= 0) return;

    // Compute mean of post-sigmoid values for this model's segment
    float sum = 0.0f;
    for (int j = 0; j < len; ++j) {
        sum += sigmoid_vals[off + j];
    }
    float mean_val = sum / (float)len;

    // Weighted contribution
    float contribution = mean_val * weights[model];
    atomicAdd(output, contribution);
}
"#;

#[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
        }
    }

    /// An adapter that always fails prediction.
    struct FailingAdapter {
        name: String,
    }

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

        fn predict(&self, _features: &FeatureVector) -> MLResult<EnsemblePrediction> {
            Err(MLError::InferenceError("intentional test failure".to_owned()))
        }

        fn is_ready(&self) -> bool {
            true
        }
    }

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

    fn test_device() -> MlDevice {
        // Use CPU for tests — no CUDA dependency in unit tests
        MlDevice::Cpu
    }

    #[test]
    fn test_stream_ensemble_empty() {
        let ensemble = StreamAwareEnsemble::new(vec![], vec![], &test_device())
            .expect("empty ensemble should create");
        assert_eq!(ensemble.adapter_count(), 0);
        assert_eq!(ensemble.ready_count(), 0);

        let pred = ensemble
            .predict(&make_features())
            .expect("empty predict should succeed");
        assert_eq!(pred.direction, 0.0);
        assert_eq!(pred.confidence, 0.0);
        assert_eq!(pred.model_name, "STREAM_ENSEMBLE(empty)");
    }

    #[test]
    fn test_stream_ensemble_single_model() {
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![Box::new(DummyAdapter {
            name: "DQN".to_owned(),
            direction: 0.7,
            confidence: 0.9,
            ready: true,
        })];

        let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0], &test_device())
            .expect("single model should create");
        assert_eq!(ensemble.adapter_count(), 1);
        assert_eq!(ensemble.ready_count(), 1);

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

        // Single model: direction comes straight through the CPU path
        assert!(
            (pred.direction - 0.7).abs() < 1e-9,
            "expected direction 0.7, got {}",
            pred.direction
        );
        assert!(
            (pred.confidence - 0.9).abs() < 1e-9,
            "expected confidence 0.9, got {}",
            pred.confidence
        );
        assert!(pred.model_name.contains("DQN"));
    }

    #[test]
    fn test_stream_ensemble_two_models_equal_weight() {
        // Model A: bullish (dir=1.0, conf=0.8)
        // Model B: bearish (dir=-1.0, conf=0.6)
        // CPU path: weighted_direction = (1.0*0.5*0.8 + (-1.0)*0.5*0.6) / (0.5*0.8 + 0.5*0.6)
        //         = (0.4 - 0.3) / (0.4 + 0.3) = 0.1 / 0.7 ~ 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 = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
            .expect("two-model ensemble should create");

        let pred = ensemble
            .predict(&make_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
        );
        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,
            "expected confidence 0.7, got {}",
            pred.confidence
        );
        assert!(pred.model_name.contains('A'));
        assert!(pred.model_name.contains('B'));
    }

    #[test]
    fn test_stream_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 = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
            .expect("ensemble should create");
        assert_eq!(ensemble.ready_count(), 1);

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

        assert!(pred.model_name.contains("Ready"));
        assert!(!pred.model_name.contains("NotReady"));
        assert!(pred.direction > 0.5);
    }

    #[test]
    fn test_stream_ensemble_skips_nan_predictions() {
        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 = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
            .expect("ensemble should create");

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

        assert!(pred.model_name.contains("Valid"));
        assert!(!pred.model_name.contains("NaN_Model"));
        assert!(
            (pred.direction - 0.8).abs() < 1e-9,
            "expected direction 0.8, got {}",
            pred.direction
        );
    }

    #[test]
    fn test_stream_ensemble_skips_failing_models() {
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "Good".to_owned(),
                direction: 0.5,
                confidence: 0.8,
                ready: true,
            }),
            Box::new(FailingAdapter {
                name: "Broken".to_owned(),
            }),
        ];

        let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
            .expect("ensemble should create");

        let pred = ensemble
            .predict(&make_features())
            .expect("predict should succeed despite one failing model");

        assert!(pred.model_name.contains("Good"));
        assert!(!pred.model_name.contains("Broken"));
    }

    #[test]
    fn test_stream_ensemble_weight_normalisation() {
        // Weights [3.0, 1.0] should normalise to [0.75, 0.25]
        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 ensemble = StreamAwareEnsemble::new(adapters, vec![3.0, 1.0], &test_device())
            .expect("ensemble should create");

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

        // Heavy (0.75) bullish should dominate over Light (0.25) bearish
        assert!(
            pred.direction > 0.3,
            "heavy bullish model should dominate, got {}",
            pred.direction
        );
    }

    #[test]
    fn test_stream_ensemble_all_models_fail_returns_neutral() {
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(FailingAdapter {
                name: "Fail1".to_owned(),
            }),
            Box::new(FailingAdapter {
                name: "Fail2".to_owned(),
            }),
        ];

        let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &test_device())
            .expect("ensemble should create");

        let pred = ensemble
            .predict(&make_features())
            .expect("predict should return neutral when all fail");

        assert_eq!(pred.direction, 0.0);
        assert_eq!(pred.confidence, 0.0);
        assert_eq!(pred.model_name, "STREAM_ENSEMBLE(empty)");
    }

    #[test]
    fn test_stream_ensemble_missing_weights_default() {
        // Provide fewer weights than adapters — missing ones should default
        let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
            Box::new(DummyAdapter {
                name: "A".to_owned(),
                direction: 0.5,
                confidence: 0.8,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "B".to_owned(),
                direction: 0.5,
                confidence: 0.8,
                ready: true,
            }),
            Box::new(DummyAdapter {
                name: "C".to_owned(),
                direction: 0.5,
                confidence: 0.8,
                ready: true,
            }),
        ];

        let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0], &test_device())
            .expect("missing weights should use defaults");

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

        // All models agree on 0.5, so direction should be ~0.5
        assert!(
            (pred.direction - 0.5).abs() < 1e-9,
            "expected direction 0.5, got {}",
            pred.direction
        );
    }
}
