#![allow(unsafe_code)] // Required for CUDA kernel launches via cudarc.

//! Residual/Skip Connections for DQN Networks
//!
//! Implements residual blocks to improve gradient flow in deep Q-networks.
//! Skip connections help prevent gradient vanishing and enable deeper architectures.
//!
//! ## Weight storage
//!
//! Linear layers use `OwnedGpuLinear` (self-contained weight/bias `CudaSlice`s)
//! with cuBLAS sgemm. `LayerNorm` parameters are raw `CudaSlice`s.
//! The forward pass runs entirely on GPU via cuBLAS + CUDA kernels.

use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};

use cudarc::driver::CudaSlice;
use ml_core::cuda_autograd::{ActivationKernels, OwnedGpuLinear, GpuTensor};
use ml_core::MLError;

/// Configuration for residual blocks
#[derive(Debug, Clone)]
pub struct ResidualConfig {
    /// Hidden dimension (must match input/output for skip connection)
    pub hidden_dim: usize,
    /// Dropout probability (applied on hot path; cold path is identity)
    pub dropout: f64,
    /// `LayerNorm` epsilon for numerical stability
    pub layer_norm_eps: f64,
}

impl Default for ResidualConfig {
    fn default() -> Self {
        Self {
            hidden_dim: 128,
            dropout: 0.2,
            layer_norm_eps: 1e-5,
        }
    }
}

// Residual layernorm kernel is precompiled to cubin by build.rs.

/// Residual block with skip connection using native CUDA weight storage.
///
/// Architecture:
/// ```text
/// input --> fc1 --> GELU --> LayerNorm --> fc2 --> (+) --> GELU --> output
///   |                                                ^
///   +------------------------------------------------+
/// ```
///
/// All operations run on GPU -- zero host downloads in forward pass.
#[allow(missing_debug_implementations)]
pub struct ResidualBlock {
    fc1: OwnedGpuLinear,
    fc2: OwnedGpuLinear,
    activations: ActivationKernels,
    ln_kernel: CudaFunction,
    norm_weight: CudaSlice<f32>,
    norm_bias: CudaSlice<f32>,
    stream: Arc<CudaStream>,
    cublas: CudaBlas,
    normalized_shape: usize,
    eps: f32,
    _dropout: f64,
}

impl ResidualBlock {
    /// Create a new residual block with native CUDA weight storage.
    pub fn new(
        stream: Arc<CudaStream>,
        config: &ResidualConfig,
        name: &str,
    ) -> Result<Self, MLError> {
        // Two linear layers with same input/output dimension for skip connection
        let fc1 = OwnedGpuLinear::kaiming(&format!("{name}_fc1"), config.hidden_dim, config.hidden_dim, &stream)?;
        let fc2 = OwnedGpuLinear::kaiming(&format!("{name}_fc2"), config.hidden_dim, config.hidden_dim, &stream)?;

        // LayerNorm parameters: weight (ones) and bias (zeros)
        let ones = vec![1.0_f32; config.hidden_dim];
        let norm_weight = ml_core::cuda_autograd::init::upload_to_gpu(&ones, &stream)?;
        let norm_bias = ml_core::cuda_autograd::init::zeros(config.hidden_dim, &stream)?;

        // Pre-compile activation kernels
        let activations = ActivationKernels::new(&stream)?;

        // Load precompiled LayerNorm kernel from cubin
        static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/residual_kernels.cubin"));
        let module = stream.context()
            .load_cubin(CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("residual cubin load: {e}")))?;
        let ln_kernel = module.load_function("residual_layernorm_forward")
            .map_err(|e| MLError::ModelError(format!("residual_layernorm_forward: {e}")))?;

        let cublas = CudaBlas::new(Arc::clone(&stream))
            .map_err(|e| MLError::ModelError(format!("ResidualBlock cuBLAS init: {e}")))?;

