//! TLOB inference adapter for ensemble prediction
//!
//! Wraps a GPU projection matching TLOBTrainableAdapter's flatten+project
//! architecture.  Buffers `seq_len` feature vectors, flattens to
//! `[1, seq_len * feature_dim]`, and projects through a 3-layer MLP to a
//! scalar output.  The scalar is mapped via sigmoid to direction + confidence
//! for ensemble aggregation.

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

use cudarc::cublas::CudaBlas;
use cudarc::driver::CudaStream;

use ml_core::cuda_autograd::{GpuLinear, GpuTensor, GpuVarStore, ActivationKernels};

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

/// 3-layer MLP projection: flat_dim -> hidden_dim -> hidden_dim/2 -> 1
struct TlobProjection {
    linear1: GpuLinear,
    linear2: GpuLinear,
    linear3: GpuLinear,
    activations: ActivationKernels,
    var_store: GpuVarStore,
    cublas: CudaBlas,
    stream: Arc<CudaStream>,
}

impl TlobProjection {
    fn new(flat_dim: usize, hidden_dim: usize, stream: &Arc<CudaStream>) -> MLResult<Self> {
        let mid_dim = hidden_dim / 2;
        let mut var_store = GpuVarStore::new(Arc::clone(stream));
        let linear1 = var_store.linear("tlob_linear1", flat_dim, hidden_dim)
            .map_err(|e| MLError::ModelError(format!("TLOB linear1 init: {e}")))?;
        let linear2 = var_store.linear("tlob_linear2", hidden_dim, mid_dim)
            .map_err(|e| MLError::ModelError(format!("TLOB linear2 init: {e}")))?;
        let linear3 = var_store.linear("tlob_linear3", mid_dim, 1)
            .map_err(|e| MLError::ModelError(format!("TLOB linear3 init: {e}")))?;
        let activations = ActivationKernels::new(stream)
            .map_err(|e| MLError::ModelError(format!("TLOB activations init: {e}")))?;
        let cublas = CudaBlas::new(Arc::clone(stream))
            .map_err(|e| MLError::ModelError(format!("TLOB cublas init: {e}")))?;
        Ok(Self {
            linear1,
            linear2,
            linear3,
            activations,
            var_store,
            cublas,
            stream: Arc::clone(stream),
        })
    }

    fn forward(&self, input: &GpuTensor) -> MLResult<GpuTensor> {
        let (h, _act1) = self.linear1.forward(input, &self.var_store, &self.cublas, &self.stream)
            .map_err(|e| MLError::ModelError(format!("TLOB forward linear1: {e}")))?;
        let (h, _mask1) = self.activations.relu_fwd(&h, &self.stream)
            .map_err(|e| MLError::ModelError(format!("TLOB forward relu1: {e}")))?;
        let (h, _act2) = self.linear2.forward(&h, &self.var_store, &self.cublas, &self.stream)
            .map_err(|e| MLError::ModelError(format!("TLOB forward linear2: {e}")))?;
        let (h, _mask2) = self.activations.relu_fwd(&h, &self.stream)
            .map_err(|e| MLError::ModelError(format!("TLOB forward relu2: {e}")))?;
        let (out, _act3) = self.linear3.forward(&h, &self.var_store, &self.cublas, &self.stream)
            .map_err(|e| MLError::ModelError(format!("TLOB forward linear3: {e}")))?;
        Ok(out)
    }
}

/// Inference adapter for the TLOB (Transformer Limit Order Book) model.
///
/// Maintains a ring buffer of `sequence_length` feature vectors.  While the
/// buffer is filling, returns a neutral (direction=0, confidence=0) prediction.
/// Once full, flattens the buffer to `[1, seq_len * feature_dim]`, runs it
/// through the 3-layer MLP projection, and maps the scalar output via sigmoid
/// to direction [-1, 1] and confidence [0, 1].
#[allow(missing_debug_implementations)]
pub struct TlobInferenceAdapter {
    model: Mutex<TlobProjection>,
    buffer: Mutex<VecDeque<Vec<f32>>>,
    stream: Arc<CudaStream>,
    sequence_length: usize,
    feature_dim: usize,
}

