//! PPO inference adapter for ensemble prediction
//!
//! Wraps a loaded PPO model (MLP actor) and normalizes its action
//! probability output into a directional signal + confidence for
//! ensemble aggregation.

use std::sync::Mutex;

use ml_core::device::MlDevice;

use crate::ensemble::inference_adapter::{
    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
};
use crate::gpu::DeviceConfig;
use crate::ppo::ppo::{PPOConfig, PPO};
use crate::{MLError, MLResult};

/// Inference adapter that wraps a PPO model for ensemble prediction.
///
/// Converts the actor's softmax action probabilities into a normalized
/// directional signal [-1, 1] using probability-weighted action values,
/// and uses the max probability as confidence [0, 1].
#[allow(missing_debug_implementations)]
pub struct PpoInferenceAdapter {
    model: Mutex<PPO>,
    state_dim: usize,
    #[allow(dead_code)]
    device: MlDevice,
}

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

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

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

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

        let padded = self.pad_features(&features.values);

        // Get action probabilities from the actor network (softmax output)
        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("PPO model lock poisoned: {e}")))?;
        let probs = model.actor.forward_host(&padded, 1)?;

        let num_actions = probs.len();
        if num_actions == 0 {
            return Err(MLError::InferenceError(
                "PPO produced zero-length probability vector".to_owned(),
            ));
        }

        // Compute directional signal: weighted sum of probs * centered action values
        let center = (num_actions as f64 - 1.0) / 2.0;
        let mut weighted_sum = 0.0_f64;
        let mut max_prob = 0.0_f32;
        for (i, &p) in probs.iter().enumerate() {
            if center > 0.0 {
                let action_val = (i as f64 - center) / center;
                weighted_sum += p as f64 * action_val;
            }
            if p > max_prob {
                max_prob = p;
            }
        }
        let direction = if center > 0.0 {
            weighted_sum.clamp(-1.0, 1.0)
        } else {
            0.0
        };

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

        Ok(EnsemblePrediction {
            model_name: "PPO".to_owned(),
            direction,
            confidence: confidence.clamp(0.0, 1.0),
            metadata: PredictionMeta {
                latency_us,
                quantiles: None,
                attention_weights: None,
                q_values: None,
            },
        })
    }

    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.state_dim);
        for fv in batch {
            flat.extend_from_slice(&self.pad_features(&fv.values));
        }

        let model = self
            .model
            .lock()
            .map_err(|e| MLError::LockError(format!("PPO model lock poisoned: {e}")))?;
        let probs = model.actor.forward_host(&flat, n)?;
        drop(model);

        let num_actions = if n > 0 { probs.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(
                    "PPO produced zero-length probability vector".to_owned(),
                ));
            }

            let row_start = batch_idx * num_actions;
            let center = (num_actions as f64 - 1.0) / 2.0;
            let mut weighted_sum = 0.0_f64;
            let mut max_prob = 0.0_f32;
            for i in 0..num_actions {
                let p = probs.get(row_start + i).copied().unwrap_or(0.0);
                if center > 0.0 {
                    let action_val = (i as f64 - center) / center;
                    weighted_sum += p as f64 * action_val;
                }
                if p > max_prob {
                    max_prob = p;
                }
            }
            let direction = if center > 0.0 {
                weighted_sum.clamp(-1.0, 1.0)
            } else {
                0.0
            };

            let confidence = max_prob as f64;

            results.push(EnsemblePrediction {
                model_name: "PPO".to_owned(),
                direction,
                confidence: confidence.clamp(0.0, 1.0),
                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::*;

    fn test_config() -> PPOConfig {
        PPOConfig {
            state_dim: 64,
            num_actions: 45,
            policy_hidden_dims: vec![64, 64],
            value_hidden_dims: vec![64, 64],
            ..Default::default()
        }
    }

    #[test]
    fn test_ppo_adapter_creation() {
        let adapter = PpoInferenceAdapter::new(test_config()).unwrap();
        assert_eq!(adapter.model_name(), "PPO");
        assert!(adapter.is_ready());
        assert_eq!(adapter.state_dim, 64);
    }

    #[test]
    fn test_ppo_adapter_pads_input_to_64() {
        let adapter = PpoInferenceAdapter::new(test_config()).unwrap();
        // 51-dim input (DQN canonical size) should be zero-padded to 64
        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_ppo_adapter_deterministic() {
        let adapter = PpoInferenceAdapter::new(test_config()).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"
        );
    }
}
