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

//! GPU Multi-Head Feature Attention for DQN state preprocessing.
//!
//! Applies 4-head self-attention over the state feature dimension
//! before feeding into the DQN trunk. This teaches the model which
//! features are most relevant for the current market state.
//!
//! Architecture: state -> MultiHead_Attn(state) -> attended_state -> DQN trunk
//! Residual connection ensures gradient flow even with frozen attention weights.
//!
//! Phase B: Full backward pass with gradient flow through attention weights.
//! The backward kernel computes gradients for all attention parameters
//! (W_Q, W_K, W_V, W_O, biases, LayerNorm gamma/beta) and propagates
//! gradients to the input. Attention weights are updated via a dedicated
//! Adam optimizer (separate from the DQN trunk Adam).

use std::sync::{Arc, OnceLock};
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use tracing::info;
use crate::MLError;

static ATTN_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
static ATTN_BWD_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();

/// Configuration for the GPU attention layer.
#[derive(Debug, Clone)]
pub struct GpuAttentionConfig {
    pub state_dim: usize,
    pub num_heads: usize,
    pub batch_size: usize,
}

impl Default for GpuAttentionConfig {
    fn default() -> Self {
        Self {
            state_dim: 72,
            num_heads: 4,
            batch_size: 256,
        }
    }
}

impl GpuAttentionConfig {
    /// Total number of attention parameters.
    /// Layout: 3*D*D (W_QKV) + 3*D (b_QKV) + D*D (W_O) + D (b_O) + D (gamma) + D (beta)
    pub fn total_params(&self) -> usize {
        let d = self.state_dim;
        4 * d * d + 5 * d
    }
}

/// GPU multi-head feature attention layer with full backward pass.
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct GpuAttention {
    config: GpuAttentionConfig,
    stream: Arc<CudaStream>,
    forward_kernel: CudaFunction,
    backward_kernel: CudaFunction,
    grad_norm_kernel: CudaFunction,
    adam_kernel: CudaFunction,
    /// Attention parameters (Xavier-initialized, trainable).
    params: CudaSlice<f32>,
    /// Output buffer for attended states [batch_size, state_dim].
    output_buf: CudaSlice<f32>,
    /// Saved input states from forward pass (for backward recomputation).
    saved_input: CudaSlice<f32>,
    /// Scratch buffer for input gradient [batch_size, state_dim].
    /// Written by backward kernel via atomicAdd. Not propagated further
    /// (attention is the first layer — no upstream to backprop into).
    d_input_scratch: CudaSlice<f32>,
    /// Gradient accumulator for attention parameters [total_params].
    d_params: CudaSlice<f32>,
    /// Gradient norm output buffer [1].
    grad_norm_buf: CudaSlice<f32>,
    /// Adam first moment (m) for attention parameters [total_params].
    attn_m: CudaSlice<f32>,
    /// Adam second moment (v) for attention parameters [total_params].
    attn_v: CudaSlice<f32>,
    /// Adam step counter (1-indexed).
    attn_adam_step: i32,
    /// Total attention parameters count.
    total_params: usize,
}

