//! Ensemble Q-Network for better uncertainty estimation
//!
//! Provides an ensemble of Q-networks that can be used together to:
//! - Estimate uncertainty via Q-value variance
//! - Make more robust predictions via voting/averaging
//! - Improve exploration through disagreement-based bonuses
//!
//! # Architecture
//!
//! The ensemble consists of multiple independent Q-networks with identical
//! architectures but different random initializations. This diversity allows
//! the ensemble to capture model uncertainty (epistemic uncertainty).

use serde::{Deserialize, Serialize};

use crate::network::{QNetwork, QNetworkConfig};
use ml_core::MLError;

/// Configuration for ensemble Q-network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsembleConfig {
    /// Base Q-network configuration
    pub base_config: QNetworkConfig,
    /// Number of networks in the ensemble
    pub num_networks: usize,
    /// Whether to use different random seeds for each network
    pub use_different_seeds: bool,
}

impl Default for EnsembleConfig {
    fn default() -> Self {
        Self {
            base_config: QNetworkConfig::default(),
            num_networks: 5,
            use_different_seeds: true,
        }
    }
}

/// Ensemble of Q-Networks for uncertainty estimation
///
/// The ensemble maintains multiple independent Q-networks and provides
/// methods to compute ensemble statistics (mean, variance, standard deviation).
#[allow(missing_debug_implementations)]
pub struct EnsembleQNetwork {
    /// Individual Q-networks in the ensemble
    networks: Vec<QNetwork>,
    /// Number of networks
    num_networks: usize,
    /// Configuration
    config: EnsembleConfig,
}

impl EnsembleQNetwork {
    /// Create new ensemble Q-network
    ///
    /// # Arguments
    ///
    /// * `config` - Base Q-network configuration (used for all networks)
    /// * `num_networks` - Number of networks in the ensemble (typically 3-10)
    ///
    /// # Returns
    ///
    /// Initialized ensemble with `num_networks` independent Q-networks
    ///
    /// # Errors
    ///
    /// Returns error if network creation fails
    pub fn new(
        config: QNetworkConfig,
        num_networks: usize,
    ) -> Result<Self, MLError> {
        if num_networks == 0 {
            return Err(MLError::InvalidInput(
                "Ensemble must have at least one network".to_owned(),
            ));
        }

        let mut networks = Vec::with_capacity(num_networks);

        // Create independent networks (diversity comes from random weight init, not epsilon)
        for _i in 0..num_networks {
            let net_config = config.clone();
            let network = QNetwork::new(net_config)?;
            networks.push(network);
        }

        let ensemble_config = EnsembleConfig {
            base_config: config,
            num_networks,
            use_different_seeds: true,
        };

        Ok(Self {
            networks,
            num_networks,
            config: ensemble_config,
        })
    }

    /// Forward pass through all networks in the ensemble
    ///
    /// # Arguments
    ///
    /// * `state` - Input state vector
    ///
    /// # Returns
    ///
    /// Vector of Q-value vectors, one per network.
    /// Each inner vector has shape [`num_actions`].
    pub fn forward(&self, state: &[f32]) -> Result<Vec<Vec<f32>>, MLError> {
        let mut q_values = Vec::with_capacity(self.num_networks);

        for network in &self.networks {
            let q_vals = network.forward(state)?;
            q_values.push(q_vals);
        }

        Ok(q_values)
    }

    /// Forward pass through all networks in the ensemble (batch API)
    ///
    /// # Arguments
    ///
    /// * `states` - Batch of state vectors
    ///
    /// # Returns
    ///
    /// Vector of batch Q-value vectors, one per network.
    /// Each inner Vec<Vec<f32>> has shape [`batch_size`][`num_actions`].
    pub fn forward_batch(&self, states: &[Vec<f32>]) -> Result<Vec<Vec<Vec<f32>>>, MLError> {
        let mut q_values = Vec::with_capacity(self.num_networks);

        for network in &self.networks {
            let batch_q = network.forward_batch(states)?;
            q_values.push(batch_q);
        }

        Ok(q_values)
    }

    /// Compute mean Q-values across ensemble
    ///
    /// # Arguments
    ///
    /// * `state` - Input state vector
    ///
    /// # Returns
    ///
    /// Mean Q-values across all networks (shape: [`num_actions`])
    pub fn mean_q(&self, state: &[f32]) -> Result<Vec<f32>, MLError> {
        let q_values = self.forward(state)?;

        if q_values.is_empty() {
            return Err(MLError::ModelError("No Q-values computed".to_owned()));
        }

        let num_actions = q_values[0].len();
        let mut mean_q = vec![0.0; num_actions];

        // Sum Q-values across all networks
        for q_vals in &q_values {
            for (i, &q) in q_vals.iter().enumerate() {
                mean_q[i] += q;
            }
        }

        // Divide by number of networks to get mean
        for q in &mut mean_q {
            *q /= self.num_networks as f32;
        }

        Ok(mean_q)
    }