        Ok(Self {
            fc1,
            fc2,
            activations,
            ln_kernel,
            norm_weight,
            norm_bias,
            stream,
            cublas,
            normalized_shape: config.hidden_dim,
            eps: config.layer_norm_eps as f32,
            _dropout: config.dropout,
        })
    }

    /// Forward pass through residual block -- GPU-only.
    ///
    /// All operations (linear, GELU, `LayerNorm`, residual add) execute on GPU
    /// with zero host downloads.
    pub fn forward(&self, x: &GpuTensor, _train: bool) -> Result<GpuTensor, MLError> {
        // fc1 forward via cuBLAS (GPU)
        let (h, _acts) = self.fc1.forward(x, &self.cublas, &self.stream)?;

        // GELU activation (GPU kernel)
        let (h_gelu, _saved) = self.activations.gelu_fwd(&h, &self.stream)?;

        // LayerNorm (GPU kernel)
        let h_norm = self.layernorm_forward(&h_gelu)?;

        // Dropout skipped on cold path (identity)

        // fc2 forward via cuBLAS (GPU)
        let (h_fc2, _acts2) = self.fc2.forward(&h_norm, &self.cublas, &self.stream)?;

        // Skip connection: residual add (GPU kernel)
        let h_residual = h_fc2.add(x, &self.stream)?;

        // Final GELU activation (GPU kernel)
        let (output, _saved2) = self.activations.gelu_fwd(&h_residual, &self.stream)?;

        Ok(output)
    }

    /// GPU-native `LayerNorm` via compiled CUDA kernel.
    fn layernorm_forward(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let dim = self.normalized_shape;
        let batch = x.numel() / dim;

        let y = GpuTensor::zeros(x.shape(), &self.stream)?;
        let batch_i32 = batch as i32;
        let features_i32 = dim as i32;

        let threads = 256_u32.min(dim as u32);
        let cfg = LaunchConfig {
            grid_dim: (batch as u32, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: (threads as usize * std::mem::size_of::<f32>() * 2) as u32,
        };

        // SAFETY: x, y, norm_weight, norm_bias are valid device allocations. Kernel indices bounds-checked.
        unsafe {
            self.stream
                .launch_builder(&self.ln_kernel)
                .arg(x.data())
                .arg(y.data())
                .arg(&self.norm_weight)
                .arg(&self.norm_bias)
                .arg(&batch_i32)
                .arg(&features_i32)
                .arg(&self.eps)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("residual layernorm_forward: {e}")))?;
        }

        Ok(y)
    }

}

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

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
            .expect("CUDA required")
            .new_stream()
            .expect("CUDA stream")
    }

    #[test]
    fn test_residual_config_default() {
        let config = ResidualConfig::default();
        assert_eq!(config.hidden_dim, 128);
        assert_eq!(config.dropout, 0.2);
        assert_eq!(config.layer_norm_eps, 1e-5);
    }

    #[test]
    fn test_residual_block_creation() -> anyhow::Result<()> {
        let config = ResidualConfig { hidden_dim: 64, dropout: 0.1, layer_norm_eps: 1e-5 };
        let block = ResidualBlock::new(make_stream(), &config, "test_block");
        assert!(block.is_ok());
        Ok(())
    }

    #[test]
    fn test_residual_block_forward_train() -> anyhow::Result<()> {
        let stream = make_stream();
        let config = ResidualConfig { hidden_dim: 32, dropout: 0.1, layer_norm_eps: 1e-5 };
        let block = ResidualBlock::new(Arc::clone(&stream), &config, "test_block")?;
        let input_data: Vec<f32> = (0..2 * 32).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&input_data, vec![2, 32], &stream)?;
        let output = block.forward(&input, true)?;
        assert_eq!(output.shape(), &[2, 32]);
        Ok(())
    }

    #[test]
    fn test_residual_block_forward_eval() -> anyhow::Result<()> {
        let stream = make_stream();
        let config = ResidualConfig { hidden_dim: 32, dropout: 0.1, layer_norm_eps: 1e-5 };
        let block = ResidualBlock::new(Arc::clone(&stream), &config, "test_block")?;
        let input_data: Vec<f32> = (0..2 * 32).map(|i| (i as f32 * 0.01).sin()).collect();
        let input = GpuTensor::from_host(&input_data, vec![2, 32], &stream)?;
        let output = block.forward(&input, false)?;
        assert_eq!(output.shape(), &[2, 32]);
        Ok(())
    }

    #[test]
    fn test_residual_skip_connection_identity() -> anyhow::Result<()> {
        let stream = make_stream();
        let config = ResidualConfig { hidden_dim: 16, dropout: 0.0, layer_norm_eps: 1e-5 };
        let block = ResidualBlock::new(Arc::clone(&stream), &config, "test_block")?;
        let input_data = vec![1.0_f32; 16];
        let input = GpuTensor::from_host(&input_data, vec![1, 16], &stream)?;
        let output = block.forward(&input, false)?;
        let output_host = output.to_host(&stream)?;
        assert!(output_host.iter().any(|&x| x != 0.0));
        Ok(())
    }

    #[test]
    fn test_residual_batch_processing() -> anyhow::Result<()> {
        let stream = make_stream();
        let config = ResidualConfig { hidden_dim: 64, dropout: 0.1, layer_norm_eps: 1e-5 };
        let block = ResidualBlock::new(Arc::clone(&stream), &config, "test_block")?;
        for batch_size in [1, 4, 8, 16] {
            let input_data: Vec<f32> = (0..batch_size * 64).map(|i| (i as f32 * 0.01).sin()).collect();
            let input = GpuTensor::from_host(&input_data, vec![batch_size, 64], &stream)?;
            let output = block.forward(&input, true)?;
            assert_eq!(output.shape(), &[batch_size, 64]);
        }
        Ok(())
    }

    #[test]
    fn test_residual_different_dimensions() -> anyhow::Result<()> {
        let stream = make_stream();
        for hidden_dim in [16, 32, 64, 128, 256] {
            let config = ResidualConfig { hidden_dim, dropout: 0.1, layer_norm_eps: 1e-5 };
            let block = ResidualBlock::new(
                Arc::clone(&stream), &config, &format!("block_{hidden_dim}"),
            )?;
            let input_data: Vec<f32> = (0..2 * hidden_dim).map(|i| (i as f32 * 0.01).sin()).collect();
            let input = GpuTensor::from_host(&input_data, vec![2, hidden_dim], &stream)?;
            let output = block.forward(&input, true)?;
            assert_eq!(output.shape(), &[2, hidden_dim]);
        }
        Ok(())
    }

    #[test]
    fn test_residual_numerical_stability() -> anyhow::Result<()> {
        let stream = make_stream();
        let config = ResidualConfig { hidden_dim: 32, dropout: 0.1, layer_norm_eps: 1e-5 };
        let block = ResidualBlock::new(Arc::clone(&stream), &config, "test_block")?;
        let input_data = vec![100.0_f32; 32];
        let input = GpuTensor::from_host(&input_data, vec![1, 32], &stream)?;
        let output = block.forward(&input, false)?;
        let output_host = output.to_host(&stream)?;
        assert!(output_host.iter().all(|&x| x.is_finite()));
        Ok(())
    }
}
