//! Distributional Dueling Q-Network Architecture
//!
//! Combines Dueling Networks (Wang et al., 2016) with Distributional RL (Bellemare et al., 2017)
//! as implemented in Rainbow DQN (Hessel et al., 2018).
//!
//! ## Architecture
//!
//! ```text
//! State [state_dim] → Shared Features [hidden_dim]
//!                     ↓                ↓
//!                Value V(s)         Advantage A(s,a)
//!                [num_atoms]        [num_actions, num_atoms]
//!                     ↓                ↓
//!    Z(s,a) = V(s) + A(s,a) - mean(A(s,·)) [DISTRIBUTIONAL SPACE]
//!                     ↓
//!    Output: [batch, num_actions, num_atoms]
//! ```
//!
//! ## Key Features
//!
//! - **Distributional Value Stream**: V(s) as a distribution over returns [`num_atoms`]
//! - **Distributional Advantage Stream**: A(s,a) as distributions per action [`num_actions`, `num_atoms`]
//! - **Mean Subtraction in Distribution Space**: Ensures identifiability (zero mean advantage)
//! - **Full Return Distributions**: Outputs probability distributions instead of scalar Q-values
//!
//! ## Mathematical Formulation
//!
//! For each atom `z_i` in the return distribution:
//! `Z(s,a,z_i)` = `V(s,z_i)` + [`A(s,a,z_i)` - (1/|A|) * `Σ_a'` A(s,a',`z_i`)]
//!
//! Where:
//! - `V(s,z_i)`: Probability of atom `z_i` for state value
//! - `A(s,a,z_i)`: Probability of atom `z_i` for advantage of action a
//! - `mean(A(s,·,z_i))`: Average advantage across all actions for atom `z_i`
//!
//! ## References
//!
//! - Rainbow: Combining Improvements in Deep RL (Hessel et al., 2018)
//! - Dueling Network Architectures (Wang et al., 2016)
//! - A Distributional Perspective on Reinforcement Learning (Bellemare et al., 2017)

use std::sync::Arc;

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

use crate::rmsnorm::RMSNorm;
use ml_core::MLError;

/// Configuration for Distributional Dueling Q-Network
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributionalDuelingConfig {
    /// Input state dimension
    pub state_dim: usize,
    /// Number of actions
    pub num_actions: usize,
    /// Number of atoms for value distribution (Rainbow DQN standard: 51)
    pub num_atoms: usize,
    /// Shared feature layer dimensions
    pub shared_hidden_dims: Vec<usize>,
    /// Value stream hidden dimension
    pub value_hidden_dim: usize,
    /// Advantage stream hidden dimension
    pub advantage_hidden_dim: usize,
    /// `LeakyReLU` negative slope
    pub leaky_relu_alpha: f64,
}

impl DistributionalDuelingConfig {
    /// Create default distributional dueling config from basic parameters
    pub const fn new(
        state_dim: usize,
        num_actions: usize,
        num_atoms: usize,
        shared_hidden_dims: Vec<usize>,
        value_hidden_dim: usize,
        advantage_hidden_dim: usize,
    ) -> Self {
        Self {
            state_dim,
            num_actions,
            num_atoms,
            shared_hidden_dims,
            value_hidden_dim,
            advantage_hidden_dim,
            leaky_relu_alpha: 0.01,
        }
    }

    /// Create from `DQNConfig` parameters
    ///
    /// All `hidden_dims` become shared feature layers. The value and advantage
    /// heads branch off after the last shared layer. This matches the CUDA
    /// experience-collection kernel which hardcodes 2 shared layers.
    pub fn from_dqn_params(
        state_dim: usize,
        num_actions: usize,
        num_atoms: usize,
        hidden_dims: &[usize],
        dueling_hidden_dim: usize,
        leaky_relu_alpha: f64,
    ) -> Self {
        Self {
            state_dim,
            num_actions,
            num_atoms,
            shared_hidden_dims: hidden_dims.to_vec(),
            value_hidden_dim: dueling_hidden_dim,
            advantage_hidden_dim: dueling_hidden_dim,
            leaky_relu_alpha,
        }
    }
}

