//! TGGN inference adapter for ensemble prediction
//!
//! Builds a 2-layer GPU projection matching the TGGNTrainableAdapter
//! architecture (input->hidden->1) for checkpoint-compatible inference.

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};

/// 2-layer GPU projection matching TGGNTrainableAdapter's architecture.
struct TggnProjection {
    linear1: GpuLinear,
    linear2: GpuLinear,
    activations: ActivationKernels,
    var_store: GpuVarStore,
    cublas: CudaBlas,
    stream: Arc<CudaStream>,
}

impl TggnProjection {
    fn new(input_dim: usize, hidden_dim: usize, stream: &Arc<CudaStream>) -> MLResult<Self> {
        let mut var_store = GpuVarStore::new(Arc::clone(stream));
        let linear1 = var_store.linear("tggn_linear1", input_dim, hidden_dim)
            .map_err(|e| MLError::ModelError(format!("TGGN linear1 init: {e}")))?;
        let linear2 = var_store.linear("tggn_linear2", hidden_dim, 1)
            .map_err(|e| MLError::ModelError(format!("TGGN linear2 init: {e}")))?;
        let activations = ActivationKernels::new(stream)
            .map_err(|e| MLError::ModelError(format!("TGGN activations init: {e}")))?;
        let cublas = CudaBlas::new(Arc::clone(stream))
            .map_err(|e| MLError::ModelError(format!("TGGN cublas init: {e}")))?;
        Ok(Self { linear1, linear2, 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!("TGGN forward linear1: {e}")))?;
        let (h, _mask) = self.activations.relu_fwd(&h, &self.stream)
            .map_err(|e| MLError::ModelError(format!("TGGN forward relu: {e}")))?;
        let (out, _act2) = self.linear2.forward(&h, &self.var_store, &self.cublas, &self.stream)
            .map_err(|e| MLError::ModelError(format!("TGGN forward linear2: {e}")))?;
        Ok(out)
    }
}

/// Inference adapter for the TGGN (Temporal Graph Gated Network).
///
/// Since the base TGGN model uses ndarray (not candle), this adapter
/// replicates the 2-layer projection from `TGGNTrainableAdapter`
/// (`input_dim -> hidden_dim -> 1`) for checkpoint-compatible inference.
/// The scalar output is mapped via sigmoid to direction [-1, 1] and
/// confidence [0, 1] for ensemble aggregation.
#[allow(missing_debug_implementations)]
pub struct TggnInferenceAdapter {
    model: Mutex<TggnProjection>,
    stream: Arc<CudaStream>,
    input_dim: usize,
}

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

impl TggnInferenceAdapter {
    /// Create a new TGGN inference adapter.
    ///
    /// # Arguments
    /// * `input_dim` - Feature vector dimension (typically 51)
    /// * `hidden_dim` - Hidden layer width (32-64 for RTX 3050 Ti)
    pub fn new(input_dim: usize, hidden_dim: 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 projection = TggnProjection::new(input_dim, hidden_dim, &stream)?;

        Ok(Self {
            model: Mutex::new(projection),
            stream,
            input_dim,
        })
    }

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

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

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

        let padded = self.pad_features(&features.values);
        let input = GpuTensor::from_host(&padded, vec![1, self.input_dim], &self.stream)?;

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("TGGN 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;

        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: "TGGN".to_owned(),
            direction,
            confidence,
            metadata: PredictionMeta {
                latency_us,
                quantiles: None,
                attention_weights: None,
                q_values: None,
            },
        })
    }

    fn predict_raw(&self, features: &FeatureVector) -> MLResult<RawPrediction> {
        let padded = self.pad_features(&features.values);
        let input = GpuTensor::from_host(&padded, vec![1, self.input_dim], &self.stream)?;

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("TGGN 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 predict_batch(&self, batch: &[FeatureVector]) -> MLResult<Vec<EnsemblePrediction>> {
        if batch.is_empty() {
            return Ok(Vec::new());
        }
        let start = std::time::Instant::now();
        let n = batch.len();

        let mut flat = Vec::with_capacity(n * self.input_dim);
        for fv in batch {
            flat.extend_from_slice(&self.pad_features(&fv.values));
        }

        let input = GpuTensor::from_host(&flat, vec![n, self.input_dim], &self.stream)?;

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

        let output_host = output.to_host(&self.stream)?;
        let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64;

        let mut results = Vec::with_capacity(n);
        for i in 0..n {
            let raw_val = output_host.get(i).copied().unwrap_or(0.0) as f64;
            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);

            results.push(EnsemblePrediction {
                model_name: "TGGN".to_owned(),
                direction,
                confidence,
                metadata: PredictionMeta {
                    latency_us,
                    quantiles: None,
                    attention_weights: None,
                    q_values: None,
                },
            });
        }

        Ok(results)
    }

    fn is_ready(&self) -> bool {
        self.model.lock().is_ok()
    }
}

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

    #[test]
    fn test_tggn_adapter_creation() {
        let adapter = TggnInferenceAdapter::new(51, 64);
        assert!(
            adapter.is_ok(),
            "TGGN adapter creation failed: {:?}",
            adapter.err()
        );
        let adapter = adapter.unwrap();
        assert_eq!(adapter.model_name(), "TGGN");
        assert!(adapter.is_ready());
    }

    #[test]
    fn test_tggn_adapter_predict_direction_range() {
        let adapter = TggnInferenceAdapter::new(51, 64).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1700000000_000_000,
        };
        let pred = adapter.predict(&fv).unwrap();
        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
        );
    }

    #[test]
    fn test_tggn_adapter_deterministic() {
        let adapter = TggnInferenceAdapter::new(51, 64).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1700000000_000_000,
        };
        let pred1 = adapter.predict(&fv).unwrap();
        let pred2 = adapter.predict(&fv).unwrap();
        assert_eq!(
            pred1.direction, pred2.direction,
            "Deterministic predictions should have same direction"
        );
    }
}
