//! DQN inference adapter for ensemble prediction
//!
//! Wraps a loaded DQN model and normalizes its Q-value output
//! into a directional signal + confidence for ensemble aggregation.

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

use cudarc::driver::CudaStream;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::device::MlDevice;

use crate::dqn::dqn::{DQNConfig, DQN};
use crate::ensemble::inference_adapter::{
    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
};
use crate::gpu::DeviceConfig;
use crate::{MLError, MLResult};

/// Inference adapter that wraps a DQN model for ensemble prediction.
///
/// Converts raw Q-values into a normalized directional signal [-1, 1]
/// and a confidence score [0, 1] based on softmax probability of the
/// best action.
#[allow(missing_debug_implementations)]
pub struct DqnInferenceAdapter {
    model: Mutex<DQN>,
    device: MlDevice,
    stream: Arc<CudaStream>,
}

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

impl DqnInferenceAdapter {
    /// Create a new DQN inference adapter from configuration.
    pub fn new(config: DQNConfig) -> MLResult<Self> {
        let device = DeviceConfig::Auto.resolve().map_err(|e| MLError::DeviceError(format!("CUDA required: {e}")))?;
        let stream = device.cuda_stream()?.clone();
        let model = DQN::new(config)?;
        Ok(Self {
            model: Mutex::new(model),
            device,
            stream,
        })
    }

    /// Create a DQN inference adapter on a specific device.
    #[cfg(test)]
    pub(crate) fn new_on_device(config: DQNConfig, device: MlDevice) -> MLResult<Self> {
        let stream = device.cuda_stream()?.clone();
        let model = DQN::new_on_device(config, device.clone())?;
        Ok(Self {
            model: Mutex::new(model),
            device,
            stream,
        })
    }

    /// Create a DQN inference adapter and load weights from a safetensors checkpoint.
    pub fn from_checkpoint(config: DQNConfig, path: &str) -> MLResult<Self> {
        let device = DeviceConfig::Auto.resolve().map_err(|e| MLError::DeviceError(format!("CUDA required: {e}")))?;
        let stream = device.cuda_stream()?.clone();
        let mut model = DQN::new(config)?;
        model.load_from_safetensors(path)?;
        Ok(Self {
            model: Mutex::new(model),
            device,
            stream,
        })
    }
}

impl ModelInferenceAdapter for DqnInferenceAdapter {
    fn model_name(&self) -> &str {
        "DQN"
    }

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

        let f32_values: Vec<f32> = features.values.iter().map(|&v| v as f32).collect();
        let len = f32_values.len();

        // Create GPU tensor [1, feature_dim]
        let input = GpuTensor::from_host(&f32_values, vec![1, len], &self.stream)?;

