//! TFT inference adapter for ensemble prediction
//!
//! Wraps a Temporal Fusion Transformer model and normalizes its quantile
//! output into a directional signal + confidence for ensemble aggregation.
//!
//! Unlike single-step models (DQN, PPO), TFT requires a sequence of
//! feature vectors before it can produce a prediction. This adapter
//! maintains an internal ring buffer that collects `sequence_length`
//! observations before running the first forward pass.

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::tft::{TFTConfig, TemporalFusionTransformer};
use crate::{MLError, MLResult};

/// Inference adapter that wraps a TFT model for ensemble prediction.
///
/// Buffers incoming feature vectors until `sequence_length` observations
/// are available, then runs the full TFT forward pass.  The median
/// quantile is mapped to a directional signal and the IQR provides a
/// confidence estimate.
#[allow(missing_debug_implementations)]
pub struct TftInferenceAdapter {
    model: Mutex<TemporalFusionTransformer>,
    buffer: Mutex<VecDeque<FeatureVector>>,
    stream: Arc<CudaStream>,
    sequence_length: usize,
    num_static: usize,
    num_known: usize,
    num_unknown: usize,
    prediction_horizon: usize,
    num_quantiles: usize,
}

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

impl TftInferenceAdapter {
    /// Create a new TFT inference adapter from configuration.
    ///
    /// # Arguments
    /// * `config` - TFT model configuration.  `input_dim` **must** equal
    ///   `num_static_features + num_known_features + num_unknown_features`.
    /// * `sequence_length` - Number of feature vectors to buffer before
    ///   the first forward pass.  Overrides `config.sequence_length` for
    ///   the buffer size (the config value is still used internally by
    ///   the model).
    pub fn new(config: TFTConfig, 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 num_static = config.num_static_features;
        let num_known = config.num_known_features;
        let num_unknown = config.num_unknown_features;
        let prediction_horizon = config.prediction_horizon;
        let num_quantiles = config.num_quantiles;

        let model = TemporalFusionTransformer::new_with_stream(config, Some(Arc::clone(&stream)))?;

        Ok(Self {
            model: Mutex::new(model),
            buffer: Mutex::new(VecDeque::with_capacity(sequence_length + 1)),
            stream,
            sequence_length,
            num_static,
            num_known,
            num_unknown,
            prediction_horizon,
            num_quantiles,
        })
    }

    /// Extract a fixed-length f32 slice from a feature vector starting at
    /// `offset`, zero-padding if the source is too short.
    fn extract_features(values: &[f64], offset: usize, count: usize) -> Vec<f32> {
        let mut out = vec![0.0_f32; count];
        for i in 0..count {
            if let Some(&v) = values.get(offset.wrapping_add(i)) {
                if let Some(slot) = out.get_mut(i) {
                    *slot = v as f32;
                }
            }
        }
        out
    }

    /// Build calendar-based future features for `prediction_horizon` steps.
    ///
    /// Layout per step (zero-padded to `num_known`):
    ///   0: hour_sin,  1: hour_cos
    ///   2: day_sin,   3: day_cos
    ///   4: minute_sin, 5: minute_cos
    fn build_future_features(&self, base_timestamp: i64) -> Vec<f32> {
        let mut out = Vec::with_capacity(self.prediction_horizon * self.num_known);
        for step in 0..self.prediction_horizon {
            // Advance timestamp by step (assume 1-second ticks)
            let ts = base_timestamp.wrapping_add(step as i64);

            // Derive simple cyclical features from timestamp
            let seconds_in_day = 86400.0_f64;
            let seconds_in_hour = 3600.0_f64;
            let seconds_in_minute = 60.0_f64;

            let day_frac = (ts as f64).rem_euclid(seconds_in_day) / seconds_in_day;
            let hour_frac = (ts as f64).rem_euclid(seconds_in_hour) / seconds_in_hour;
            let minute_frac = (ts as f64).rem_euclid(seconds_in_minute) / seconds_in_minute;

            let two_pi = std::f64::consts::TAU;

            let mut step_feats = vec![0.0_f32; self.num_known];
            // hour sin/cos
            if let Some(slot) = step_feats.get_mut(0) {
                *slot = (hour_frac * two_pi).sin() as f32;
            }
            if let Some(slot) = step_feats.get_mut(1) {
                *slot = (hour_frac * two_pi).cos() as f32;
            }
            // day sin/cos
            if let Some(slot) = step_feats.get_mut(2) {
                *slot = (day_frac * two_pi).sin() as f32;
            }
            if let Some(slot) = step_feats.get_mut(3) {
                *slot = (day_frac * two_pi).cos() as f32;
            }
            // minute sin/cos
            if let Some(slot) = step_feats.get_mut(4) {
                *slot = (minute_frac * two_pi).sin() as f32;
            }
            if let Some(slot) = step_feats.get_mut(5) {
                *slot = (minute_frac * two_pi).cos() as f32;
            }

            out.extend_from_slice(&step_feats);
        }
        out
    }

