//! Rainbow DQN Network with Dueling Architecture and C51 Distributional Output
//!
//! This implementation combines:
//! - Dueling networks (Wang et al., 2016)
//! - Distributional RL with C51 (Bellemare et al., 2017)
//! - Noisy networks for exploration (Fortunato et al., 2018)

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::cuda_autograd::GpuTensor;
use serde::{Deserialize, Serialize};

use super::distributional::{CategoricalDistribution, DistributionalConfig};
use super::noisy_layers::NoisyLinear;
use ml_core::MLError;

/// Activation function types
#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
pub enum ActivationType {
    #[default]
    ReLU,
    LeakyReLU,
    Swish,
    ELU,
    GELU,
    Mish,
}

/// Rainbow network configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RainbowNetworkConfig {
    pub input_size: usize,
    pub hidden_sizes: Vec<usize>,
    pub num_actions: usize,
    pub activation: ActivationType,
    pub dropout_rate: f64,
    pub distributional: DistributionalConfig,
    /// Initial noise std dev for noisy layers (Rainbow DQN default: 0.5, scaled by 1/sqrt(in))
    pub noisy_sigma_init: f64,
    pub dueling: bool,
    pub use_spectral_norm: bool, // Q-value stability feature
    pub spectral_norm_iterations: usize,
}

impl Default for RainbowNetworkConfig {
    fn default() -> Self {
        Self {
            input_size: 64,
            hidden_sizes: vec![512, 512],
            num_actions: 4,
            activation: ActivationType::ReLU,
            dropout_rate: 0.1,
            distributional: DistributionalConfig::default(),
            noisy_sigma_init: 0.5,
            dueling: true,
            use_spectral_norm: false, // Disabled by default, enable for unstable training
            spectral_norm_iterations: 1,
        }
    }
}

/// Rainbow `DQN` network with distributional outputs
#[allow(missing_debug_implementations)]
pub struct RainbowNetwork {
    config: RainbowNetworkConfig,

    // Shared feature extractor (NoisyLinear layers)
    feature_layers: Vec<NoisyLinear>,

    // Dueling architecture (NoisyLinear layers)
    value_stream: Vec<NoisyLinear>,
    advantage_stream: Vec<NoisyLinear>,

    // Final distributional layers
    value_distribution: NoisyLinear,
    advantage_distribution: NoisyLinear,

    // Distribution handler
    categorical_dist: CategoricalDistribution,

    /// Dropout rate (0.0 = identity)
    _dropout_rate: f32,

    /// CUDA stream for GPU operations
    stream: Arc<CudaStream>,
}

impl RainbowNetwork {
    pub fn new(stream: Arc<CudaStream>, config: RainbowNetworkConfig) -> Result<Self, MLError> {
        let categorical_dist = CategoricalDistribution::new(&config.distributional, &stream)?;

        // Create feature extraction layers (always NoisyLinear)
        let mut feature_layers = Vec::new();
        let mut current_size = config.input_size;

        for &hidden_size in config.hidden_sizes.iter() {
            let noisy_layer = NoisyLinear::new(current_size, hidden_size, stream.clone(), config.noisy_sigma_init)?;
            feature_layers.push(noisy_layer);
            current_size = hidden_size;
        }

        let final_feature_size = current_size;

        // Create dueling streams if enabled (always NoisyLinear)
        let (value_stream, advantage_stream) = if config.dueling {
            let value_hidden = final_feature_size / 2;
            let value_layer = NoisyLinear::new(final_feature_size, value_hidden, stream.clone(), config.noisy_sigma_init)?;

            let advantage_hidden = final_feature_size / 2;
            let advantage_layer = NoisyLinear::new(final_feature_size, advantage_hidden, stream.clone(), config.noisy_sigma_init)?;

            (vec![value_layer], vec![advantage_layer])
        } else {
            (Vec::new(), Vec::new())
        };

        // Final distributional output layers
        let num_atoms = config.distributional.num_atoms;

        let value_dist_in = if config.dueling { final_feature_size / 2 } else { final_feature_size };
        let value_distribution = NoisyLinear::new(value_dist_in, num_atoms, stream.clone(), config.noisy_sigma_init)?;

        let adv_dist_in = if config.dueling { final_feature_size / 2 } else { final_feature_size };
        let adv_dist_out = config.num_actions * num_atoms;
        let advantage_distribution = NoisyLinear::new(adv_dist_in, adv_dist_out, stream.clone(), config.noisy_sigma_init)?;

        let dropout_rate_f32 = config.dropout_rate as f32;

        Ok(Self {
            config,
            feature_layers,
            value_stream,
            advantage_stream,
            value_distribution,
            advantage_distribution,
            categorical_dist,
            _dropout_rate: dropout_rate_f32,
            stream,
        })
    }