impl GpuAttention {
    pub fn new(
        stream: Arc<CudaStream>,
        config: GpuAttentionConfig,
    ) -> Result<Self, MLError> {
        let context = stream.context();
        let d = config.state_dim;
        let total_params = config.total_params();
        let b = config.batch_size;

        // Compile forward attention kernel
        let defines = format!(
            "#define ATTN_NUM_HEADS {}\n#define ATTN_STATE_DIM {}\n",
            config.num_heads, d
        );
        let kernel_src = include_str!("attention_kernel.cu");
        let full_source = format!("{defines}{kernel_src}");
        let ptx = ATTN_PTX.get_or_init(|| {
            crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        });
        let ptx = ptx.as_ref().map_err(|e| MLError::ModelError(format!("attention PTX: {e}")))?;
        let module = context.load_module(ptx.clone())
            .map_err(|e| MLError::ModelError(format!("attention module: {e}")))?;
        let forward_kernel = module.load_function("multihead_feature_attention")
            .map_err(|e| MLError::ModelError(format!("attention kernel load: {e}")))?;

        // Compile backward attention kernel (separate PTX for independent caching)
        let bwd_kernel_src = include_str!("attention_backward_kernel.cu");
        let bwd_full_source = format!("{defines}{bwd_kernel_src}");
        let bwd_ptx = ATTN_BWD_PTX.get_or_init(|| {
            crate::cuda_pipeline::compile_ptx_for_device(&bwd_full_source, &context)
        });
        let bwd_ptx = bwd_ptx.as_ref().map_err(|e| MLError::ModelError(format!("attention backward PTX: {e}")))?;
        let bwd_module = context.load_module(bwd_ptx.clone())
            .map_err(|e| MLError::ModelError(format!("attention backward module: {e}")))?;
        let backward_kernel = bwd_module.load_function("attention_backward_kernel")
            .map_err(|e| MLError::ModelError(format!("attention backward kernel load: {e}")))?;
        let grad_norm_kernel = bwd_module.load_function("attn_grad_norm_kernel")
            .map_err(|e| MLError::ModelError(format!("attention grad_norm kernel load: {e}")))?;
        let adam_kernel = bwd_module.load_function("attn_adam_kernel")
            .map_err(|e| MLError::ModelError(format!("attention adam kernel load: {e}")))?;

        // Xavier initialization for attention weights
        let mut host_params = vec![0.0_f32; total_params];
        let fan_in = d as f32;
        let fan_out = d as f32;
        let xavier_std = (2.0 / (fan_in + fan_out)).sqrt();
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        let mut hasher = DefaultHasher::new();
        42u64.hash(&mut hasher);
        let mut rng_state = hasher.finish();

        // Initialize W_Q, W_K, W_V, W_O with Xavier
        let weight_count = 4 * d * d;
        for i in 0..weight_count {
            // Simple LCG PRNG for initialization
            rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            let u = (rng_state >> 33) as f32 / (1u64 << 31) as f32; // [0, 1)
            // Box-Muller approximation
            let normal = (u - 0.5) * 3.46; // ~ uniform -> normal approximation
            host_params[i] = normal as f32 * xavier_std;
        }

        // Biases = 0, LayerNorm gamma = 1, beta = 0
        // Layout: 3*D*D (W_QKV) + 3*D (b_QKV) + D*D (W_O) + D (b_O) + D (gamma) + D (beta)
        // weight_count = 4*D*D covers W_QKV + W_O. Biases = 3*D + D = 4*D.
        let ln_gamma_start = 4 * d * d + 4 * d; // after all weights + biases
        for i in 0..d {
            host_params[ln_gamma_start + i] = 1.0; // gamma = 1
        }
        // beta stays 0

        let mut params = stream.alloc_zeros::<f32>(total_params)
            .map_err(|e| MLError::ModelError(format!("attention params alloc: {e}")))?;
        stream.memcpy_htod(&host_params, &mut params)
            .map_err(|e| MLError::ModelError(format!("attention params upload: {e}")))?;

        let output_buf = stream.alloc_zeros::<f32>(b * d)
            .map_err(|e| MLError::ModelError(format!("attention output alloc: {e}")))?;

        // Saved state buffer for backward pass (input states)
        let saved_input = stream.alloc_zeros::<f32>(b * d)
            .map_err(|e| MLError::ModelError(format!("attention saved_input alloc: {e}")))?;

        // Gradient and optimizer buffers
        let d_input_scratch = stream.alloc_zeros::<f32>(b * d)
            .map_err(|e| MLError::ModelError(format!("attention d_input_scratch alloc: {e}")))?;
        let d_params = stream.alloc_zeros::<f32>(total_params)
            .map_err(|e| MLError::ModelError(format!("attention d_params alloc: {e}")))?;
        let grad_norm_buf = stream.alloc_zeros::<f32>(1)
            .map_err(|e| MLError::ModelError(format!("attention grad_norm alloc: {e}")))?;
        let attn_m = stream.alloc_zeros::<f32>(total_params)
            .map_err(|e| MLError::ModelError(format!("attention adam_m alloc: {e}")))?;
        let attn_v = stream.alloc_zeros::<f32>(total_params)
            .map_err(|e| MLError::ModelError(format!("attention adam_v alloc: {e}")))?;

        // VRAM: params(4) + d_params + m + v + output + saved_input + d_input_scratch + grad_norm
        let vram_bytes = (total_params * 4 + b * d * 3 + 1) * 4;
        let vram_kb = vram_bytes / 1024;
        info!(
            state_dim = d,
            num_heads = config.num_heads,
            head_dim = d / config.num_heads,
            total_params,
            vram_kb,
            "GpuAttention initialized: multi-head feature attention with backward pass"
        );

        Ok(Self {
            config,
            stream,
            forward_kernel,
            backward_kernel,
            grad_norm_kernel,
            adam_kernel,
            params,
            output_buf,
            saved_input,
            d_input_scratch,
            d_params,
            grad_norm_buf,
            attn_m,
            attn_v,
            attn_adam_step: 0,
            total_params,
        })
    }