    /// Map raw median quantile value to a directional signal in [-1, 1]
    /// using a smooth saturating function: sign(x) * (1 - exp(-|x|)).
    fn median_to_direction(median: f64) -> f64 {
        if median == 0.0 {
            return 0.0;
        }
        let sign = if median > 0.0 { 1.0 } else { -1.0 };
        let magnitude = 1.0 - (-median.abs()).exp();
        (sign * magnitude).clamp(-1.0, 1.0)
    }

    /// Map IQR relative to median magnitude to a confidence score [0, 1].
    ///
    /// Tight quantile spread (low IQR) relative to a strong signal yields
    /// high confidence.  When the median is near zero the model is
    /// indecisive, so confidence should be low.
    fn iqr_to_confidence(median: f64, iqr: f64) -> f64 {
        let abs_median = median.abs();
        if abs_median < 1e-10 {
            return 0.0;
        }
        // confidence = median_strength * certainty
        // certainty  = 1 / (1 + iqr / |median|)
        let certainty = 1.0 / (1.0 + iqr / abs_median);
        let strength = 1.0 - (-abs_median).exp();
        (strength * certainty).clamp(0.0, 1.0)
    }
}

impl ModelInferenceAdapter for TftInferenceAdapter {
    fn model_name(&self) -> &str {
        "TFT"
    }

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

        // --- buffer phase (short-lived lock) ---
        let buffered_data: Option<(Vec<FeatureVector>, i64)> = {
            let mut buf = self
                .buffer
                .lock()
                .map_err(|e| MLError::LockError(format!("TFT buffer lock poisoned: {e}")))?;

            buf.push_back(features.clone());

            // Trim to sequence_length (keep most recent)
            while buf.len() > self.sequence_length {
                buf.pop_front();
            }

            if buf.len() < self.sequence_length {
                None
            } else {
                // Snapshot the buffer contents so we can drop the lock
                let snapshot: Vec<FeatureVector> = buf.iter().cloned().collect();
                let last_ts = snapshot
                    .last()
                    .map(|fv| fv.timestamp)
                    .unwrap_or(features.timestamp);
                Some((snapshot, last_ts))
            }
        }; // buffer lock dropped here

        // If we don't have enough data yet, return a neutral prediction
        let (snapshot, last_ts) = match buffered_data {
            Some(data) => data,
            None => {
                let latency_us = start.elapsed().as_micros() as u64;
                return Ok(EnsemblePrediction {
                    model_name: "TFT".to_owned(),
                    direction: 0.0,
                    confidence: 0.0,
                    metadata: PredictionMeta {
                        latency_us,
                        quantiles: None,
                        attention_weights: None,
                        q_values: None,
                    },
                });
            }
        };

        // --- tensor construction ---

        // Static features: extract from the first FV at offset 40
        let static_offset = 40;
        let first_fv_values = snapshot
            .first()
            .map(|fv| fv.values.as_slice())
            .ok_or_else(|| MLError::InferenceError("Empty snapshot".to_owned()))?;
        let static_f32 = Self::extract_features(first_fv_values, static_offset, self.num_static);

        // Historical (unknown) features: from each FV at offset 0
        let mut hist_f32 = Vec::with_capacity(self.sequence_length * self.num_unknown);
        for fv in &snapshot {
            let row = Self::extract_features(&fv.values, 0, self.num_unknown);
            hist_f32.extend_from_slice(&row);
        }

        // Future (known) features: calendar-derived
        let future_f32 = self.build_future_features(last_ts);

        // Build StreamTensors from f32 data (all compute is F32 on GPU)
        // ml-supervised uses StreamTensor aliased as GpuTensor internally.
        let static_tensor = StreamTensor::from_vec(
            static_f32,
            &[1, self.num_static],
            &self.stream,
        )?;

        let hist_tensor = StreamTensor::from_vec(
            hist_f32,
            &[1, self.sequence_length, self.num_unknown],
            &self.stream,
        )?;

        let future_tensor = StreamTensor::from_vec(
            future_f32,
            &[1, self.prediction_horizon, self.num_known],
            &self.stream,
        )?;

        // --- forward pass (model lock) ---
        let quantile_preds = {
            let mut model = self
                .model
                .lock()
                .map_err(|e| MLError::LockError(format!("TFT model lock poisoned: {e}")))?;
            model.forward(&static_tensor, &hist_tensor, &future_tensor)?
        };

        // --- output processing ---
        // quantile_preds shape: [1, prediction_horizon, num_quantiles]
        // Download to host and extract first horizon step
        let output_host = quantile_preds.to_vec()?; // gpu-exit: prediction quantiles
        let total_quantiles = self.prediction_horizon * self.num_quantiles;
        if output_host.len() < self.num_quantiles || total_quantiles == 0 {
            return Err(MLError::InferenceError("TFT produced empty output".to_owned()));
        }

        // First horizon step: elements [0..num_quantiles]
        let median_idx = self.num_quantiles / 2;
        let q25_idx = self.num_quantiles / 4;
        let q75_idx = (self.num_quantiles * 3) / 4;