// SAFETY: TlobProjection 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 TlobInferenceAdapter {}
#[allow(unsafe_code)]
unsafe impl Sync for TlobInferenceAdapter {}

impl TlobInferenceAdapter {
    /// Create a new TLOB inference adapter.
    ///
    /// # Arguments
    /// * `feature_dim` - Feature vector dimension (typically 51)
    /// * `hidden_dim` - Hidden layer width (32-64 for RTX 3050 Ti)
    /// * `sequence_length` - Number of feature vectors to buffer before inference
    pub fn new(feature_dim: usize, hidden_dim: usize, 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 flat_dim = sequence_length * feature_dim;
        let projection = TlobProjection::new(flat_dim, hidden_dim, &stream)?;

        Ok(Self {
            model: Mutex::new(projection),
            buffer: Mutex::new(VecDeque::with_capacity(sequence_length)),
            stream,
            sequence_length,
            feature_dim,
        })
    }

    /// Load from a safetensors checkpoint.
    pub fn from_checkpoint(
        feature_dim: usize,
        hidden_dim: usize,
        sequence_length: usize,
        _path: &str,
    ) -> MLResult<Self> {
        // TODO: Load GpuLinear weights from safetensors checkpoint
        tracing::warn!("TLOB checkpoint loading not yet supported for GpuLinear weights");
        Self::new(feature_dim, hidden_dim, sequence_length)
    }

    /// Pad or truncate feature values to match `feature_dim`, converting f64 to f32.
    fn pad_features(&self, values: &[f64]) -> Vec<f32> {
        let mut padded = vec![0.0_f32; self.feature_dim];
        let copy_len = values.len().min(self.feature_dim);
        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 TlobInferenceAdapter {
    fn model_name(&self) -> &str {
        "TLOB"
    }

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

        // Buffer phase (short-lived lock)
        let buffer_ready = {
            let mut buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("TLOB buffer lock poisoned: {e}")))?;
            buf.push_back(padded);
            while buf.len() > self.sequence_length {
                buf.pop_front();
            }
            buf.len() >= self.sequence_length
        };
        // buffer lock is dropped here

        if !buffer_ready {
            let latency_us = start.elapsed().as_micros() as u64;
            return Ok(EnsemblePrediction {
                model_name: "TLOB".to_owned(),
                direction: 0.0,
                confidence: 0.0,
                metadata: PredictionMeta {
                    latency_us,
                    ..Default::default()
                },
            });
        }