    /// Compute mean Q-values across ensemble (batch API)
    ///
    /// # Arguments
    ///
    /// * `states` - Batch of state vectors
    ///
    /// # Returns
    ///
    /// Mean Q-values across all networks (shape: [`batch_size`][`num_actions`])
    pub fn mean_q_batch(&self, states: &[Vec<f32>]) -> Result<Vec<Vec<f32>>, MLError> {
        let q_values = self.forward_batch(states)?;

        if q_values.is_empty() {
            return Err(MLError::ModelError("No Q-values computed".to_owned()));
        }

        let batch_size = states.len();
        let num_actions = q_values[0].first().map_or(0, |r| r.len());

        let mut mean_q = vec![vec![0.0_f32; num_actions]; batch_size];

        for batch_q in &q_values {
            for (b, row) in batch_q.iter().enumerate() {
                for (a, &q) in row.iter().enumerate() {
                    mean_q[b][a] += q;
                }
            }
        }

        let n = self.num_networks as f32;
        for row in &mut mean_q {
            for v in row.iter_mut() {
                *v /= n;
            }
        }

        Ok(mean_q)
    }

    /// Compute standard deviation of Q-values across ensemble
    ///
    /// # Arguments
    ///
    /// * `state` - Input state vector
    ///
    /// # Returns
    ///
    /// Standard deviation of Q-values (shape: [`num_actions`])
    pub fn std_q(&self, state: &[f32]) -> Result<Vec<f32>, MLError> {
        let q_values = self.forward(state)?;

        if q_values.is_empty() {
            return Err(MLError::ModelError("No Q-values computed".to_owned()));
        }

        let num_actions = q_values[0].len();

        // First compute mean
        let mean_q = self.mean_q(state)?;

        // Compute variance: Var[X] = E[(X - E[X])^2]
        let mut variance = vec![0.0; num_actions];

        for q_vals in &q_values {
            for (i, &q) in q_vals.iter().enumerate() {
                let diff = q - mean_q[i];
                variance[i] += diff * diff;
            }
        }

        // Divide by number of networks
        for v in &mut variance {
            *v /= self.num_networks as f32;
        }

        // Standard deviation is sqrt(variance)
        let std: Vec<f32> = variance.iter().map(|v| v.sqrt()).collect();

        Ok(std)
    }

    /// Get number of networks in the ensemble
    pub const fn num_networks(&self) -> usize {
        self.num_networks
    }

    /// Get reference to a specific network
    ///
    /// # Arguments
    ///
    /// * `index` - Network index (0 to num_networks-1)
    ///
    /// # Returns
    ///
    /// Reference to the network, or None if index out of bounds
    pub fn get_network(&self, index: usize) -> Option<&QNetwork> {
        self.networks.get(index)
    }

