//! Mamba2 SSM inference adapter for ensemble prediction
//!
//! Wraps a loaded Mamba2 state-space model with a sequence buffer.
//! Feature vectors are zero-padded to `d_model` dimensions and accumulated
//! in a ring buffer.  Once `sequence_length` frames are collected the SSM
//! runs a forward pass and the last timestep's sigmoid output is mapped to
//! a directional signal + confidence for ensemble aggregation.

use std::collections::VecDeque;
use std::sync::{Arc, Mutex};

use cudarc::driver::CudaStream;

use ml_core::cuda_autograd::stream_ops::StreamTensor;

use crate::ensemble::inference_adapter::{
    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
};
use crate::mamba::{Mamba2Config, Mamba2SSM};
use crate::{MLError, MLResult};

/// Inference adapter that wraps a Mamba2 SSM for ensemble prediction.
///
/// Maintains an internal sequence buffer of zero-padded feature vectors.
/// When the buffer reaches `sequence_length`, each call to [`predict`]
/// builds a `[1, seq_len, d_model]` GpuTensor, runs the SSM forward pass,
/// and converts the sigmoid output of the last timestep into:
///
/// - **direction**: `2 * prob - 1` mapping `[0,1] -> [-1,1]`
/// - **confidence**: `|prob - 0.5| * 2` (highest at extremes, lowest at 0.5)
#[allow(missing_debug_implementations)]
pub struct Mamba2InferenceAdapter {
    model: Mutex<Mamba2SSM>,
    buffer: Mutex<VecDeque<Vec<f32>>>,
    stream: Arc<CudaStream>,
    sequence_length: usize,
    d_model: usize,
}

// SAFETY: Mamba2SSM internally uses GpuTensor/CudaSlice which are Send+Sync.
// The Mutex provides exclusive access for inference calls and buffer mutation.
#[allow(unsafe_code)]
unsafe impl Send for Mamba2InferenceAdapter {}
#[allow(unsafe_code)]
unsafe impl Sync for Mamba2InferenceAdapter {}

impl Mamba2InferenceAdapter {
    /// Create a new Mamba2 inference adapter from configuration.
    pub fn new(config: Mamba2Config, sequence_length: usize) -> MLResult<Self> {
        let ctx = cudarc::driver::CudaContext::new(0)
            .map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?;
        let stream = ctx
            .new_stream()
            .map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;

        let d_model = config.d_model;
        let model = Mamba2SSM::new(config, &stream)?;
        Ok(Self {
            model: Mutex::new(model),
            buffer: Mutex::new(VecDeque::with_capacity(sequence_length)),
            stream,
            sequence_length,
            d_model,
        })
    }

    /// Zero-pad a feature vector to `d_model` dimensions (f64 -> f32).
    fn pad_to_d_model(&self, values: &[f64]) -> Vec<f32> {
        let mut padded = vec![0.0_f32; self.d_model];
        let copy_len = values.len().min(self.d_model);
        for i in 0..copy_len {
            if let Some(v) = values.get(i) {
                if let Some(slot) = padded.get_mut(i) {
                    *slot = *v as f32;
                }
            }
        }
        padded
    }
}

impl ModelInferenceAdapter for Mamba2InferenceAdapter {
    fn model_name(&self) -> &str {
        "MAMBA-2"
    }

    fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
        let start = std::time::Instant::now();

        let padded = self.pad_to_d_model(&features.values);