/// Distributional Dueling Q-Network with separate value and advantage streams
/// operating in distributional space
#[allow(missing_debug_implementations)]
pub struct DistributionalDuelingQNetwork {
    /// Shared feature extraction layers
    shared_layers: Vec<GpuLinear>,

    /// `RMSNorm` after each shared hidden layer
    shared_norms: Vec<RMSNorm>,

    /// Value stream layers (outputs distribution)
    value_fc: GpuLinear,
    value_out: GpuLinear, // Output: [batch, num_atoms]

    /// `RMSNorm` after value stream hidden layer
    value_norm: RMSNorm,

    /// Advantage stream layers (outputs distributions per action)
    advantage_fc: GpuLinear,
    advantage_out: GpuLinear, // Output: [batch, num_actions * num_atoms]

    /// `RMSNorm` after advantage stream hidden layer
    advantage_norm: RMSNorm,

    /// Configuration
    config: DistributionalDuelingConfig,

    /// `GpuVarStore` for weight management
    vars: GpuVarStore,

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

impl DistributionalDuelingQNetwork {
    /// Create new distributional dueling Q-network
    ///
    /// # Arguments
    ///
    /// * `config` - Distributional dueling network configuration
    /// * `stream` - CUDA stream for GPU operations
    ///
    /// # Returns
    ///
    /// New `DistributionalDuelingQNetwork` instance with Xavier-initialized weights
    pub fn new(config: DistributionalDuelingConfig, stream: Arc<CudaStream>) -> Result<Self, MLError> {
        let mut vars = GpuVarStore::new(stream.clone());

        // Build shared feature layers with RMSNorm after each
        let mut shared_layers = Vec::new();
        let mut shared_norms = Vec::new();
        let mut current_dim = config.state_dim;

        for (i, &hidden_dim) in config.shared_hidden_dims.iter().enumerate() {
            let layer_name = format!("shared_{}", i);
            let layer = vars.linear_xavier(&layer_name, current_dim, hidden_dim)?;
            shared_layers.push(layer);

            let norm_stream = stream.clone();
            let norm = RMSNorm::new_default(norm_stream, hidden_dim)?;
            shared_norms.push(norm);

            current_dim = hidden_dim;
        }

        // Value stream (outputs distribution over returns)
        let value_fc = vars.linear_xavier("value_fc", current_dim, config.value_hidden_dim)?;
        let value_norm = RMSNorm::new_default(stream.clone(), config.value_hidden_dim)?;

        // Near-zero init for output layers
        let value_out = vars.linear_near_zero("value_out", config.value_hidden_dim, config.num_atoms)?;

        // Advantage stream (outputs distributions per action)
        let advantage_fc = vars.linear_xavier("advantage_fc", current_dim, config.advantage_hidden_dim)?;
        let advantage_norm = RMSNorm::new_default(stream.clone(), config.advantage_hidden_dim)?;

        let advantage_out = vars.linear_near_zero(
            "advantage_out",
            config.advantage_hidden_dim,
            config.num_actions * config.num_atoms,
        )?;

        Ok(Self {
            shared_layers,
            shared_norms,
            value_fc,
            value_out,
            value_norm,
            advantage_fc,
            advantage_out,
            advantage_norm,
            config,
            vars,
            stream,
        })
    }

