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

//! GPU-accelerated IQL (Implicit Q-Learning) value network trainer.
//!
//! Implements the value network component of IQL (Kostrikov et al., 2021)
//! as a fused CUDA kernel pipeline:
//!
//! 1. **Forward + Expectile Loss** -- V(s) MLP forward + asymmetric loss
//! 2. **Gradient Norm** -- L2 norm for gradient clipping
//! 3. **Backward** -- backprop through SiLU MLP, atomicAdd gradients
//! 4. **Adam** -- AdamW update with gradient clipping
//!
//! ## Architecture
//!
//! V(s) is a 2-hidden-layer MLP with SiLU activation:
//!   state -> Linear(STATE_DIM, H) -> SiLU -> Linear(H, H) -> SiLU -> Linear(H, 1)
//!
//! ## Integration
//!
//! Called after the main DQN training step in `FusedTrainingCtx::run_full_step()`.
//! The value loss trains V(s) to approximate the expectile of Q(s,a), and the
//! advantage weights `exp(beta * (Q(s,a) - V(s)))` can modulate the DQN policy.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use tracing::info;

use crate::MLError;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

/// Configuration for the GPU IQL value trainer.
#[derive(Debug, Clone)]
pub struct GpuIqlConfig {
    /// Expectile parameter tau (0.5 = symmetric MSE, 0.7-0.9 = optimistic).
    pub expectile_tau: f32,
    /// Temperature beta for advantage-weighted policy extraction.
    pub advantage_temperature: f32,
    /// Hidden dimension for V(s) MLP layers.
    pub value_hidden_dim: usize,
    /// Input state dimension (tensor-core-aligned, e.g. 48 or 56).
    pub state_dim: usize,
    /// Training batch size (must match DQN trainer batch size).
    pub batch_size: usize,
    /// Learning rate for V network Adam optimizer.
    pub lr: f32,
    /// Adam beta1 (first moment decay).
    pub beta1: f32,
    /// Adam beta2 (second moment decay).
    pub beta2: f32,
    /// Adam epsilon (numerical stability).
    pub epsilon: f32,
    /// Decoupled weight decay (AdamW).
    pub weight_decay: f32,
    /// Maximum gradient L2 norm for clipping.
    pub max_grad_norm: f32,
}

impl Default for GpuIqlConfig {
    fn default() -> Self {
        Self {
            expectile_tau: 0.7,
            advantage_temperature: 3.0,
            value_hidden_dim: 128,
            state_dim: 72,
            batch_size: 256,
            lr: 3e-4,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
            weight_decay: 1e-5,
            max_grad_norm: 1.0,
        }
    }
}

impl GpuIqlConfig {
    /// Total number of parameters in the V(s) MLP.
    fn total_params(&self) -> usize {
        let h = self.value_hidden_dim;
        let sd = self.state_dim;
        // w1[H*SD] + b1[H] + w2[H*H] + b2[H] + w3[H] + b3[1]
        h * sd + h + h * h + h + h + 1
    }
}

// ---------------------------------------------------------------------------
// GpuIqlTrainer
// ---------------------------------------------------------------------------

/// GPU-accelerated IQL value network trainer.
///
/// Owns pre-allocated GPU buffers for the V(s) network weights, Adam state,
/// activation saves, and the compiled CUDA kernels. Operates entirely on GPU
/// with no per-step CPU-GPU data transfers (states and Q-values are already
/// on device from the DQN training step).
#[allow(missing_debug_implementations)]
pub struct GpuIqlTrainer {
    config: GpuIqlConfig,
    stream: Arc<CudaStream>,

    // ── Compiled kernels ────────────────────────────────────────────
    forward_loss_kernel: CudaFunction,
    backward_kernel: CudaFunction,
    grad_norm_kernel: CudaFunction,
    adam_kernel: CudaFunction,
    forward_kernel: CudaFunction,

    // ── V network parameters (flat f32 on GPU) ─────────────────────
    params_buf: CudaSlice<f32>,

    // ── Adam optimizer state ────────────────────────────────────────
    m_buf: CudaSlice<f32>,        // first moment
    v_buf: CudaSlice<f32>,        // second moment
    grad_buf: CudaSlice<f32>,     // gradient accumulator
    grad_norm_buf: CudaSlice<f32>, // [1] gradient L2 norm

    // ── Activation save buffers (forward -> backward) ───────────────
    save_pre1: CudaSlice<f32>,    // [B, H] pre-activation layer 1
    save_pre2: CudaSlice<f32>,    // [B, H] pre-activation layer 2
    save_h1: CudaSlice<f32>,      // [B, H] post-activation layer 1
    save_h2: CudaSlice<f32>,      // [B, H] post-activation layer 2