        // Run forward pass through the Q-network
        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("DQN model lock poisoned: {e}")))?;
        let q_output = model.forward(&input)?;

        // Extract Q-values to host
        let q_vec_f32 = q_output.to_host(&self.stream)?;
        let q_vec: Vec<f64> = q_vec_f32.iter().map(|&v| v as f64).collect();

        let num_actions = q_vec.len();
        if num_actions == 0 {
            return Err(MLError::InferenceError(
                "DQN produced zero-length Q-value vector".to_owned(),
            ));
        }

        // Find best action index (argmax of Q-values)
        let best_idx = q_vec
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(idx, _)| idx)
            .ok_or_else(|| MLError::InferenceError("Failed to find max Q-value".to_owned()))?;

        // Compute directional signal: map action index to [-1, 1]
        let center = (num_actions as f64 - 1.0) / 2.0;
        let direction = if center > 0.0 {
            ((best_idx as f64 - center) / center).clamp(-1.0, 1.0)
        } else {
            0.0
        };

        // Compute confidence via softmax probability of the best action
        let max_q = q_vec
            .iter()
            .copied()
            .fold(f64::NEG_INFINITY, f64::max);
        let exp_sum: f64 = q_vec.iter().map(|&q| (q - max_q).exp()).sum();
        let best_q = q_vec
            .get(best_idx)
            .copied()
            .ok_or_else(|| MLError::InferenceError("Best index out of bounds".to_owned()))?;
        let confidence = (best_q - max_q).exp() / exp_sum;

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

        Ok(EnsemblePrediction {
            model_name: "DQN".to_owned(),
            direction,
            confidence,
            metadata: PredictionMeta {
                latency_us,
                quantiles: None,
                attention_weights: None,
                q_values: Some(q_vec),
            },
        })
    }

    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 feature_dim = batch
            .first()
            .map(|f| f.values.len())
            .unwrap_or(0);

        // Single flat buffer: N x feature_dim f32 values -> one GPU upload
        let mut flat = Vec::with_capacity(n * feature_dim);
        for fv in batch {
            for &v in &fv.values {
                flat.push(v as f32);
            }
        }

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

        // One forward pass for the entire batch: [N, feature_dim] -> [N, num_actions]
        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("DQN model lock poisoned: {e}")))?;
        let q_output = model.forward(&input)?;
        drop(model);

        let q_flat = q_output.to_host(&self.stream)?;
        let num_actions = if n > 0 { q_flat.len() / n } else { 0 };
        let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64;

        let mut results = Vec::with_capacity(n);
        for batch_idx in 0..n {
            if num_actions == 0 {
                return Err(MLError::InferenceError(
                    "DQN produced zero-length Q-value vector".to_owned(),
                ));
            }

            let row_start = batch_idx * num_actions;
            let q_vec: Vec<f64> = (0..num_actions)
                .map(|i| q_flat.get(row_start + i).copied().unwrap_or(0.0) as f64)
                .collect();

            let best_idx = q_vec
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                .map(|(idx, _)| idx)
                .ok_or_else(|| MLError::InferenceError("argmax failed".to_owned()))?;

            let center = (num_actions as f64 - 1.0) / 2.0;
            let direction = if center > 0.0 {
                ((best_idx as f64 - center) / center).clamp(-1.0, 1.0)
            } else {
                0.0
            };

            let max_q = q_vec.iter().copied().fold(f64::NEG_INFINITY, f64::max);
            let exp_sum: f64 = q_vec.iter().map(|&q| (q - max_q).exp()).sum();
            let best_q = q_vec
                .get(best_idx)
                .copied()
                .ok_or_else(|| MLError::InferenceError("best_idx OOB".to_owned()))?;
            let confidence = (best_q - max_q).exp() / exp_sum;

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

        Ok(results)
    }

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

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

    static SHARED_CUDA: OnceLock<MlDevice> = OnceLock::new();

    fn shared_device() -> MlDevice {
        SHARED_CUDA
            .get_or_init(|| {
                DeviceConfig::Auto
                    .resolve()
                    .expect("CUDA required")
            })
            .clone()
    }

    fn test_config() -> DQNConfig {
        DQNConfig {
            state_dim: 56,
            num_actions: 9,
            hidden_dims: vec![64, 64],
            ..Default::default()
        }
    }

    #[test]
    fn test_dqn_adapter_creation() {
        let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap();
        assert_eq!(adapter.model_name(), "DQN");
        assert!(adapter.is_ready());
    }

    #[test]
    fn test_dqn_adapter_predict_direction_range() {
        let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 56],
            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
        );
        assert!(
            pred.metadata.q_values.is_some(),
            "q_values metadata should be Some"
        );
    }

    #[test]
    fn test_dqn_adapter_deterministic() {
        let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 56],
            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"
        );
    }

    #[test]
    #[ignore = "branching always on: checkpoint round-trip needs branching weight save/load (not just base q_network)"]
    fn test_dqn_checkpoint_round_trip() {
        let config = test_config();
        let adapter = DqnInferenceAdapter::new_on_device(config.clone(), shared_device());
        assert!(
            adapter.is_ok(),
            "failed to create adapter: {:?}",
            adapter.err()
        );
        let adapter = adapter.unwrap();

        let fv = FeatureVector {
            values: vec![0.3; 56],
            timestamp: 1700000000,
        };
        let pred1 = adapter.predict(&fv);
        assert!(pred1.is_ok(), "first predict failed: {:?}", pred1.err());
        let pred1 = pred1.unwrap();

        // Save checkpoint to temp dir
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("dqn_test.safetensors");
        {
            let model = adapter.model.lock().unwrap();
            let vars = model.get_q_network_vars();
            let host_data = vars.export_to_host().unwrap(); // gpu-exit: test checkpoint
            let arch_metadata = Some(model.config.checkpoint_metadata());

            // Build safetensors TensorView map from host data
            let mut tensor_map = std::collections::HashMap::new();
            // Keep byte buffers alive for the duration of serialize_to_file
            let byte_vecs: Vec<(String, Vec<usize>, Vec<u8>)> = host_data
                .into_iter()
                .map(|(name, (shape, data))| {
                    let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
                    (name, shape, bytes)
                })
                .collect();
            for (name, shape, bytes) in &byte_vecs {
                let view = safetensors::tensor::TensorView::new(
                    safetensors::Dtype::F32,
                    shape.clone(),
                    bytes,
                )
                .unwrap();
                tensor_map.insert(name.as_str(), view);
            }
            safetensors::serialize_to_file(
                tensor_map,
                arch_metadata,
                path.as_path(),
            )
            .unwrap();
        }

        // Load into new adapter
        let path_str = path.to_str().unwrap();
        let adapter2 = DqnInferenceAdapter::from_checkpoint(config, path_str);
        assert!(
            adapter2.is_ok(),
            "from_checkpoint failed: {:?}",
            adapter2.err()
        );
        let adapter2 = adapter2.unwrap();
        let pred2 = adapter2.predict(&fv);
        assert!(
            pred2.is_ok(),
            "second predict failed: {:?}",
            pred2.err()
        );
        let pred2 = pred2.unwrap();

        assert!(
            (pred1.direction - pred2.direction).abs() < 1e-4,
            "round-trip mismatch: {} vs {}",
            pred1.direction,
            pred2.direction
        );
    }
}
