#![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 layer weights are stored in a [`GpuVarStore`] with [`GpuLinear`] layers
//! using cuBLAS sgemm. `LayerNorm` parameters are also in the `GpuVarStore`.
//! 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::nvrtc::Ptx;

use ml_core::cuda_autograd::{ActivationKernels, GpuLinear, GpuTensor, GpuVarStore};
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,
        }
    }
}

// ── LayerNorm CUDA kernel source (for residual block) ──────────────────

const LAYERNORM_CUDA_SRC: &str = r#"
extern "C" __global__
void residual_layernorm_forward(const float* __restrict__ x,
                                float* __restrict__ y,
                                const float* __restrict__ gamma,
                                const float* __restrict__ beta,
                                int batch,
                                int features,
                                float eps) {
    int row = blockIdx.x;
    if (row >= batch) return;

    const float* x_row = x + row * features;
    float* y_row = y + row * features;

    extern __shared__ float smem[];
    float* s_sum = smem;
    float* s_sq  = smem + blockDim.x;

    float local_sum = 0.0f;
    float local_sq  = 0.0f;
    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float val = x_row[j];
        local_sum += val;
        local_sq  += val * val;
    }

    s_sum[threadIdx.x] = local_sum;
    s_sq[threadIdx.x]  = local_sq;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            s_sum[threadIdx.x] += s_sum[threadIdx.x + stride];
            s_sq[threadIdx.x]  += s_sq[threadIdx.x + stride];
        }
        __syncthreads();
    }

    float mean = s_sum[0] / (float)features;
    float var  = s_sq[0] / (float)features - mean * mean;
    float inv_std = rsqrtf(var + eps);

    for (int j = threadIdx.x; j < features; j += blockDim.x) {
        float normalized = (x_row[j] - mean) * inv_std;
        y_row[j] = normalized * gamma[j] + beta[j];
    }
}
"#;

/// 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: GpuLinear,
    fc2: GpuLinear,
    activations: ActivationKernels,
    ln_kernel: CudaFunction,
    store: GpuVarStore,
    stream: Arc<CudaStream>,
    cublas: CudaBlas,
    norm_weight_name: String,
    norm_bias_name: String,
    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> {
        let mut store = GpuVarStore::new(Arc::clone(&stream));

        // Two linear layers with same input/output dimension for skip connection
        let fc1 = store.linear_kaiming(&format!("{name}_fc1"), config.hidden_dim, config.hidden_dim)?;
        let fc2 = store.linear_kaiming(&format!("{name}_fc2"), config.hidden_dim, config.hidden_dim)?;

        // LayerNorm parameters: weight (ones) and bias (zeros)
        let norm_weight_name = format!("{name}_norm_weight");
        let norm_bias_name = format!("{name}_norm_bias");
        let ones = vec![1.0_f32; config.hidden_dim];
        let w_data = ml_core::cuda_autograd::init::upload_to_gpu(&ones, &stream)?;
        store.register(norm_weight_name.clone(), w_data, vec![config.hidden_dim])?;
        let b_data = ml_core::cuda_autograd::init::zeros(config.hidden_dim, &stream)?;
        store.register(norm_bias_name.clone(), b_data, vec![config.hidden_dim])?;

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

        // Eagerly compile LayerNorm kernel
        let context = stream.context();
        let ptx: Ptx = ml_core::cuda_compile::compile_ptx_for_device(LAYERNORM_CUDA_SRC, context)
            .map_err(|e| MLError::ModelError(format!("residual layernorm compile: {e}")))?;
        let module = context.load_module(ptx)
            .map_err(|e| MLError::ModelError(format!("residual layernorm module: {e}")))?;
        let ln_kernel = module.load_function("residual_layernorm_forward")
            .map_err(|e| MLError::ModelError(format!("residual layernorm load: {e}")))?;

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

        Ok(Self {
            fc1,
            fc2,
            activations,
            ln_kernel,
            store,
            stream,
            cublas,
            norm_weight_name,
            norm_bias_name,
            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.store, &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.store, &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 w_param = self.store.get(&self.norm_weight_name).ok_or_else(|| {
            MLError::ModelError(format!("norm weight '{}' not found", self.norm_weight_name))
        })?;
        let b_param = self.store.get(&self.norm_bias_name).ok_or_else(|| {
            MLError::ModelError(format!("norm bias '{}' not found", self.norm_bias_name))
        })?;

        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, w_param, b_param are valid device allocations. Kernel indices bounds-checked.
        unsafe {
            self.stream
                .launch_builder(&self.ln_kernel)
                .arg(x.data())
                .arg(y.data())
                .arg(&w_param.data)
                .arg(&b_param.data)
                .arg(&batch_i32)
                .arg(&features_i32)
                .arg(&self.eps)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("residual layernorm_forward: {e}")))?;
        }

        Ok(y)
    }

    /// Reference to the underlying `GpuVarStore`.
    pub const fn store(&self) -> &GpuVarStore { &self.store }
}

#[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(())
    }
}