        let median = output_host.get(median_idx).copied().unwrap_or(0.0) as f64;
        let q25 = output_host.get(q25_idx).copied().unwrap_or(0.0) as f64;
        let q75 = output_host.get(q75_idx).copied().unwrap_or(0.0) as f64;
        let iqr = (q75 - q25).abs();

        let direction = Self::median_to_direction(median);
        let confidence = Self::iqr_to_confidence(median, iqr);

        // Collect all quantiles for metadata (first horizon)
        let mut quantile_values = Vec::with_capacity(self.num_quantiles);
        for qi in 0..self.num_quantiles {
            let qv = output_host.get(qi).copied().unwrap_or(0.0) as f64;
            quantile_values.push(qv);
        }

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

        Ok(EnsemblePrediction {
            model_name: "TFT".to_owned(),
            direction,
            confidence,
            metadata: PredictionMeta {
                latency_us,
                quantiles: Some(quantile_values),
                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::*;

    /// Small config for fast tests.
    /// input_dim = num_static + num_known + num_unknown = 6 + 6 + 8 = 20
    fn test_config() -> TFTConfig {
        TFTConfig {
            input_dim: 20,
            hidden_dim: 32,
            num_heads: 2,
            num_layers: 1,
            prediction_horizon: 5,
            sequence_length: 4,
            num_quantiles: 9,
            num_static_features: 6,
            num_known_features: 6,
            num_unknown_features: 8,
            dropout_rate: 0.0,
            ..Default::default()
        }
    }

    fn make_fv(ts: i64) -> FeatureVector {
        // 51 values -- more than our small config needs, adapter will
        // extract/zero-pad as required.
        FeatureVector {
            values: vec![0.1; 51],
            timestamp: ts,
        }
    }

    #[test]
    fn test_tft_adapter_creation() {
        let adapter = TftInferenceAdapter::new(test_config(), 4)
            .expect("TftInferenceAdapter::new should succeed");
        assert_eq!(adapter.model_name(), "TFT");
        // Buffer is empty, so the adapter is NOT ready yet
        assert!(
            !adapter.is_ready(),
            "Adapter should not be ready with an empty buffer"
        );
    }

    #[test]
    fn test_tft_adapter_buffers_and_predicts() {
        let seq_len = 4;
        let adapter = TftInferenceAdapter::new(test_config(), seq_len)
            .expect("TftInferenceAdapter::new should succeed");

        // Feed seq_len - 1 feature vectors; each should return neutral
        for i in 0..(seq_len - 1) {
            let fv = make_fv(1_700_000_000 + i as i64);
            let pred = adapter.predict(&fv).expect("predict should not error");
            assert_eq!(
                pred.direction, 0.0,
                "Should return neutral direction while buffering"
            );
            assert_eq!(
                pred.confidence, 0.0,
                "Should return zero confidence while buffering"
            );
        }

        assert!(
            !adapter.is_ready(),
            "Adapter should still not be ready with {} / {} entries",
            seq_len - 1,
            seq_len
        );

        // Feed the final FV to fill the buffer
        let fv_final = make_fv(1_700_000_000 + seq_len as i64);
        let pred = adapter
            .predict(&fv_final)
            .expect("predict should succeed once buffer is full");

        // Now the adapter should be ready
        assert!(
            adapter.is_ready(),
            "Adapter should be ready after receiving seq_len entries"
        );

        // Direction and confidence should be bounded
        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
        );

        // Quantile metadata should be present
        assert!(
            pred.metadata.quantiles.is_some(),
            "TFT predictions should include quantile metadata"
        );

        // One more prediction to ensure repeated calls work
        let fv_extra = make_fv(1_700_000_000 + (seq_len + 1) as i64);
        let pred2 = adapter
            .predict(&fv_extra)
            .expect("subsequent predict should succeed");
        assert!(
            pred2.direction >= -1.0 && pred2.direction <= 1.0,
            "second direction {} out of [-1,1]",
            pred2.direction
        );
    }

    #[test]
    fn test_tft_adapter_deterministic() {
        let seq_len = 4;
        // Test: same model, same input → same output (GPU determinism)
        let adapter = TftInferenceAdapter::new(test_config(), seq_len)
            .expect("adapter creation");

        // Fill the buffer
        for i in 0..seq_len {
            let fv = make_fv(1_700_000_000 + i as i64);
            let _ = adapter.predict(&fv);
        }

        // Two predictions with the same input on the same model
        let fv = make_fv(1_700_000_000 + seq_len as i64);
        let pred1 = adapter.predict(&fv).expect("predict 1");
        let pred2 = adapter.predict(&fv).expect("predict 2");

        // Same model + same input → must produce identical output
        assert!(
            (pred1.direction - pred2.direction).abs() < 1e-6,
            "Same-model predictions diverged: {} vs {}",
            pred1.direction, pred2.direction
        );
        assert!(
            (pred1.confidence - pred2.confidence).abs() < 1e-6,
            "Same-model confidences diverged: {} vs {}",
            pred1.confidence, pred2.confidence
        );
    }
}