    /// Forward pass through distributional dueling network
    ///
    /// # Arguments
    ///
    /// * `state` - State tensor [`batch_size`, `state_dim`]
    ///
    /// # Returns
    ///
    /// Distribution tensor [`batch_size`, `num_actions`, `num_atoms`]
    /// representing probability distributions over returns for each action
    pub fn forward(&self, state: &GpuTensor) -> Result<GpuTensor, MLError> {
        // Distributional Dueling forward pass (cold path).
        // 1. Shared features: linear + LeakyReLU + RMSNorm
        // 2. Value stream: [batch, num_atoms]
        // 3. Advantage stream: [batch, num_actions * num_atoms]
        // 4. Combine: Z(s,a) = V(s) + A(s,a) - mean(A)
        // 5. Softmax over atoms dimension
        let cublas = cudarc::cublas::CudaBlas::new(self.stream.clone()).map_err(|e| {
            MLError::ModelError(format!("DistDueling forward cublas init: {e}"))
        })?;
        let act_kernels = ml_core::cuda_autograd::ActivationKernels::new(&self.stream)?;
        let alpha = self.config.leaky_relu_alpha as f32;

        // Shared feature extraction
        let mut h = state.gpu_clone(&self.stream)?;
        for (layer, norm) in self.shared_layers.iter().zip(self.shared_norms.iter()) {
            let (out, _acts) = layer.forward(&h, &self.vars, &cublas, &self.stream)?;
            let (activated, _mask) = act_kernels.leaky_relu_fwd(&out, alpha, &self.stream)?;
            h = norm.forward(&activated)?;
        }

        // Value stream
        let (v_hidden, _) = self.value_fc.forward(&h, &self.vars, &cublas, &self.stream)?;
        let (v_act, _) = act_kernels.leaky_relu_fwd(&v_hidden, alpha, &self.stream)?;
        let v_normed = self.value_norm.forward(&v_act)?;
        let (v_logits, _) = self.value_out.forward(&v_normed, &self.vars, &cublas, &self.stream)?;
        // v_logits: [batch, num_atoms]

        // Advantage stream
        let (a_hidden, _) = self.advantage_fc.forward(&h, &self.vars, &cublas, &self.stream)?;
        let (a_act, _) = act_kernels.leaky_relu_fwd(&a_hidden, alpha, &self.stream)?;
        let a_normed = self.advantage_norm.forward(&a_act)?;
        let (a_logits, _) = self.advantage_out.forward(&a_normed, &self.vars, &cublas, &self.stream)?;
        // a_logits: [batch, num_actions * num_atoms]

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

        // Combine and softmax on host (cold path)
        let v_host = v_logits.to_host(&self.stream)?;
        let a_host = a_logits.to_host(&self.stream)?;
        let batch = v_host.len() / num_atoms;

        // Z(s,a,z_i) = V(s,z_i) + A(s,a,z_i) - mean_a'(A(s,a',z_i))
        let mut combined = 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 {
                // Softmax over atoms for this action
                let mut logits = Vec::with_capacity(num_atoms);
                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);
                    logits.push(v_val + a_val - mean_val);
                }

                // Numerically stable softmax
                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 {
                    combined.push(e * inv_sum);
                }
            }
        }

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

    /// Get `GpuVarStore` for weight serialization
    pub const fn vars(&self) -> &GpuVarStore {
        &self.vars
    }

    /// Mutable reference to the `GpuVarStore` (for checkpoint import).
    pub const fn vars_mut(&mut self) -> &mut GpuVarStore {
        &mut self.vars
    }

    /// Get CUDA stream
    pub const fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

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

    /// Copy weights from another distributional dueling network
    pub fn copy_weights_from(
        &mut self,
        other: &DistributionalDuelingQNetwork,
    ) -> Result<(), MLError> {
        let host_data = other.vars.export_to_host()?; // gpu-exit: checkpoint weight export
        self.vars.import_from_host(&host_data)
    }
}

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

    #[test]
    fn test_distributional_dueling_from_dqn_params() -> anyhow::Result<()> {
        let config = DistributionalDuelingConfig::from_dqn_params(
            32,              // state_dim
            5,               // num_actions (5 exposure levels)
            51,              // num_atoms
            &[256, 128, 64], // hidden_dims
            64,              // dueling_hidden_dim
            0.01,            // leaky_relu_alpha
        );

        assert_eq!(config.state_dim, 32);
        assert_eq!(config.num_actions, 5);
        assert_eq!(config.num_atoms, 51);
        assert_eq!(config.shared_hidden_dims, vec![256, 128, 64]);
        assert_eq!(config.value_hidden_dim, 64);
        assert_eq!(config.advantage_hidden_dim, 64);

        Ok(())
    }
}