    /// Get configuration
    pub const fn config(&self) -> &EnsembleConfig {
        &self.config
    }
}

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

    #[test]
    fn test_ensemble_creation() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![32, 16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 5)?;

        assert_eq!(ensemble.num_networks(), 5);

        Ok(())
    }

    #[test]
    fn test_ensemble_zero_networks_error() {
        let config = QNetworkConfig::default();
        let result = EnsembleQNetwork::new(config, 0);

        assert!(result.is_err());
        match result {
            Err(MLError::InvalidInput(msg)) => {
                assert!(msg.contains("at least one network"));
            }
            Err(e) => {
                // Wrong error variant
                assert!(false, "Expected InvalidInput error, got: {}", e);
            }
            Ok(_) => {
                assert!(false, "Expected error but got Ok");
            }
        }
    }

    #[test]
    fn test_forward_pass() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 3)?;
        let state = vec![1.0, 2.0, 3.0, 4.0];

        let q_values = ensemble.forward(&state)?;

        // Should have Q-values from 3 networks
        assert_eq!(q_values.len(), 3);
        // Each network should output 3 Q-values (one per action)
        for q_vals in &q_values {
            assert_eq!(q_vals.len(), 3);
        }

        Ok(())
    }

    #[test]
    fn test_mean_q_single_network() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        // Single network ensemble - mean should equal the network's output
        let ensemble = EnsembleQNetwork::new(config, 1)?;
        let state = vec![1.0, 2.0, 3.0, 4.0];

        let q_values = ensemble.forward(&state)?;
        let mean_q = ensemble.mean_q(&state)?;

        assert_eq!(mean_q.len(), 3);

        // Mean should equal the single network's output
        for (i, &q) in mean_q.iter().enumerate() {
            assert!((q - q_values[0][i]).abs() < 1e-6);
        }

        Ok(())
    }

    #[test]
    fn test_mean_q_multiple_networks() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 5)?;
        let state = vec![1.0, 2.0, 3.0, 4.0];

        let q_values = ensemble.forward(&state)?;
        let mean_q = ensemble.mean_q(&state)?;

        assert_eq!(mean_q.len(), 3);

        // Manually verify mean computation
        for action_idx in 0..3 {
            let sum: f32 = q_values.iter().map(|q| q[action_idx]).sum();
            let expected_mean = sum / 5.0;
            assert!((mean_q[action_idx] - expected_mean).abs() < 1e-5);
        }

        Ok(())
    }

    #[test]
    fn test_std_q_single_network() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        // Single network - std should be zero
        let ensemble = EnsembleQNetwork::new(config, 1)?;
        let state = vec![1.0, 2.0, 3.0, 4.0];

        let std_q = ensemble.std_q(&state)?;

        assert_eq!(std_q.len(), 3);

        // All standard deviations should be zero (single network)
        for &std in &std_q {
            assert!(std < 1e-6, "Expected zero std for single network, got {}", std);
        }

        Ok(())
    }

    #[test]
    fn test_std_q_multiple_networks() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 5)?;
        let state = vec![1.0, 2.0, 3.0, 4.0];

        let q_values = ensemble.forward(&state)?;
        let std_q = ensemble.std_q(&state)?;
        let mean_q = ensemble.mean_q(&state)?;

        assert_eq!(std_q.len(), 3);

        // Manually verify std computation
        for action_idx in 0..3 {
            let variance: f32 = q_values
                .iter()
                .map(|q| {
                    let diff = q[action_idx] - mean_q[action_idx];
                    diff * diff
                })
                .sum::<f32>()
                / 5.0;

            let expected_std = variance.sqrt();
            assert!(
                (std_q[action_idx] - expected_std).abs() < 1e-4,
                "Std mismatch at action {}: got {}, expected {}",
                action_idx,
                std_q[action_idx],
                expected_std
            );
        }

        Ok(())
    }

    #[test]
    fn test_get_network() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 3)?;

        // Valid indices
        assert!(ensemble.get_network(0).is_some());
        assert!(ensemble.get_network(1).is_some());
        assert!(ensemble.get_network(2).is_some());

        // Invalid index
        assert!(ensemble.get_network(3).is_none());
        assert!(ensemble.get_network(100).is_none());

        Ok(())
    }

    #[test]
    fn test_forward_batch_api() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 3)?;

        // Create batch of 2 states
        let states = vec![
            vec![1.0_f32, 2.0, 3.0, 4.0],
            vec![5.0_f32, 6.0, 7.0, 8.0],
        ];

        let q_values = ensemble.forward_batch(&states)?;

        // Should have Q-values from 3 networks
        assert_eq!(q_values.len(), 3);

        // Each network should produce 2 rows (batch_size=2) of 3 Q-values (num_actions=3)
        for batch_q in &q_values {
            assert_eq!(batch_q.len(), 2);
            for row in batch_q {
                assert_eq!(row.len(), 3);
            }
        }

        Ok(())
    }

    #[test]
    fn test_mean_q_batch_api() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 5)?;

        // Create batch of 2 states
        let states = vec![
            vec![1.0_f32, 2.0, 3.0, 4.0],
            vec![5.0_f32, 6.0, 7.0, 8.0],
        ];

        let mean_q = ensemble.mean_q_batch(&states)?;

        // Should have shape [2, 3] (batch_size=2, num_actions=3)
        assert_eq!(mean_q.len(), 2);
        for row in &mean_q {
            assert_eq!(row.len(), 3);
        }

        Ok(())
    }

    #[test]
    fn test_batch_consistency_with_single_api() -> Result<(), MLError> {
        let config = QNetworkConfig {
            state_dim: 4,
            num_actions: 3,
            hidden_dims: vec![16],
            ..QNetworkConfig::default()
        };

        let ensemble = EnsembleQNetwork::new(config, 3)?;

        let state_vec = vec![1.0_f32, 2.0, 3.0, 4.0];

        // Single API
        let mean_single = ensemble.mean_q(&state_vec)?;
        let std_single = ensemble.std_q(&state_vec)?;

        // Batch API (batch of 1)
        let states = vec![state_vec.clone()];
        let mean_batch = ensemble.mean_q_batch(&states)?;

        // Compare results (should be very close)
        for i in 0..3 {
            assert!(
                (mean_single[i] - mean_batch[0][i]).abs() < 1e-4,
                "Mean mismatch at index {}: single={}, batch={}",
                i,
                mean_single[i],
                mean_batch[0][i]
            );
        }

        // std_q is only for single states (no batch variant needed in common usage)
        assert_eq!(std_single.len(), 3);

        Ok(())
    }
}