        // Flatten buffer to [1, seq_len * feature_dim]
        let flat_data = {
            let buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("TLOB buffer lock poisoned: {e}")))?;
            let mut data = Vec::with_capacity(self.sequence_length * self.feature_dim);
            for frame in buf.iter() {
                data.extend_from_slice(frame);
            }
            data
        };
        // buffer lock is dropped here

        let flat_dim = self.sequence_length * self.feature_dim;
        let input = GpuTensor::from_host(&flat_data, vec![1, flat_dim], &self.stream)?;

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

        let output_host = output.to_host(&self.stream)?;
        let raw_val = output_host.first().copied().unwrap_or(0.0) as f64;

        // Sigmoid -> direction + confidence
        let prob = 1.0 / (1.0 + (-raw_val).exp());
        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: "TLOB".to_owned(),
            direction,
            confidence,
            metadata: PredictionMeta {
                latency_us,
                ..Default::default()
            },
        })
    }

    fn predict_raw(&self, features: &FeatureVector) -> MLResult<RawPrediction> {
        let padded = self.pad_features(&features.values);

        // Buffer phase (short-lived lock)
        let buffer_ready = {
            let mut buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("TLOB 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 {
            return Ok(RawPrediction {
                direction_scalar: 0.0,
                confidence: 0.0,
                logits: None,
            });
        }

        // Flatten buffer to [1, seq_len * feature_dim]
        let flat_data = {
            let buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("TLOB buffer lock poisoned: {e}")))?;
            let mut data = Vec::with_capacity(self.sequence_length * self.feature_dim);
            for frame in buf.iter() {
                data.extend_from_slice(frame);
            }
            data
        };

        let flat_dim = self.sequence_length * self.feature_dim;
        let input = GpuTensor::from_host(&flat_data, vec![1, flat_dim], &self.stream)?;

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

        let output_host = output.to_host(&self.stream)?;
        let raw_f32 = output_host.first().copied().unwrap_or(0.0);
        let prob = 1.0 / (1.0 + (-raw_f32 as f64).exp());
        let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);

        Ok(RawPrediction {
            direction_scalar: 0.0,
            confidence,
            logits: Some(vec![raw_f32]),
        })
    }

    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::*;

    const FEAT_DIM: usize = 51;
    const HIDDEN: usize = 64;
    const SEQ_LEN: usize = 4;

    #[test]
    fn test_tlob_adapter_creation() {
        let adapter = TlobInferenceAdapter::new(FEAT_DIM, HIDDEN, SEQ_LEN);
        assert!(
            adapter.is_ok(),
            "TLOB adapter creation failed: {:?}",
            adapter.err()
        );
        let adapter = match adapter {
            Ok(a) => a,
            Err(e) => panic!("TLOB adapter creation failed: {e}"),
        };
        assert_eq!(adapter.model_name(), "TLOB");
        assert!(
            !adapter.is_ready(),
            "Should not be ready with empty buffer"
        );
    }

    #[test]
    fn test_tlob_adapter_buffers_and_predicts() {
        let adapter = match TlobInferenceAdapter::new(FEAT_DIM, HIDDEN, SEQ_LEN) {
            Ok(a) => a,
            Err(e) => panic!("TLOB adapter creation failed: {e}"),
        };

        for i in 0..SEQ_LEN {
            let fv = FeatureVector {
                values: vec![0.1 * (i as f64 + 1.0); FEAT_DIM],
                timestamp: 1700000000_000_000 + i as i64,
            };
            let pred = match adapter.predict(&fv) {
                Ok(p) => p,
                Err(e) => panic!("predict failed at step {i}: {e}"),
            };

            if i < SEQ_LEN - 1 {
                assert_eq!(
                    pred.direction, 0.0,
                    "Expected neutral direction before buffer full"
                );
                assert_eq!(
                    pred.confidence, 0.0,
                    "Expected zero confidence before buffer full"
                );
            } 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(), "Adapter should be ready after filling buffer");
    }

    #[test]
    fn test_tlob_adapter_deterministic() {
        let adapter = match TlobInferenceAdapter::new(FEAT_DIM, HIDDEN, SEQ_LEN) {
            Ok(a) => a,
            Err(e) => panic!("TLOB adapter creation failed: {e}"),
        };

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

        // Now make two predictions with the same input (buffer is full and stable)
        let fv = FeatureVector {
            values: vec![0.1; FEAT_DIM],
            timestamp: 1700000000_000_000,
        };
        let pred1 = match adapter.predict(&fv) {
            Ok(p) => p,
            Err(e) => panic!("predict1 failed: {e}"),
        };
        let pred2 = match adapter.predict(&fv) {
            Ok(p) => p,
            Err(e) => panic!("predict2 failed: {e}"),
        };

        assert!(
            (pred1.direction - pred2.direction).abs() < 1e-6,
            "Deterministic predictions should have same direction: {} vs {}",
            pred1.direction,
            pred2.direction
        );
        assert!(
            (pred1.confidence - pred2.confidence).abs() < 1e-6,
            "Deterministic predictions should have same confidence: {} vs {}",
            pred1.confidence,
            pred2.confidence
        );
    }
}