    /// Apply multi-head attention to batch of states.
    ///
    /// Saves input states and pre-layernorm output for the backward pass.
    /// Returns a reference to the output buffer (attended states).
    pub fn forward(&mut self, states: &CudaSlice<f32>, batch_size: usize) -> Result<&CudaSlice<f32>, MLError> {
        let b = batch_size;
        let d = self.config.state_dim;

        // Save input states for backward pass (DtoD copy)
        let n_bytes = b * d * std::mem::size_of::<f32>();
        let (src_ptr, _sg) = states.device_ptr(&self.stream);
        let src_ptr_val = src_ptr;
        let _sg = std::mem::ManuallyDrop::new(_sg);
        let (dst_ptr, _dg) = self.saved_input.device_ptr(&self.stream);
        let dst_ptr_val = dst_ptr;
        let _dg = std::mem::ManuallyDrop::new(_dg);
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr_val, src_ptr_val, n_bytes, self.stream.cu_stream()
            ).map_err(|e| MLError::ModelError(format!("attention save input DtoD: {e}")))?;
        }

        let launch_cfg = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: (d * 4) as u32, // shared_concat
        };
        let b_i32 = b as i32;

        unsafe {
            self.stream
                .launch_builder(&self.forward_kernel)
                .arg(states)
                .arg(&self.params)
                .arg(&mut self.output_buf)
                .arg(&b_i32)
                .launch(launch_cfg)
                .map_err(|e| MLError::ModelError(format!("attention forward: {e}")))?;
        }

        // The backward kernel recomputes the full forward pass from saved_input
        // + params (concat, pre_ln). No intermediate activations need saving
        // beyond the input states (already copied above).

        Ok(&self.output_buf)
    }

    /// Full backward pass: computes gradients for all attention parameters.
    ///
    /// The backward kernel recomputes the full forward pass from saved_input + params
    /// to recover intermediate activations (concat, pre_ln). This avoids storing
    /// large intermediate buffers and is acceptable since attention runs once per step.
    ///
    /// `d_output` is the gradient w.r.t. the attended states [B, D] from the
    /// DQN trunk backward pass (bw_d_h_s2).
    ///
    /// After this call, `d_params` contains accumulated weight gradients.
    /// Call `adam_step()` after this to update the attention weights.
    /// Input gradients are written to an internal scratch buffer (not propagated
    /// further since attention is the first layer).
    pub fn backward(
        &mut self,
        d_output: &CudaSlice<f32>,
        batch_size: usize,
    ) -> Result<(), MLError> {
        let b = batch_size;
        let d = self.config.state_dim;
        let f32_size = std::mem::size_of::<f32>();

        // Zero the gradient accumulator before backward
        let (dp_ptr, _dpg) = self.d_params.device_ptr(&self.stream);
        let dp_ptr_val = dp_ptr;
        let _dpg = std::mem::ManuallyDrop::new(_dpg);
        unsafe {
            cudarc::driver::result::memset_d8_async(
                dp_ptr_val, 0u8, self.total_params * f32_size, self.stream.cu_stream()
            ).map_err(|e| MLError::ModelError(format!("attention d_params zero: {e}")))?;
        }

        // Zero d_input_scratch (backward kernel uses atomicAdd for input gradients)
        let (di_ptr, _dig) = self.d_input_scratch.device_ptr(&self.stream);
        let di_ptr_val = di_ptr;
        let _dig = std::mem::ManuallyDrop::new(_dig);
        unsafe {
            cudarc::driver::result::memset_d8_async(
                di_ptr_val, 0u8, b * d * f32_size, self.stream.cu_stream()
            ).map_err(|e| MLError::ModelError(format!("attention d_input zero: {e}")))?;
        }

        // Launch backward kernel: one warp per sample
        // Shared memory: concat[D] + d_proj[D] = 2*D*sizeof(float)
        let shared_mem = (d * 4 * 2) as u32;
        let launch_cfg = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: shared_mem,
        };
        let b_i32 = b as i32;

        unsafe {
            self.stream
                .launch_builder(&self.backward_kernel)
                .arg(d_output)
                .arg(&self.saved_input)
                .arg(&self.params)
                .arg(&mut self.d_input_scratch)
                .arg(&mut self.d_params)
                .arg(&b_i32)
                .launch(launch_cfg)
                .map_err(|e| MLError::ModelError(format!("attention backward: {e}")))?;
        }

        Ok(())
    }

    /// Run gradient norm + Adam update on attention parameters.
    ///
    /// Must be called after `backward()`. Uses the same grad_norm + Adam
    /// kernel pattern as IQN (separate from DQN trunk Adam).
    pub fn adam_step(&mut self, lr: f32, max_grad_norm: f32) -> Result<(), MLError> {
        let tp = self.total_params;
        let tp_i32 = tp as i32;

        // 1. Zero grad_norm_buf
        let (gn_ptr, _gng) = self.grad_norm_buf.device_ptr(&self.stream);
        let gn_ptr_val = gn_ptr;
        let _gng = std::mem::ManuallyDrop::new(_gng);
        unsafe {
            cudarc::driver::result::memset_d8_async(
                gn_ptr_val, 0u8, std::mem::size_of::<f32>(), self.stream.cu_stream()
            ).map_err(|e| MLError::ModelError(format!("attention grad_norm zero: {e}")))?;
        }

        // 2. Gradient norm kernel
        let norm_blocks = (tp + 255) / 256;
        let norm_cfg = LaunchConfig {
            grid_dim: (norm_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_kernel)
                .arg(&self.d_params)
                .arg(&mut self.grad_norm_buf)
                .arg(&tp_i32)
                .launch(norm_cfg)
                .map_err(|e| MLError::ModelError(format!("attention grad_norm kernel: {e}")))?;
        }

        // 3. Adam update kernel
        self.attn_adam_step += 1;
        let adam_blocks = (tp + 255) / 256;
        let adam_cfg = LaunchConfig {
            grid_dim: (adam_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let beta1 = 0.9_f32;
        let beta2 = 0.999_f32;
        let eps = 1e-8_f32;
        let wd = 1e-5_f32;
        let t = self.attn_adam_step;

        unsafe {
            self.stream
                .launch_builder(&self.adam_kernel)
                .arg(&mut self.params)
                .arg(&mut self.d_params)
                .arg(&mut self.attn_m)
                .arg(&mut self.attn_v)
                .arg(&self.grad_norm_buf)
                .arg(&lr)
                .arg(&beta1)
                .arg(&beta2)
                .arg(&eps)
                .arg(&wd)
                .arg(&max_grad_norm)
                .arg(&t)
                .arg(&tp_i32)
                .launch(adam_cfg)
                .map_err(|e| MLError::ModelError(format!("attention adam kernel: {e}")))?;
        }

        Ok(())
    }

    /// Current Adam step count.
    pub fn adam_step_count(&self) -> i32 {
        self.attn_adam_step
    }

    /// Get the output buffer reference.
    pub fn output(&self) -> &CudaSlice<f32> {
        &self.output_buf
    }
}
