//! Diffusion model inference adapter for ensemble prediction
//!
//! Wraps a `Denoiser` and runs a single forward pass at t=1
//! (minimal noise level). The mean of the denoiser output is
//! mapped via sigmoid to direction + confidence.

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

use cudarc::driver::CudaStream;

use crate::diffusion::config::DiffusionConfig;
use crate::diffusion::denoiser::Denoiser;

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

use ml_supervised::gpu_tensor::{gpu_mean_all, GpuTensor};

/// Inference adapter for the Diffusion (DDPM/DDIM) denoising model.
///
/// Uses the `Denoiser` network with a minimal timestep (t=1) to process
/// features. The mean of the denoiser output provides the raw signal
/// which is mapped via sigmoid to direction [-1, 1] and confidence [0, 1].
#[allow(missing_debug_implementations)]
pub struct DiffusionInferenceAdapter {
    model: Mutex<Denoiser>,
    stream: Arc<CudaStream>,
    data_dim: usize,
}

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

impl DiffusionInferenceAdapter {
    /// Create a new Diffusion inference adapter.
    pub fn new(config: DiffusionConfig) -> MLResult<Self> {
        let data_dim = config.seq_len * config.feature_dim;

        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 denoiser = Denoiser::new(
            data_dim,
            config.hidden_dim,
            config.num_layers,
            config.time_embed_dim,
            &stream,
        )?;

        Ok(Self {
            model: Mutex::new(denoiser),
            stream,
            data_dim,
        })
    }

    /// Load from a checkpoint (placeholder -- weight loading TBD).
    pub fn from_checkpoint(config: DiffusionConfig, _path: &str) -> MLResult<Self> {
        // For now, create fresh model. Checkpoint loading will be implemented
        // when the binary weight format is finalized.
        Self::new(config)
    }

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

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

        let padded = self.pad_features(&features.values);
        let input = GpuTensor::from_vec(padded, &[1, self.data_dim], &self.stream)?;

        let t = vec![1_u32];

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

        let mean_val = gpu_mean_all(&output)?;

        let raw_val = mean_val as f64;
        let prob = 1.0 / (1.0 + (-raw_val).exp());
        let direction = (2.0_f64 * 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: "Diffusion".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_vec(padded, &[1, self.data_dim], &self.stream)?;

        let t = vec![1_u32];

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

        let raw_f32 = gpu_mean_all(&output)?;
        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.data_dim);
        for fv in batch {
            flat.extend_from_slice(&self.pad_features(&fv.values));
        }

        let input = GpuTensor::from_vec(flat, &[n, self.data_dim], &self.stream)?;
        let t = vec![1_u32; n];

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

        // Output is [N, data_dim]; take mean per row
        let output_host = output.to_vec()?; // gpu-exit: batch prediction
        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 row_start = i * self.data_dim;
            let row_end = row_start + self.data_dim;
            let row = output_host.get(row_start..row_end).unwrap_or(&[]);
            let row_mean: f32 = if row.is_empty() {
                0.0
            } else {
                row.iter().sum::<f32>() / row.len() as f32
            };

            let raw_val = row_mean as f64;
            let prob = 1.0 / (1.0 + (-raw_val).exp());
            let direction = (2.0_f64 * 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: "Diffusion".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::*;
    use crate::diffusion::config::NoiseSchedule;

    fn test_config() -> DiffusionConfig {
        DiffusionConfig {
            feature_dim: 51,
            seq_len: 1,
            hidden_dim: 64,
            num_layers: 2,
            time_embed_dim: 32,
            num_timesteps: 100,
            sampling_steps: 5,
            schedule: NoiseSchedule::Cosine,
            ..Default::default()
        }
    }

    #[test]
    fn test_diffusion_adapter_creation() {
        let adapter = DiffusionInferenceAdapter::new(test_config());
        assert!(
            adapter.is_ok(),
            "Diffusion creation failed: {:?}",
            adapter.err()
        );
        let adapter = adapter.unwrap();
        assert_eq!(adapter.model_name(), "Diffusion");
        assert!(adapter.is_ready());
    }

    #[test]
    fn test_diffusion_adapter_predict_direction_range() {
        let adapter = DiffusionInferenceAdapter::new(test_config()).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1_700_000_000_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_diffusion_adapter_deterministic() {
        let adapter = DiffusionInferenceAdapter::new(test_config()).unwrap();
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1_700_000_000_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"
        );
    }
}