    // ── Output buffers ──────────────────────────────────────────────
    v_out_buf: CudaSlice<f32>,    // [B] V(s) predictions
    loss_buf: CudaSlice<f32>,     // [B] per-sample loss
    total_loss_buf: CudaSlice<f32>, // [1] batch-mean loss

    // ── Training state ──────────────────────────────────────────────
    adam_step: i32,
    total_params: usize,
}

impl GpuIqlTrainer {
    /// Create a new GPU IQL trainer.
    ///
    /// Compiles the 5 CUDA kernels (forward+loss, backward, grad_norm, adam,
    /// forward-only), initializes V(s) weights with Xavier/Glorot uniform,
    /// and pre-allocates all GPU buffers.
    pub fn new(
        stream: Arc<CudaStream>,
        config: GpuIqlConfig,
    ) -> Result<Self, MLError> {
        let total_params = config.total_params();
        let b = config.batch_size;
        let h = config.value_hidden_dim;

        // Compile all 5 kernels
        let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_kernel, forward_kernel) =
            compile_iql_kernels(&stream, &config)?;

        // Allocate parameter buffer and initialize with Xavier/Glorot
        let params_buf = init_xavier_weights(&stream, &config)?;

        // Allocate Adam state (zeroed)
        let m_buf = alloc_f32(&stream, total_params, "iql_m")?;
        let v_buf = alloc_f32(&stream, total_params, "iql_v")?;
        let grad_buf = alloc_f32(&stream, total_params, "iql_grad")?;
        let grad_norm_buf = alloc_f32(&stream, 1, "iql_grad_norm")?;

        // Allocate activation save buffers
        let save_pre1 = alloc_f32(&stream, b * h, "iql_save_pre1")?;
        let save_pre2 = alloc_f32(&stream, b * h, "iql_save_pre2")?;
        let save_h1 = alloc_f32(&stream, b * h, "iql_save_h1")?;
        let save_h2 = alloc_f32(&stream, b * h, "iql_save_h2")?;

        // Allocate output buffers
        let v_out_buf = alloc_f32(&stream, b, "iql_v_out")?;
        let loss_buf = alloc_f32(&stream, b, "iql_loss")?;
        let total_loss_buf = alloc_f32(&stream, 1, "iql_total_loss")?;

        let vram_bytes = (total_params * 4       // params + m + v + grad = 4 copies
            + b * h * 4                          // 4 activation buffers
            + b * 2 + 2                          // v_out + loss + total_loss + grad_norm
        ) * 4;

        info!(
            state_dim = config.state_dim,
            hidden_dim = config.value_hidden_dim,
            batch_size = config.batch_size,
            total_params,
            expectile_tau = config.expectile_tau,
            advantage_temperature = config.advantage_temperature,
            vram_kb = vram_bytes / 1024,
            "GpuIqlTrainer initialized: 5 kernels compiled, V(s) weights Xavier-initialized"
        );