    pub fn forward(&self, input: &GpuTensor) -> Result<GpuTensor, MLError> {
        // Rainbow forward pass (cold path):
        // 1. Feature extraction: NoisyLinear + activation
        // 2. Dueling split: value stream + advantage stream
        // 3. Distributional output: [batch, num_actions, num_atoms]
        // 4. Softmax over atoms dimension
        let activations = ml_core::cuda_autograd::ActivationKernels::new(&self.stream)?;

        // Feature extraction
        let mut x = input.gpu_clone(&self.stream)?;
        for layer in &self.feature_layers {
            x = layer.forward(&x)?;
            // Apply activation (ReLU by default for Rainbow)
            let (activated, _mask) = activations.relu_fwd(&x, &self.stream)?;
            x = activated;
        }

        let num_atoms = self.config.distributional.num_atoms;
        let num_actions = self.config.num_actions;

        if self.config.dueling {
            // Value stream
            let mut v = x.gpu_clone(&self.stream)?;
            for layer in &self.value_stream {
                v = layer.forward(&v)?;
                let (activated, _mask) = activations.relu_fwd(&v, &self.stream)?;
                v = activated;
            }
            v = self.value_distribution.forward(&v)?;
            // v: [batch, num_atoms]

            // Advantage stream
            let mut a = x;
            for layer in &self.advantage_stream {
                a = layer.forward(&a)?;
                let (activated, _mask) = activations.relu_fwd(&a, &self.stream)?;
                a = activated;
            }
            a = self.advantage_distribution.forward(&a)?;
            // a: [batch, num_actions * num_atoms]

            // Combine: Z(s,a) = V(s) + A(s,a) - mean(A)
            // Host-side for cold path
            let v_host = v.to_host(&self.stream)?;
            let a_host = a.to_host(&self.stream)?;
            let batch = v_host.len() / num_atoms;

            let mut output = Vec::with_capacity(batch * num_actions * num_atoms);
            for b in 0..batch {
                // Compute mean advantage per atom
                let mut mean_a = vec![0.0_f32; num_atoms];
                for act in 0..num_actions {
                    for atom in 0..num_atoms {
                        let idx = b * num_actions * num_atoms + act * num_atoms + atom;
                        let val = a_host.get(idx).copied().unwrap_or(0.0);
                        if let Some(m) = mean_a.get_mut(atom) {
                            *m += val;
                        }
                    }
                }
                for m in &mut mean_a {
                    *m /= num_actions as f32;
                }

                for act in 0..num_actions {
                    for atom in 0..num_atoms {
                        let v_val = v_host.get(b * num_atoms + atom).copied().unwrap_or(0.0);
                        let a_idx = b * num_actions * num_atoms + act * num_atoms + atom;
                        let a_val = a_host.get(a_idx).copied().unwrap_or(0.0);
                        let mean_val = mean_a.get(atom).copied().unwrap_or(0.0);
                        output.push(v_val + a_val - mean_val);
                    }
                }
            }

            // Apply softmax over atoms dimension (per action)
            let mut softmax_output = Vec::with_capacity(output.len());
            for b in 0..batch {
                for act in 0..num_actions {
                    let start = b * num_actions * num_atoms + act * num_atoms;
                    // Numerically stable softmax
                    let logits: Vec<f32> = (0..num_atoms)
                        .map(|i| output.get(start + i).copied().unwrap_or(0.0))
                        .collect();
                    let max_val = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                    let exps: Vec<f32> = logits.iter().map(|&v| (v - max_val).exp()).collect();
                    let sum_exp: f32 = exps.iter().sum();
                    let inv_sum = if sum_exp > 0.0 { 1.0 / sum_exp } else { 1.0 / num_atoms as f32 };
                    for &e in &exps {
                        softmax_output.push(e * inv_sum);
                    }
                }
            }

            GpuTensor::from_host(
                &softmax_output,
                vec![batch, num_actions, num_atoms],
                &self.stream,
            )
        } else {
            // Non-dueling: direct output through advantage distribution
            let a = self.advantage_distribution.forward(&x)?;
            let a_host = a.to_host(&self.stream)?;
            let batch = a_host.len() / (num_actions * num_atoms);

            // Apply softmax over atoms per action
            let mut softmax_output = Vec::with_capacity(a_host.len());
            for b in 0..batch {
                for act in 0..num_actions {
                    let start = b * num_actions * num_atoms + act * num_atoms;
                    let logits: Vec<f32> = (0..num_atoms)
                        .map(|i| a_host.get(start + i).copied().unwrap_or(0.0))
                        .collect();
                    let max_val = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                    let exps: Vec<f32> = logits.iter().map(|&v| (v - max_val).exp()).collect();
                    let sum_exp: f32 = exps.iter().sum();
                    let inv_sum = if sum_exp > 0.0 { 1.0 / sum_exp } else { 1.0 / num_atoms as f32 };
                    for &e in &exps {
                        softmax_output.push(e * inv_sum);
                    }
                }
            }

            GpuTensor::from_host(
                &softmax_output,
                vec![batch, num_actions, num_atoms],
                &self.stream,
            )
        }
    }