        // Push to buffer
        let buffer_ready = {
            let mut buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("Mamba2 buffer lock poisoned: {e}")))?;
            buf.push_back(padded);
            while buf.len() > self.sequence_length {
                buf.pop_front();
            }
            buf.len() >= self.sequence_length
        };

        if !buffer_ready {
            let latency_us = start.elapsed().as_micros() as u64;
            return Ok(EnsemblePrediction {
                model_name: "MAMBA-2".to_owned(),
                direction: 0.0,
                confidence: 0.0,
                metadata: PredictionMeta {
                    latency_us,
                    quantiles: None,
                    attention_weights: None,
                    q_values: None,
                },
            });
        }

        // Build input GpuTensor from buffer snapshot
        let flat_data: Vec<f32> = {
            let buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("Mamba2 buffer lock poisoned: {e}")))?;
            let mut data = Vec::with_capacity(self.sequence_length * self.d_model);
            for frame in buf.iter() {
                data.extend_from_slice(frame);
            }
            data
        };

        // [1, seq_len, d_model]
        let input = StreamTensor::from_vec(
            flat_data,
            &[1, self.sequence_length, self.d_model],
            &self.stream,
        )?;

        // Run forward pass
        let mut model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("Mamba2 model lock poisoned: {e}")))?;
        let output = model.forward(&input)?;

        // Output shape: [1, seq_len, 1] or similar -- extract last timestep
        let output_host = output.to_vec()?; // gpu-exit: scalar prediction
        let prob = output_host.last().copied().unwrap_or(0.5) as f64;

        // Map sigmoid output [0, 1] to direction [-1, 1] and confidence [0, 1]
        let direction = (2.0 * prob - 1.0).clamp(-1.0, 1.0);
        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);

        let latency_us = start.elapsed().as_micros() as u64;

        Ok(EnsemblePrediction {
            model_name: "MAMBA-2".to_owned(),
            direction,
            confidence,
            metadata: PredictionMeta {
                latency_us,
                quantiles: None,
                attention_weights: None,
                q_values: None,
            },
        })
    }

    fn is_ready(&self) -> bool {
        let buf_ok = self
            .buffer
            .lock()
            .map(|buf| buf.len() >= self.sequence_length)
            .unwrap_or(false);
        let model_ok = self.model.lock().is_ok();
        buf_ok && model_ok
    }
}

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

    fn test_config() -> Mamba2Config {
        Mamba2Config {
            d_model: 32,
            d_state: 8,
            d_head: 8,
            num_heads: 2,
            expand: 2,
            num_layers: 1,
            max_seq_len: 8,
            dropout: 0.0,
            ..Default::default()
        }
    }

    const TEST_SEQ_LEN: usize = 4;

    #[test]
    fn test_mamba2_adapter_creation() {
        let adapter = Mamba2InferenceAdapter::new(test_config(), TEST_SEQ_LEN);
        assert!(
            adapter.is_ok(),
            "Mamba2 adapter creation failed: {:?}",
            adapter.err()
        );
        let adapter = adapter.unwrap();
        assert_eq!(adapter.model_name(), "MAMBA-2");
        assert!(!adapter.is_ready(), "Should not be ready with empty buffer");
    }

    #[test]
    fn test_mamba2_adapter_buffers_and_predicts() {
        let adapter = Mamba2InferenceAdapter::new(test_config(), TEST_SEQ_LEN).unwrap();

        for i in 0..TEST_SEQ_LEN {
            let fv = FeatureVector {
                values: vec![0.1 * (i as f64 + 1.0); 32],
                timestamp: 1700000000_000_000 + i as i64,
            };
            let pred = adapter.predict(&fv).unwrap();

            if i < TEST_SEQ_LEN - 1 {
                assert_eq!(pred.direction, 0.0, "Neutral direction while buffering");
                assert_eq!(pred.confidence, 0.0, "Zero confidence while buffering");
            } else {
                assert!(
                    pred.direction >= -1.0 && pred.direction <= 1.0,
                    "direction {} out of [-1,1]",
                    pred.direction
                );
                assert!(
                    pred.confidence >= 0.0 && pred.confidence <= 1.0,
                    "confidence {} out of [0,1]",
                    pred.confidence
                );
            }
        }

        assert!(adapter.is_ready());
    }

    #[test]
    fn test_mamba2_adapter_deterministic() {
        let adapter = Mamba2InferenceAdapter::new(test_config(), TEST_SEQ_LEN).unwrap();

        // Fill buffer with identical feature vectors
        for _ in 0..TEST_SEQ_LEN {
            let fv = FeatureVector {
                values: vec![0.1; 32],
                timestamp: 1700000000_000_000,
            };
            let _ = adapter.predict(&fv);
        }

        let fv = FeatureVector {
            values: vec![0.1; 32],
            timestamp: 1700000000_000_000,
        };
        let pred1 = adapter.predict(&fv).unwrap();
        let pred2 = adapter.predict(&fv).unwrap();
        assert!(
            (pred1.direction - pred2.direction).abs() < 1e-6,
            "direction mismatch: {} vs {}",
            pred1.direction,
            pred2.direction
        );
    }
}