        Ok(Self {
            config,
            stream,
            forward_loss_kernel,
            backward_kernel,
            grad_norm_kernel,
            adam_kernel,
            forward_kernel,
            params_buf,
            m_buf,
            v_buf,
            grad_buf,
            grad_norm_buf,
            save_pre1,
            save_pre2,
            save_h1,
            save_h2,
            v_out_buf,
            loss_buf,
            total_loss_buf,
            adam_step: 0,
            total_params,
        })
    }

    /// Run one IQL value network training step.
    ///
    /// Executes the full forward + expectile loss + backward + Adam cycle:
    /// 1. Zero total_loss and grad_norm scalars
    /// 2. Forward + loss kernel (one warp per sample)
    /// 3. Backward kernel (one warp per sample, atomicAdd gradients)
    /// 4. Grad norm kernel (parallel reduction)
    /// 5. Adam update kernel (one thread per parameter)
    ///
    /// Returns the batch-mean expectile loss (read back from GPU).
    ///
    /// # Arguments
    ///
    /// * `states_f32` - F32 CudaSlice `[B * STATE_DIM]` on CUDA device
    /// * `q_values_f32` - F32 CudaSlice `[B]` on CUDA device (target Q-values from DQN)
    ///
    /// Takes pre-converted F32 CudaSlice buffers directly — zero Candle
    /// tensor involvement.
    pub fn train_value_step(
        &mut self,
        states_f32: &CudaSlice<f32>,
        q_values_f32: &CudaSlice<f32>,
    ) -> Result<f32, MLError> {
        let b = self.config.batch_size;

        // Zero total_loss and grad_norm before kernel launches
        let zero_f32 = [0.0_f32];
        self.stream.memcpy_htod(&zero_f32, &mut self.total_loss_buf)
            .map_err(|e| MLError::ModelError(format!("IQL zero total_loss: {e}")))?;
        self.stream.memcpy_htod(&zero_f32, &mut self.grad_norm_buf)
            .map_err(|e| MLError::ModelError(format!("IQL zero grad_norm: {e}")))?;

        let batch_size_i32 = b as i32;
        let total_params_i32 = self.total_params as i32;

        // 1. Forward + loss kernel
        let fwd_config = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: 0,
        };

        // Safety: states_f32 and q_values_f32 are valid F32 CudaSlice buffers
        // on the same CUDA context. All internal buffers are pre-allocated.
        unsafe {
            self.stream
                .launch_builder(&self.forward_loss_kernel)
                .arg(states_f32)
                .arg(q_values_f32)
                .arg(&self.params_buf)
                .arg(&mut self.v_out_buf)
                .arg(&mut self.loss_buf)
                .arg(&mut self.total_loss_buf)
                .arg(&mut self.save_pre1)
                .arg(&mut self.save_pre2)
                .arg(&mut self.save_h1)
                .arg(&mut self.save_h2)
                .arg(&batch_size_i32)
                .launch(fwd_config)
                .map_err(|e| MLError::ModelError(format!("IQL forward+loss kernel: {e}")))?;
        }

        // 2. Backward kernel
        let bwd_config = LaunchConfig {
            grid_dim: (b as u32, 1, 1),
            block_dim: (32, 1, 1),
            shared_mem_bytes: 0,
        };

        // Safety: all pointers are valid GPU allocations on the same stream.
        unsafe {
            self.stream
                .launch_builder(&self.backward_kernel)
                .arg(states_f32)
                .arg(q_values_f32)
                .arg(&self.v_out_buf)
                .arg(&self.params_buf)
                .arg(&self.save_pre1)
                .arg(&self.save_pre2)
                .arg(&self.save_h1)
                .arg(&self.save_h2)
                .arg(&mut self.grad_buf)
                .arg(&batch_size_i32)
                .launch(bwd_config)
                .map_err(|e| MLError::ModelError(format!("IQL backward kernel: {e}")))?;
        }

        // 3. Grad norm kernel
        let norm_blocks = (self.total_params + 255) / 256;
        let norm_config = LaunchConfig {
            grid_dim: (norm_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        // Safety: grad_buf and grad_norm_buf are valid allocations.
        unsafe {
            self.stream
                .launch_builder(&self.grad_norm_kernel)
                .arg(&self.grad_buf)
                .arg(&mut self.grad_norm_buf)
                .arg(&total_params_i32)
                .launch(norm_config)
                .map_err(|e| MLError::ModelError(format!("IQL grad_norm kernel: {e}")))?;
        }

        // 4. Adam update kernel
        self.adam_step += 1;
        let adam_blocks = (self.total_params + 255) / 256;
        let adam_config = LaunchConfig {
            grid_dim: (adam_blocks as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        let lr = self.config.lr;
        let beta1 = self.config.beta1;
        let beta2 = self.config.beta2;
        let eps = self.config.epsilon;
        let wd = self.config.weight_decay;
        let mgn = self.config.max_grad_norm;
        let t = self.adam_step;

        // Safety: all parameter and optimizer buffers are valid allocations.
        unsafe {
            self.stream
                .launch_builder(&self.adam_kernel)
                .arg(&mut self.params_buf)
                .arg(&mut self.grad_buf)
                .arg(&mut self.m_buf)
                .arg(&mut self.v_buf)
                .arg(&self.grad_norm_buf)
                .arg(&lr)
                .arg(&beta1)
                .arg(&beta2)
                .arg(&eps)
                .arg(&wd)
                .arg(&mgn)
                .arg(&t)
                .arg(&total_params_i32)
                .launch(adam_config)
                .map_err(|e| MLError::ModelError(format!("IQL adam kernel: {e}")))?;
        }

        // Read back total loss
        let mut loss_host = [0.0_f32];
        self.stream
            .memcpy_dtoh(&self.total_loss_buf, &mut loss_host) // gpu-exit: 4B scalar loss
            .map_err(|e| MLError::ModelError(format!("IQL loss readback: {e}")))?;

        Ok(loss_host[0])
    }

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

// ---------------------------------------------------------------------------
// Kernel compilation
// ---------------------------------------------------------------------------

/// Compile all 5 IQL CUDA kernels with dimension defines.
fn compile_iql_kernels(
    stream: &Arc<CudaStream>,
    config: &GpuIqlConfig,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
    let dim_overrides = format!(
        "#define STATE_DIM {state_dim}\n\
         #define VALUE_HIDDEN_DIM {hidden_dim}\n\
         #define IQL_EXPECTILE_TAU {tau:.6}f\n",
        state_dim = config.state_dim,
        hidden_dim = config.value_hidden_dim,
        tau = config.expectile_tau,
    );

    let kernel_src = include_str!("iql_value_kernel.cu");
    let full_source = format!("{dim_overrides}\n{kernel_src}");

    info!(
        state_dim = config.state_dim,
        hidden_dim = config.value_hidden_dim,
        expectile_tau = config.expectile_tau,
        total_params = config.total_params(),
        "GpuIqlTrainer: compiling 5 IQL kernels"
    );

    let context = stream.context();
    let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| {
            MLError::ModelError(format!("iql_value_kernel compilation failed: {e}"))
        })?;
    let module = context.load_module(ptx).map_err(|e| {
        MLError::ModelError(format!("iql_value module load: {e}"))
    })?;

    let forward_loss = module
        .load_function("iql_forward_loss_kernel")
        .map_err(|e| MLError::ModelError(format!("iql_forward_loss_kernel load: {e}")))?;
    let backward = module
        .load_function("iql_backward_kernel")
        .map_err(|e| MLError::ModelError(format!("iql_backward_kernel load: {e}")))?;
    let grad_norm = module
        .load_function("iql_grad_norm_kernel")
        .map_err(|e| MLError::ModelError(format!("iql_grad_norm_kernel load: {e}")))?;
    let adam = module
        .load_function("iql_adam_kernel")
        .map_err(|e| MLError::ModelError(format!("iql_adam_kernel load: {e}")))?;
    let forward = module
        .load_function("iql_forward_kernel")
        .map_err(|e| MLError::ModelError(format!("iql_forward_kernel load: {e}")))?;

    info!("GpuIqlTrainer: 5 kernels compiled and loaded");
    Ok((forward_loss, backward, grad_norm, adam, forward))
}

// ---------------------------------------------------------------------------
// Weight initialization
// ---------------------------------------------------------------------------

/// Initialize V(s) network weights with Xavier/Glorot uniform distribution.
///
/// For a layer with fan_in and fan_out, weights are sampled from
/// U(-sqrt(6/(fan_in+fan_out)), sqrt(6/(fan_in+fan_out))).
/// Biases are initialized to zero.
fn init_xavier_weights(
    stream: &Arc<CudaStream>,
    config: &GpuIqlConfig,
) -> Result<CudaSlice<f32>, MLError> {
    use rand::Rng;

    let total = config.total_params();
    let h = config.value_hidden_dim;
    let sd = config.state_dim;
    let mut weights = vec![0.0_f32; total];
    let mut rng = rand::thread_rng();

    // Layer 1: w1[H, SD], b1[H]
    let limit1 = (6.0_f64 / (sd + h) as f64).sqrt() as f32;
    let w1_end = h * sd;
    for w in &mut weights[..w1_end] {
        *w = rng.gen_range(-limit1..limit1);
    }
    // b1 stays zero (w1_end .. w1_end + h)

    // Layer 2: w2[H, H], b2[H]
    let w2_start = h * sd + h;
    let limit2 = (6.0_f64 / (h + h) as f64).sqrt() as f32;
    let w2_end = w2_start + h * h;
    for w in &mut weights[w2_start..w2_end] {
        *w = rng.gen_range(-limit2..limit2);
    }
    // b2 stays zero (w2_end .. w2_end + h)

    // Output layer: w3[H], b3[1]
    let w3_start = w2_end + h;
    let limit3 = (6.0_f64 / (h + 1) as f64).sqrt() as f32;
    let w3_end = w3_start + h;
    for w in &mut weights[w3_start..w3_end] {
        *w = rng.gen_range(-limit3..limit3);
    }
    // b3 stays zero

    // Upload to GPU
    let mut params_buf = alloc_f32(stream, total, "iql_params")?;
    stream
        .memcpy_htod(&weights, &mut params_buf)
        .map_err(|e| MLError::ModelError(format!("IQL weight upload: {e}")))?;

    Ok(params_buf)
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Allocate a zeroed f32 buffer on GPU.
fn alloc_f32(
    stream: &Arc<CudaStream>,
    n: usize,
    name: &str,
) -> Result<CudaSlice<f32>, MLError> {
    stream.alloc_zeros::<f32>(n).map_err(|e| {
        MLError::ModelError(format!("GpuIql alloc {name} ({n} f32): {e}"))
    })
}