    pub fn get_q_values(&self, distributions: &GpuTensor) -> Result<GpuTensor, MLError> {
        // distributions: [batch, num_actions, num_atoms]
        // Q(s,a) = sum(probs * support) for each action
        // Returns: [batch, num_actions]
        let (batch, num_actions, num_atoms) = distributions.dims3().map_err(|e| {
            MLError::ModelError(format!("get_q_values dims3: {e}"))
        })?;

        let dist_host = distributions.to_host(&self.stream)?;
        let support = &self.categorical_dist.support_host();

        let mut q_values = Vec::with_capacity(batch * num_actions);
        for b in 0..batch {
            for a in 0..num_actions {
                let mut q = 0.0_f32;
                for atom in 0..num_atoms {
                    let idx = b * num_actions * num_atoms + a * num_atoms + atom;
                    let prob = dist_host.get(idx).copied().unwrap_or(0.0);
                    let z = support.get(atom).copied().unwrap_or(0.0);
                    q += prob * z;
                }
                q_values.push(q);
            }
        }

        GpuTensor::from_host(&q_values, vec![batch, num_actions], &self.stream)
    }

    pub const fn config(&self) -> &RainbowNetworkConfig {
        &self.config
    }

    pub const fn categorical_distribution(&self) -> &CategoricalDistribution {
        &self.categorical_dist
    }
}

#[cfg(test)]
#[allow(
    clippy::map_err_ignore,
    clippy::field_reassign_with_default,
    clippy::unnecessary_wraps
)]
mod tests {
    use super::*;

    #[test]
    fn test_rainbow_config_default() -> Result<(), MLError> {
        let config = RainbowNetworkConfig::default();
        assert!(config.input_size > 0);
        assert!(config.num_actions > 0);
        assert!(!config.hidden_sizes.is_empty());
        Ok(())
    }
}
