#![allow(unsafe_code)]

//! Decision Transformer for offline RL pre-training.
//!
//! Architecture (Chen et al., 2021 "Decision Transformer"):
//!   Input sequence: [(R_1, s_1, a_1), (R_2, s_2, a_2), ..., (R_T, s_T, a_T)]
//!   where R_t = return-to-go at timestep t
//!
//! At inference: set R_1 = target_return (e.g., Sharpe=1.0)
//!   -> model predicts actions that achieve that return
//!
//! For trading:
//!   - Pre-train on historical walk-forward windows (offline data)
//!   - Condition on desired Sharpe ratio as return-to-go
//!   - Fine-tune with DQN online training (hybrid)
//!
//! ## GPU Implementation
//!
//! Uses custom CUDA kernels compiled via NVRTC for the full transformer
//! forward + backward + Adam pipeline. All operations are GPU-resident
//! with zero CPU in the hot path.
//!
//! The transformer layers are:
//!   1. Token embedding: [state_dim + 2] -> [embed_dim] per timestep
//!      (state features + return-to-go + action index)
//!   2. Positional encoding: learned [context_len, embed_dim]
//!   3. N transformer blocks, each with:
//!      a. Multi-head causal self-attention (H heads)
//!      b. LayerNorm + residual
//!      c. FFN: embed -> 4*embed -> embed
//!      d. LayerNorm + residual
//!   4. Action prediction head: [embed_dim] -> [num_actions]
//!
//! ## Integration with DQN
//!
//! The Decision Transformer can be used in two modes:
//! 1. **Pre-training**: train on offline data with return-to-go conditioning
//! 2. **Fine-tuning**: use as a state encoder for the DQN trunk
//!    - Replace the MLP trunk with transformer encoder
//!    - DQN heads (C51, IQN) attach to the transformer output
//!    - Train end-to-end with the DQN loss

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

/// Configuration for the Decision Transformer.
#[derive(Debug, Clone)]
pub struct DecisionTransformerConfig {
    /// State dimension (input features per timestep).
    pub state_dim: usize,
    /// Number of exposure actions (9 for branching DQN).
    pub num_actions: usize,
    /// Embedding dimension for transformer.
    pub embed_dim: usize,
    /// Number of transformer layers.
    pub num_layers: usize,
    /// Number of attention heads per layer.
    pub num_heads: usize,
    /// Context length (number of past timesteps to attend to).
    pub context_len: usize,
    /// Dropout probability (0.0 = disabled for inference).
    pub dropout: f32,
    /// Maximum batch size.
    pub batch_size: usize,
}

impl Default for DecisionTransformerConfig {
    fn default() -> Self {
        Self {
            state_dim: 72,
            num_actions: 9,
            embed_dim: 128,
            num_layers: 3,
            num_heads: 4,
            context_len: 20, // attend to last 20 timesteps
            dropout: 0.1,
            batch_size: 256,
        }
    }
}

impl DecisionTransformerConfig {
    /// Input dimension per timestep: state_dim + 2 (return-to-go + action).
    fn input_dim(&self) -> usize {
        self.state_dim + 2
    }

    /// Total parameter count for the transformer.
    pub fn total_params(&self) -> usize {
        let d = self.embed_dim;
        let s = self.state_dim;

        // Token embedding: (state_dim + 2) x embed_dim + embed_dim (bias)
        let token_embed = (s + 2) * d + d;

        // Positional encoding: context_len x embed_dim
        let pos_embed = self.context_len * d;

        // Per transformer layer:
        //   Multi-head attention: 4 x d x d + 4 x d (Q, K, V, O + biases)
        //   FFN: d x 4d + 4d + 4d x d + d (two linear layers + biases)
        //   LayerNorm: 2 x (d + d) = 4d (two LN layers, gamma + beta each)
        let per_layer = 4 * d * d + 4 * d  // attention
                      + d * 4 * d + 4 * d + 4 * d * d + d  // FFN
                      + 4 * d;  // layer norms

        // Action prediction head: embed_dim x num_actions + num_actions
        let action_head = d * self.num_actions + self.num_actions;

        token_embed + pos_embed + self.num_layers * per_layer + action_head
    }

    /// VRAM estimate in bytes.
    pub fn vram_estimate(&self) -> usize {
        let params = self.total_params() * 4; // f32
        let activations = self.batch_size * self.context_len * self.embed_dim * 4 * 2; // save for backward
        params + activations
    }
}

// ── Parameter layout offsets ────────────────────────────────────────────
//
// The flat parameter buffer is laid out as:
//   [token_embed] [pos_embed] [layer_0] [layer_1] ... [layer_N-1] [action_head]
//
// Each layer contains:
//   W_Q[E,E] W_K[E,E] W_V[E,E] W_O[E,E] b_Q[E] b_K[E] b_V[E] b_O[E]
//   W_ffn1[E,4E] b_ffn1[4E] W_ffn2[4E,E] b_ffn2[E]
//   ln1_gamma[E] ln1_beta[E] ln2_gamma[E] ln2_beta[E]

/// Compute the element offset for each parameter group in the flat buffer.
struct ParamLayout {
    w_embed_off: usize,
    b_embed_off: usize,
    pos_embed_off: usize,
    layers_off: usize,
    layer_size: usize,
    w_head_off: usize,
    b_head_off: usize,
    embed_dim: usize,
}

impl ParamLayout {
    fn new(config: &DecisionTransformerConfig) -> Self {
        let d = config.embed_dim;
        let s = config.state_dim;
        let input_dim = s + 2;

        let w_embed_off = 0;
        let b_embed_off = w_embed_off + input_dim * d;
        let pos_embed_off = b_embed_off + d;
        let layers_off = pos_embed_off + config.context_len * d;

        let layer_size = 4 * d * d + 4 * d
                       + d * 4 * d + 4 * d + 4 * d * d + d
                       + 4 * d;

        let w_head_off = layers_off + config.num_layers * layer_size;
        let b_head_off = w_head_off + d * config.num_actions;

        Self {
            w_embed_off,
            b_embed_off,
            pos_embed_off,
            layers_off,
            layer_size,
            w_head_off,
            b_head_off,
            embed_dim: d,
        }
    }

    /// Get element offsets within a single layer (relative to layer start).
    fn layer_offsets(&self) -> LayerOffsets {
        let d = self.embed_dim;
        let w_q = 0;
        let w_k = w_q + d * d;
        let w_v = w_k + d * d;
        let w_o = w_v + d * d;
        let b_q = w_o + d * d;
        let b_k = b_q + d;
        let b_v = b_k + d;
        let b_o = b_v + d;
        let w_ffn1 = b_o + d;
        let b_ffn1 = w_ffn1 + d * 4 * d;
        let w_ffn2 = b_ffn1 + 4 * d;
        let b_ffn2 = w_ffn2 + 4 * d * d;
        let ln1_gamma = b_ffn2 + d;
        let ln1_beta = ln1_gamma + d;
        let ln2_gamma = ln1_beta + d;
        let ln2_beta = ln2_gamma + d;

        LayerOffsets {
            w_q, w_k, w_v, w_o,
            b_q, b_k, b_v, b_o,
            w_ffn1, b_ffn1, w_ffn2, b_ffn2,
            ln1_gamma, ln1_beta, ln2_gamma, ln2_beta,
        }
    }
}

#[allow(dead_code)]
struct LayerOffsets {
    w_q: usize, w_k: usize, w_v: usize, w_o: usize,
    b_q: usize, b_k: usize, b_v: usize, b_o: usize,
    w_ffn1: usize, b_ffn1: usize, w_ffn2: usize, b_ffn2: usize,
    ln1_gamma: usize, ln1_beta: usize, ln2_gamma: usize, ln2_beta: usize,
}

// ── Raw device pointer extraction ──────────────────────────────────────
//
// Same pattern as `raw_device_ptr` in gpu_dqn_trainer.rs:
// Extract the u64 device pointer and prevent the SyncOnDrop guard from
// recording an event. Safe because all operations happen on the same stream.

fn raw_ptr_f32(slice: &CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

fn raw_ptr_f32_mut(slice: &mut CudaSlice<f32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr_mut(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

fn raw_ptr_i32(slice: &CudaSlice<i32>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

/// Byte offset from an element offset (f32 = 4 bytes).
fn byte_off(elem_off: usize) -> u64 {
    (elem_off * 4) as u64
}

// ── Compiled DT kernels (lazy init) ────────────────────────────────────

struct DtKernels {
    embed: CudaFunction,
    qkv_projection: CudaFunction,
    causal_attention: CudaFunction,
    layernorm: CudaFunction,
    ffn: CudaFunction,
    action_head: CudaFunction,
    cross_entropy: CudaFunction,
    ce_backward: CudaFunction,
    linear_backward: CudaFunction,
    zero: CudaFunction,
    residual_add: CudaFunction,
    return_to_go: CudaFunction,
    build_trajectories: CudaFunction,
    compute_rewards_actions: CudaFunction,
}

fn compile_dt_kernels(
    stream: &Arc<CudaStream>,
    config: &DecisionTransformerConfig,
) -> Result<DtKernels, MLError> {
    let defines = format!(
        "#define DT_EMBED_DIM {}\n\
         #define DT_NUM_HEADS {}\n\
         #define DT_CONTEXT_LEN {}\n\
         #define DT_INPUT_DIM {}\n\
         #define DT_NUM_ACTIONS {}\n",
        config.embed_dim,
        config.num_heads,
        config.context_len,
        config.input_dim(),
        config.num_actions,
    );

    let kernel_src = include_str!("dt_kernels.cu");
    let full_source = format!("{defines}{kernel_src}");

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

    let load = |name: &str| -> Result<CudaFunction, MLError> {
        module.load_function(name)
            .map_err(|e| MLError::ModelError(format!("dt kernel {name} load: {e}")))
    };

    Ok(DtKernels {
        embed: load("dt_embed_kernel")?,
        qkv_projection: load("dt_qkv_projection_kernel")?,
        causal_attention: load("dt_causal_attention_kernel")?,
        layernorm: load("dt_layernorm_kernel")?,
        ffn: load("dt_ffn_kernel")?,
        action_head: load("dt_action_head_kernel")?,
        cross_entropy: load("dt_cross_entropy_kernel")?,
        ce_backward: load("dt_ce_backward_kernel")?,
        linear_backward: load("dt_linear_backward_kernel")?,
        zero: load("dt_zero_kernel")?,
        residual_add: load("dt_residual_add_kernel")?,
        return_to_go: load("dt_return_to_go_kernel")?,
        build_trajectories: load("dt_build_trajectories_kernel")?,
        compute_rewards_actions: load("dt_compute_rewards_actions_kernel")?,
    })
}

// ── GPU scratch buffers for DT forward/backward ─────────────────────

struct DtScratch {
    /// Embedded tokens: [B, T, E]
    embedded: CudaSlice<f32>,
    /// Q, K, V projections: [B, T, E] each
    q_buf: CudaSlice<f32>,
    k_buf: CudaSlice<f32>,
    v_buf: CudaSlice<f32>,
    /// Attention output (pre-LN): [B, T, E]
    attn_out: CudaSlice<f32>,
    /// Post-LN1 output: [B, T, E]
    ln1_out: CudaSlice<f32>,
    /// FFN output (post-LN2): [B, T, E]
    ffn_out: CudaSlice<f32>,
    /// LN2 temp output: [B, T, E]
    ln2_out: CudaSlice<f32>,
    /// Action logits: [B*T, A]
    logits: CudaSlice<f32>,
    /// Per-sample loss: [B*T]
    per_sample_loss: CudaSlice<f32>,
    /// Scalar total loss: [1]
    total_loss: CudaSlice<f32>,
    /// Gradient of logits: [B*T, A]
    d_logits: CudaSlice<f32>,
    /// Gradient buffer for layer backward: [B, T, E]
    d_layer: CudaSlice<f32>,
    /// Second gradient buffer for embed backward input grad: [B, T, input_dim]
    d_embed_input: CudaSlice<f32>,
    /// Gradient accumulator for all parameters: [total_params]
    d_params: CudaSlice<f32>,
    /// Adam first moment: [total_params]
    adam_m: CudaSlice<f32>,
    /// Adam second moment: [total_params]
    adam_v: CudaSlice<f32>,
    /// Grad norm scratch: [1]
    grad_norm: CudaSlice<f32>,
    /// Adam step counter on device: [1]
    adam_t: CudaSlice<i32>,
}

fn alloc_dt_scratch(
    stream: &Arc<CudaStream>,
    config: &DecisionTransformerConfig,
) -> Result<DtScratch, MLError> {
    let bte = config.batch_size * config.context_len * config.embed_dim;
    let bt_a = config.batch_size * config.context_len * config.num_actions;
    let bt = config.batch_size * config.context_len;
    let bt_input = config.batch_size * config.context_len * config.input_dim();
    let total_params = config.total_params();

    let alloc_f = |n: usize, name: &str| -> Result<CudaSlice<f32>, MLError> {
        stream.alloc_zeros::<f32>(n)
            .map_err(|e| MLError::ModelError(format!("DT scratch {name} ({n} f32): {e}")))
    };

    Ok(DtScratch {
        embedded: alloc_f(bte, "embedded")?,
        q_buf: alloc_f(bte, "q_buf")?,
        k_buf: alloc_f(bte, "k_buf")?,
        v_buf: alloc_f(bte, "v_buf")?,
        attn_out: alloc_f(bte, "attn_out")?,
        ln1_out: alloc_f(bte, "ln1_out")?,
        ffn_out: alloc_f(bte, "ffn_out")?,
        ln2_out: alloc_f(bte, "ln2_out")?,
        logits: alloc_f(bt_a, "logits")?,
        per_sample_loss: alloc_f(bt, "per_sample_loss")?,
        total_loss: alloc_f(1, "total_loss")?,
        d_logits: alloc_f(bt_a, "d_logits")?,
        d_layer: alloc_f(bte, "d_layer")?,
        d_embed_input: alloc_f(bt_input, "d_embed_input")?,
        d_params: alloc_f(total_params, "d_params")?,
        adam_m: alloc_f(total_params, "adam_m")?,
        adam_v: alloc_f(total_params, "adam_v")?,
        grad_norm: alloc_f(1, "grad_norm")?,
        adam_t: stream.alloc_zeros::<i32>(1)
            .map_err(|e| MLError::ModelError(format!("DT scratch adam_t: {e}")))?,
    })
}

// ── Adam kernels (reused from DQN utility kernels) ─────────────────────

struct AdamKernels {
    grad_norm: CudaFunction,
    adam_update: CudaFunction,
}

fn compile_adam_kernels(stream: &Arc<CudaStream>, state_dim: usize) -> Result<AdamKernels, MLError> {
    let dim_overrides = format!(
        "#define STATE_DIM {state_dim}\n\
         #define MARKET_DIM 42\n\
         #define PORTFOLIO_DIM 8\n",
    );
    let common_src = include_str!("common_device_functions.cuh");
    let kernel_src = include_str!("dqn_utility_kernels.cu");
    let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}");

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
        .map_err(|e| MLError::ModelError(format!("DT adam kernels compile: {e}")))?;
    let module = context.load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("DT adam module load: {e}")))?;

    let grad_norm = module.load_function("dqn_grad_norm_kernel")
        .map_err(|e| MLError::ModelError(format!("DT grad_norm load: {e}")))?;
    let adam_update = module.load_function("dqn_adam_update_kernel")
        .map_err(|e| MLError::ModelError(format!("DT adam_update load: {e}")))?;

    Ok(AdamKernels { grad_norm, adam_update })
}

/// Decision Transformer model.
///
/// Pre-trains on offline trading data, conditions on return-to-go.
/// Can be used standalone or as a state encoder for the DQN trunk.
#[allow(missing_debug_implementations)] // CudaSlice does not implement Debug
pub struct DecisionTransformer {
    config: DecisionTransformerConfig,
    stream: Arc<CudaStream>,
    /// Flat parameter buffer (all weights + biases).
    params: CudaSlice<f32>,
    /// Context buffer: [batch_size, context_len, state_dim + 2].
    context_buf: CudaSlice<f32>,
    /// Output buffer: [batch_size, num_actions] action logits.
    action_logits: CudaSlice<f32>,
    /// Compiled DT kernels (lazy init on first pretrain_step).
    kernels: Option<DtKernels>,
    /// Adam optimizer kernels (grad_norm + adam_update).
    adam_kernels: Option<AdamKernels>,
    /// GPU scratch buffers for forward/backward.
    scratch: Option<DtScratch>,
    /// Parameter layout offsets.
    layout: ParamLayout,
    /// Adam step counter (host-side, incremented and uploaded each step).
    adam_step: i32,
}

impl DecisionTransformer {
    /// Create a new Decision Transformer with Xavier-initialized weights.
    pub fn new(
        stream: Arc<CudaStream>,
        config: DecisionTransformerConfig,
    ) -> Result<Self, MLError> {
        let total_params = config.total_params();
        let context_size = config.batch_size * config.context_len * config.input_dim();
        let output_size = config.batch_size * config.num_actions;
        let layout = ParamLayout::new(&config);

        // Xavier initialization with proper LN gamma=1, beta=0
        let mut host_params = vec![0.0_f32; total_params];
        let xavier_std = (2.0 / (config.embed_dim as f32 + config.embed_dim as f32)).sqrt();
        let mut rng = 0x1234_5678_u64;

        let lo = layout.layer_offsets();
        for (i, p) in host_params.iter_mut().enumerate() {
            rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
            let u = (rng >> 33) as f32 / (1_u64 << 31) as f32;
            *p = (u - 0.5) * 3.46 * xavier_std;

            // Override LN gamma/beta for each layer
            for layer_idx in 0..config.num_layers {
                let lb = layout.layers_off + layer_idx * layout.layer_size;
                if (lb + lo.ln1_gamma..lb + lo.ln1_gamma + config.embed_dim).contains(&i) {
                    *p = 1.0;
                } else if (lb + lo.ln1_beta..lb + lo.ln1_beta + config.embed_dim).contains(&i) {
                    *p = 0.0;
                } else if (lb + lo.ln2_gamma..lb + lo.ln2_gamma + config.embed_dim).contains(&i) {
                    *p = 1.0;
                } else if (lb + lo.ln2_beta..lb + lo.ln2_beta + config.embed_dim).contains(&i) {
                    *p = 0.0;
                } else {
                    // Keep Xavier-initialized value for all other parameters
                }
            }
        }

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

        let context_buf = stream.alloc_zeros::<f32>(context_size)
            .map_err(|e| MLError::ModelError(format!("DT context alloc: {e}")))?;
        let action_logits = stream.alloc_zeros::<f32>(output_size)
            .map_err(|e| MLError::ModelError(format!("DT action_logits alloc: {e}")))?;

        let vram_mb = config.vram_estimate() as f64 / (1024.0 * 1024.0);
        info!(
            embed_dim = config.embed_dim,
            num_layers = config.num_layers,
            num_heads = config.num_heads,
            context_len = config.context_len,
            total_params,
            vram_mb = format!("{vram_mb:.1}"),
            "DecisionTransformer initialized"
        );

        Ok(Self {
            layout,
            config,
            stream,
            params,
            context_buf,
            action_logits,
            kernels: None,
            adam_kernels: None,
            scratch: None,
            adam_step: 0,
        })
    }

    /// Lazy-init kernels and scratch buffers on first call.
    fn ensure_compiled(&mut self) -> Result<(), MLError> {
        if self.kernels.is_none() {
            info!("DecisionTransformer: compiling CUDA kernels (first call)");
            self.kernels = Some(compile_dt_kernels(&self.stream, &self.config)?);
            self.adam_kernels = Some(compile_adam_kernels(&self.stream, self.config.state_dim)?);
            self.scratch = Some(alloc_dt_scratch(&self.stream, &self.config)?);
            info!("DecisionTransformer: kernels compiled and scratch allocated");
        }
        Ok(())
    }

    /// Pre-train on a batch of offline trajectories.
    ///
    /// # Arguments
    /// * `trajectories` - GPU buffer containing `[..., batch_size, context_len, input_dim]`
    /// * `target_actions` - GPU buffer containing `[..., batch_size * context_len]` action indices
    /// * `batch_size` - Number of trajectories in this batch
    /// * `traj_elem_offset` - Element offset into trajectories buffer (f32 elements, not bytes)
    /// * `act_elem_offset` - Element offset into target_actions buffer (i32 elements, not bytes)
    ///
    /// # Returns
    /// Mean cross-entropy loss for this batch.
    pub fn pretrain_step(
        &mut self,
        trajectories: &CudaSlice<f32>,
        target_actions: &CudaSlice<i32>,
        batch_size: usize,
        traj_elem_offset: usize,
        act_elem_offset: usize,
    ) -> Result<f32, MLError> {
        self.ensure_compiled()?;

        let b = batch_size;
        let t = self.config.context_len;
        let e = self.config.embed_dim;
        let a = self.config.num_actions;
        let input_dim = self.config.input_dim();
        let bt = (b * t) as i32;
        let e_i32 = e as i32;
        let a_i32 = a as i32;
        let input_dim_i32 = input_dim as i32;
        let b_i32 = b as i32;
        let t_i32 = t as i32;
        let num_heads_i32 = self.config.num_heads as i32;
        let total_params = self.config.total_params();
        let total_params_i32 = total_params as i32;
        let num_layers = self.config.num_layers;

        let embed_block = e.min(256) as u32;
        let bt_grid = bt as u32;

        // Get raw pointers for all buffers upfront to avoid borrow issues.
        // All operations happen on self.stream — same stream, no sync needed.
        let stream = &self.stream;
        let kernels = self.kernels.as_ref()
            .ok_or_else(|| MLError::ModelError("DT kernels not compiled".to_string()))?;
        let adam = self.adam_kernels.as_ref()
            .ok_or_else(|| MLError::ModelError("Adam kernels not compiled".to_string()))?;
        let scratch = self.scratch.as_mut()
            .ok_or_else(|| MLError::ModelError("DT scratch not allocated".to_string()))?;

        let traj_ptr = raw_ptr_f32(trajectories, stream) + byte_off(traj_elem_offset);
        let target_ptr = raw_ptr_i32(target_actions, stream) + (act_elem_offset * 4) as u64;
        let params_ptr = raw_ptr_f32(&self.params, stream);

        let emb_ptr = raw_ptr_f32_mut(&mut scratch.embedded, stream);
        let q_ptr = raw_ptr_f32_mut(&mut scratch.q_buf, stream);
        let k_ptr = raw_ptr_f32_mut(&mut scratch.k_buf, stream);
        let v_ptr = raw_ptr_f32_mut(&mut scratch.v_buf, stream);
        let attn_ptr = raw_ptr_f32_mut(&mut scratch.attn_out, stream);
        let ln1_ptr = raw_ptr_f32_mut(&mut scratch.ln1_out, stream);
        let ffn_ptr = raw_ptr_f32_mut(&mut scratch.ffn_out, stream);
        let ln2_ptr = raw_ptr_f32_mut(&mut scratch.ln2_out, stream);
        let logits_ptr = raw_ptr_f32_mut(&mut scratch.logits, stream);
        let psl_ptr = raw_ptr_f32_mut(&mut scratch.per_sample_loss, stream);
        let total_loss_ptr = raw_ptr_f32_mut(&mut scratch.total_loss, stream);
        let d_logits_ptr = raw_ptr_f32_mut(&mut scratch.d_logits, stream);
        let d_layer_ptr = raw_ptr_f32_mut(&mut scratch.d_layer, stream);
        let d_embed_ptr = raw_ptr_f32_mut(&mut scratch.d_embed_input, stream);
        let d_params_ptr = raw_ptr_f32_mut(&mut scratch.d_params, stream);
        let grad_norm_ptr = raw_ptr_f32_mut(&mut scratch.grad_norm, stream);

        let lo = self.layout.layer_offsets();

        // ── FORWARD PASS ───────────────────────────────────────────────

        // 1. Token embedding + positional encoding
        {
            let w_e = params_ptr + byte_off(self.layout.w_embed_off);
            let b_e = params_ptr + byte_off(self.layout.b_embed_off);
            let p_e = params_ptr + byte_off(self.layout.pos_embed_off);
            let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (embed_block, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                stream.launch_builder(&kernels.embed)
                    .arg(&traj_ptr).arg(&w_e).arg(&b_e).arg(&p_e).arg(&emb_ptr)
                    .arg(&b_i32).arg(&t_i32).arg(&input_dim_i32).arg(&e_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT embed: {e}")))?;
            }
        }

        // 2-3. Transformer layers
        for layer_idx in 0..num_layers {
            let lb = self.layout.layers_off + layer_idx * self.layout.layer_size;

            // Layer input: embedded for layer 0, ln2_out for subsequent layers
            let layer_in = if layer_idx == 0 { emb_ptr } else { ln2_ptr };

            // QKV projection
            {
                let w_q = params_ptr + byte_off(lb + lo.w_q);
                let w_k = params_ptr + byte_off(lb + lo.w_k);
                let w_v = params_ptr + byte_off(lb + lo.w_v);
                let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (embed_block, 1, 1), shared_mem_bytes: 0 };
                unsafe {
                    stream.launch_builder(&kernels.qkv_projection)
                        .arg(&layer_in).arg(&w_q).arg(&w_k).arg(&w_v)
                        .arg(&q_ptr).arg(&k_ptr).arg(&v_ptr)
                        .arg(&b_i32).arg(&t_i32).arg(&e_i32)
                        .launch(lc)
                        .map_err(|e| MLError::ModelError(format!("DT QKV L{layer_idx}: {e}")))?;
                }
            }

            // Causal attention with output projection
            {
                let w_o = params_ptr + byte_off(lb + lo.w_o);
                let b_o = params_ptr + byte_off(lb + lo.b_o);
                let dh = e / self.config.num_heads;
                let shmem = (t * dh * 2 * 4) as u32; // K + V in shmem
                let lc = LaunchConfig {
                    grid_dim: (b as u32, self.config.num_heads as u32, 1),
                    block_dim: (t as u32, 1, 1),
                    shared_mem_bytes: shmem,
                };
                unsafe {
                    stream.launch_builder(&kernels.causal_attention)
                        .arg(&q_ptr).arg(&k_ptr).arg(&v_ptr).arg(&w_o).arg(&b_o)
                        .arg(&attn_ptr)
                        .arg(&b_i32).arg(&t_i32).arg(&e_i32).arg(&num_heads_i32)
                        .launch(lc)
                        .map_err(|e| MLError::ModelError(format!("DT attn L{layer_idx}: {e}")))?;
                }
            }

            // Residual add: attn_out += layer_input
            {
                let n_elems = (b * t * e) as i32;
                let lc = LaunchConfig {
                    grid_dim: (((n_elems as u32) + 255) / 256, 1, 1),
                    block_dim: (256, 1, 1), shared_mem_bytes: 0,
                };
                unsafe {
                    stream.launch_builder(&kernels.residual_add)
                        .arg(&attn_ptr).arg(&layer_in).arg(&n_elems)
                        .launch(lc)
                        .map_err(|e| MLError::ModelError(format!("DT res1 L{layer_idx}: {e}")))?;
                }
            }

            // LayerNorm 1: attn_out -> ln1_out
            {
                let g1 = params_ptr + byte_off(lb + lo.ln1_gamma);
                let b1 = params_ptr + byte_off(lb + lo.ln1_beta);
                let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (embed_block, 1, 1), shared_mem_bytes: 0 };
                unsafe {
                    stream.launch_builder(&kernels.layernorm)
                        .arg(&attn_ptr).arg(&g1).arg(&b1).arg(&ln1_ptr)
                        .arg(&bt).arg(&e_i32)
                        .launch(lc)
                        .map_err(|e| MLError::ModelError(format!("DT LN1 L{layer_idx}: {e}")))?;
                }
            }

            // FFN: ln1_out -> ffn_out (with residual from ln1_out)
            {
                let w1 = params_ptr + byte_off(lb + lo.w_ffn1);
                let b1 = params_ptr + byte_off(lb + lo.b_ffn1);
                let w2 = params_ptr + byte_off(lb + lo.w_ffn2);
                let b2 = params_ptr + byte_off(lb + lo.b_ffn2);
                let ffn_shmem = (e * 4 * 4) as u32; // [4E] floats hidden
                let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (embed_block, 1, 1), shared_mem_bytes: ffn_shmem };
                unsafe {
                    stream.launch_builder(&kernels.ffn)
                        .arg(&ln1_ptr).arg(&w1).arg(&b1).arg(&w2).arg(&b2).arg(&ffn_ptr)
                        .arg(&bt).arg(&e_i32)
                        .launch(lc)
                        .map_err(|e| MLError::ModelError(format!("DT FFN L{layer_idx}: {e}")))?;
                }
            }

            // LayerNorm 2: ffn_out -> ln2_out
            {
                let g2 = params_ptr + byte_off(lb + lo.ln2_gamma);
                let b2 = params_ptr + byte_off(lb + lo.ln2_beta);
                let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (embed_block, 1, 1), shared_mem_bytes: 0 };
                unsafe {
                    stream.launch_builder(&kernels.layernorm)
                        .arg(&ffn_ptr).arg(&g2).arg(&b2).arg(&ln2_ptr)
                        .arg(&bt).arg(&e_i32)
                        .launch(lc)
                        .map_err(|e| MLError::ModelError(format!("DT LN2 L{layer_idx}: {e}")))?;
                }
            }
        }

        // 4. Action head: ln2_out -> logits
        {
            let w_h = params_ptr + byte_off(self.layout.w_head_off);
            let b_h = params_ptr + byte_off(self.layout.b_head_off);
            let action_block = a.min(256).max(1) as u32;
            let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (action_block, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                stream.launch_builder(&kernels.action_head)
                    .arg(&ln2_ptr).arg(&w_h).arg(&b_h).arg(&logits_ptr)
                    .arg(&bt).arg(&e_i32).arg(&a_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT action head: {e}")))?;
            }
        }

        // 5. Cross-entropy loss
        {
            // Zero total_loss
            let one = 1_i32;
            let lc_z = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                stream.launch_builder(&kernels.zero)
                    .arg(&total_loss_ptr).arg(&one)
                    .launch(lc_z)
                    .map_err(|e| MLError::ModelError(format!("DT zero loss: {e}")))?;
            }

            let action_block = a.min(256).max(1) as u32;
            let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (action_block, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                stream.launch_builder(&kernels.cross_entropy)
                    .arg(&logits_ptr).arg(&target_ptr)
                    .arg(&psl_ptr).arg(&total_loss_ptr)
                    .arg(&bt).arg(&a_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT CE loss: {e}")))?;
            }
        }

        // ── BACKWARD PASS (Phase 1 — simplified) ──────────────────────

        // Zero gradient accumulator
        {
            let lc = LaunchConfig {
                grid_dim: (((total_params as u32) + 255) / 256, 1, 1),
                block_dim: (256, 1, 1), shared_mem_bytes: 0,
            };
            unsafe {
                stream.launch_builder(&kernels.zero)
                    .arg(&d_params_ptr).arg(&total_params_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT zero d_params: {e}")))?;
            }
        }

        // CE backward: d_logits = softmax(logits) - one_hot(target)
        {
            let action_block = a.min(256).max(1) as u32;
            let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (action_block, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                stream.launch_builder(&kernels.ce_backward)
                    .arg(&logits_ptr).arg(&target_ptr).arg(&d_logits_ptr)
                    .arg(&bt).arg(&a_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT CE bw: {e}")))?;
            }
        }

        // Action head backward: d_logits -> d_layer
        {
            let w_h = params_ptr + byte_off(self.layout.w_head_off);
            let dw_h = d_params_ptr + byte_off(self.layout.w_head_off);
            let db_h = d_params_ptr + byte_off(self.layout.b_head_off);
            let io_max = e.max(a).min(256) as u32;
            let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (io_max, 1, 1), shared_mem_bytes: 0 };

            // Input to action head was ln2_out (or embedded if 0 layers)
            let head_input = if num_layers > 0 { ln2_ptr } else { emb_ptr };

            unsafe {
                stream.launch_builder(&kernels.linear_backward)
                    .arg(&d_logits_ptr)  // d_output [N, A]
                    .arg(&head_input)    // input [N, E]
                    .arg(&w_h)           // W [E, A]
                    .arg(&dw_h)          // dW
                    .arg(&db_h)          // db
                    .arg(&d_layer_ptr)   // d_input [N, E]
                    .arg(&bt).arg(&e_i32).arg(&a_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT head bw: {e}")))?;
            }
        }

        // Transformer layers backward: Phase 1 = residual passthrough.
        // d_layer already contains the gradient from the action head backward.
        // The gradient passes through unchanged via residual connections.
        // No parameter gradients for transformer layers in Phase 1.

        // Embed backward: d_layer -> dW_embed, db_embed
        {
            let w_e = params_ptr + byte_off(self.layout.w_embed_off);
            let dw_e = d_params_ptr + byte_off(self.layout.w_embed_off);
            let db_e = d_params_ptr + byte_off(self.layout.b_embed_off);
            let io_max = e.max(input_dim).min(256) as u32;
            let lc = LaunchConfig { grid_dim: (bt_grid, 1, 1), block_dim: (io_max, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                stream.launch_builder(&kernels.linear_backward)
                    .arg(&d_layer_ptr)   // d_output [N, E]
                    .arg(&traj_ptr)      // input [N, input_dim]
                    .arg(&w_e)           // W [input_dim, E]
                    .arg(&dw_e)          // dW
                    .arg(&db_e)          // db
                    .arg(&d_embed_ptr)   // d_input [N, input_dim] (unused but kernel writes it)
                    .arg(&bt).arg(&input_dim_i32).arg(&e_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT embed bw: {e}")))?;
            }
        }

        // ── OPTIMIZER ──────────────────────────────────────────────────

        // Increment Adam step
        self.adam_step += 1;
        let step_host = [self.adam_step];
        stream.memcpy_htod(&step_host, &mut scratch.adam_t)
            .map_err(|e| MLError::ModelError(format!("DT adam_t upload: {e}")))?;

        let adam_t_ptr = raw_ptr_i32(&scratch.adam_t, stream);

        // Zero grad_norm
        {
            let one = 1_i32;
            let lc = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), shared_mem_bytes: 0 };
            unsafe {
                stream.launch_builder(&kernels.zero)
                    .arg(&grad_norm_ptr).arg(&one)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT zero gnorm: {e}")))?;
            }
        }

        // Grad norm reduction
        {
            let lc = LaunchConfig {
                grid_dim: (((total_params as u32) + 255) / 256, 1, 1),
                block_dim: (256, 1, 1), shared_mem_bytes: 0,
            };
            unsafe {
                stream.launch_builder(&adam.grad_norm)
                    .arg(&d_params_ptr).arg(&grad_norm_ptr).arg(&total_params_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT grad norm: {e}")))?;
            }
        }

        // Adam update
        {
            let lr = 3e-4_f32;
            let beta1 = 0.9_f32;
            let beta2 = 0.999_f32;
            let epsilon = 1e-8_f32;
            let weight_decay = 1e-5_f32;
            let max_grad_norm = 1.0_f32;

            let params_ptr_mut = raw_ptr_f32_mut(&mut self.params, stream);
            let adam_m_ptr = raw_ptr_f32_mut(&mut scratch.adam_m, stream);
            let adam_v_ptr = raw_ptr_f32_mut(&mut scratch.adam_v, stream);

            let lc = LaunchConfig {
                grid_dim: (((total_params as u32) + 255) / 256, 1, 1),
                block_dim: (256, 1, 1), shared_mem_bytes: 0,
            };

            unsafe {
                stream.launch_builder(&adam.adam_update)
                    .arg(&params_ptr_mut)
                    .arg(&d_params_ptr)
                    .arg(&adam_m_ptr)
                    .arg(&adam_v_ptr)
                    .arg(&grad_norm_ptr)
                    .arg(&lr).arg(&beta1).arg(&beta2).arg(&epsilon)
                    .arg(&weight_decay).arg(&max_grad_norm)
                    .arg(&adam_t_ptr)
                    .arg(&total_params_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT adam: {e}")))?;
            }
        }

        // ── Read back loss ─────────────────────────────────────────────
        let mut loss_host = [0.0_f32];
        stream.memcpy_dtoh(&scratch.total_loss, &mut loss_host)
            .map_err(|e| MLError::ModelError(format!("DT loss readback: {e}")))?;
        stream.synchronize()
            .map_err(|e| MLError::ModelError(format!("DT sync: {e}")))?;

        Ok(loss_host[0])
    }

    /// Get the configuration.
    pub fn config(&self) -> &DecisionTransformerConfig {
        &self.config
    }

    /// Get the action logits buffer.
    pub fn action_logits(&self) -> &CudaSlice<f32> {
        &self.action_logits
    }

    /// Build trajectory batches from GPU-resident features and targets.
    ///
    /// Performs all trajectory construction on GPU with zero CPU in the hot path:
    /// 1. Compute per-bar rewards and expert actions via GPU kernel
    /// 2. Compute return-to-go via reverse cumulative sum GPU kernel
    /// 3. Pack `[rtg, state, action]` tokens into trajectory tensor via GPU kernel
    ///
    /// # Arguments
    /// * `features_gpu` - GPU buffer `[num_bars, feat_dim]` (42-dim feature vectors)
    /// * `targets_gpu` - GPU buffer `[num_bars, 4]` (OHLC per bar)
    /// * `num_bars` - Number of bars in the dataset
    /// * `gamma` - Discount factor for return-to-go
    ///
    /// # Returns
    /// `(trajectories, target_actions, num_complete_batches)` where:
    /// - `trajectories` is `[num_episodes, T, input_dim]` on GPU
    /// - `target_actions` is `[num_episodes * T]` on GPU (i32)
    /// - `num_complete_batches` is how many full batches of `batch_size` episodes are available
    pub fn build_dt_trajectories(
        &mut self,
        features_gpu: &CudaSlice<f32>,
        targets_gpu: &CudaSlice<f32>,
        num_bars: usize,
        gamma: f64,
    ) -> Result<(CudaSlice<f32>, CudaSlice<i32>, usize), MLError> {
        self.ensure_compiled()?;

        let stream = &self.stream;
        let kernels = self.kernels.as_ref()
            .ok_or_else(|| MLError::ModelError("DT kernels not compiled".to_string()))?;

        let context_len = self.config.context_len;
        let feat_dim = 42_usize; // FeatureVector is [f64; 42]
        let input_dim = self.config.input_dim(); // feat_dim + 2
        let batch_size = self.config.batch_size;

        // Sliding window episodes with stride = context_len / 2 (50% overlap)
        let stride = context_len.max(2) / 2;
        if num_bars < context_len {
            return Err(MLError::ModelError(format!(
                "DT trajectory build: not enough bars ({num_bars}) for context_len ({context_len})"
            )));
        }
        let num_episodes = (num_bars - context_len) / stride + 1;
        if num_episodes == 0 {
            return Err(MLError::ModelError(
                "DT trajectory build: zero episodes from training data".to_string()
            ));
        }

        let num_complete_batches = num_episodes / batch_size;

        info!(
            num_bars,
            context_len,
            stride,
            num_episodes,
            num_complete_batches,
            batch_size,
            "DT trajectory builder: computing episodes"
        );

        // ── Step 1: Compute rewards and expert actions on GPU ──
        let mut rewards_gpu = stream.alloc_zeros::<f32>(num_bars)
            .map_err(|e| MLError::ModelError(format!("DT rewards alloc: {e}")))?;
        let mut actions_all_gpu = stream.alloc_zeros::<i32>(num_bars)
            .map_err(|e| MLError::ModelError(format!("DT actions alloc: {e}")))?;

        let targets_ptr = raw_ptr_f32(targets_gpu, stream);
        let rewards_ptr = raw_ptr_f32_mut(&mut rewards_gpu, stream);
        let actions_all_ptr = {
            let (ptr, guard) = actions_all_gpu.device_ptr_mut(stream);
            let _no_drop = std::mem::ManuallyDrop::new(guard);
            ptr
        };

        let num_bars_i32 = num_bars as i32;
        let action_threshold = 0.0001_f32; // 1 bps threshold for Long/Short

        {
            let lc = LaunchConfig {
                grid_dim: (((num_bars as u32) + 255) / 256, 1, 1),
                block_dim: (256, 1, 1),
                shared_mem_bytes: 0,
            };
            unsafe {
                stream.launch_builder(&kernels.compute_rewards_actions)
                    .arg(&targets_ptr)
                    .arg(&rewards_ptr)
                    .arg(&actions_all_ptr)
                    .arg(&num_bars_i32)
                    .arg(&action_threshold)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT compute_rewards_actions: {e}")))?;
            }
        }

        // ── Step 2: Build episode start indices on CPU (tiny: num_episodes ints) ──
        let episode_starts: Vec<i32> = (0..num_episodes)
            .map(|i| (i * stride) as i32)
            .collect();

        let episode_starts_gpu = stream.clone_htod(&episode_starts)
            .map_err(|e| MLError::ModelError(format!("DT episode_starts upload: {e}")))?;

        // ── Step 3: Compute per-episode return-to-go on GPU ──
        // We need rewards reshaped as [num_episodes, context_len].
        // Instead of copying, we'll compute RTG per-episode using the global rewards array.
        // Build a temporary [num_episodes, context_len] reward slice and RTG.
        let ep_total = num_episodes * context_len;
        let mut ep_rewards_gpu = stream.alloc_zeros::<f32>(ep_total)
            .map_err(|e| MLError::ModelError(format!("DT ep_rewards alloc: {e}")))?;
        let mut rtg_gpu = stream.alloc_zeros::<f32>(ep_total)
            .map_err(|e| MLError::ModelError(format!("DT rtg alloc: {e}")))?;

        // Copy per-episode reward windows from the global reward array.
        // Use a simple gather kernel (reuse build_trajectories_kernel pattern).
        // For simplicity, do this via a small CPU-side H2D of episode reward windows.
        // This is NOT in the hot path (one-time pre-training data build).
        //
        // Actually, we can avoid the extra copy by folding RTG into build_trajectories_kernel.
        // But for clarity: use the return_to_go kernel on per-episode windows.
        //
        // Gather episode rewards: for each episode, copy context_len rewards from global array.
        // We'll use a simple kernel (dt_build_trajectories_kernel accesses features by start index,
        // so let's do the same for rewards using a small gather pass).
        //
        // To avoid adding yet another kernel, compute episode rewards on CPU (num_episodes * T floats).
        // This is the ONE-TIME trajectory build, not per-step.
        {
            let mut ep_rewards_host = vec![0.0_f32; ep_total];
            let mut global_rewards = vec![0.0_f32; num_bars];
            stream.memcpy_dtoh(&rewards_gpu, &mut global_rewards)
                .map_err(|e| MLError::ModelError(format!("DT rewards readback: {e}")))?;
            stream.synchronize()
                .map_err(|e| MLError::ModelError(format!("DT sync rewards: {e}")))?;

            for (ep_idx, &start) in episode_starts.iter().enumerate() {
                let start_u = start as usize;
                for t in 0..context_len {
                    let bar = (start_u + t).min(num_bars - 1);
                    ep_rewards_host[ep_idx * context_len + t] = global_rewards[bar];
                }
            }

            stream.memcpy_htod(&ep_rewards_host, &mut ep_rewards_gpu)
                .map_err(|e| MLError::ModelError(format!("DT ep_rewards upload: {e}")))?;
        }

        // RTG reverse cumulative sum
        {
            let ep_rewards_ptr = raw_ptr_f32(&ep_rewards_gpu, stream);
            let rtg_ptr = raw_ptr_f32_mut(&mut rtg_gpu, stream);
            let gamma_f32 = gamma as f32;
            let context_len_i32 = context_len as i32;
            let num_episodes_i32 = num_episodes as i32;

            let lc = LaunchConfig {
                grid_dim: (num_episodes as u32, 1, 1),
                block_dim: (1, 1, 1),
                shared_mem_bytes: 0,
            };
            unsafe {
                stream.launch_builder(&kernels.return_to_go)
                    .arg(&ep_rewards_ptr)
                    .arg(&rtg_ptr)
                    .arg(&gamma_f32)
                    .arg(&context_len_i32)
                    .arg(&num_episodes_i32)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT return_to_go: {e}")))?;
            }
        }

        // ── Step 4: Build trajectory tensor on GPU ──
        let traj_total = num_episodes * context_len * input_dim;
        let mut trajectories_gpu = stream.alloc_zeros::<f32>(traj_total)
            .map_err(|e| MLError::ModelError(format!("DT trajectories alloc: {e}")))?;
        let mut target_actions_gpu = stream.alloc_zeros::<i32>(ep_total)
            .map_err(|e| MLError::ModelError(format!("DT target_actions alloc: {e}")))?;

        {
            let features_ptr = raw_ptr_f32(features_gpu, stream);
            let rtg_ptr = raw_ptr_f32(&rtg_gpu, stream);
            let actions_all_ptr_rd = {
                let (ptr, guard) = actions_all_gpu.device_ptr(stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            let starts_ptr = {
                let (ptr, guard) = episode_starts_gpu.device_ptr(stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };
            let traj_ptr = raw_ptr_f32_mut(&mut trajectories_gpu, stream);
            let target_actions_ptr = {
                let (ptr, guard) = target_actions_gpu.device_ptr_mut(stream);
                let _no_drop = std::mem::ManuallyDrop::new(guard);
                ptr
            };

            let t_i32 = context_len as i32;
            let feat_dim_i32 = feat_dim as i32;
            let input_dim_i32 = input_dim as i32;
            let num_episodes_i32 = num_episodes as i32;
            let num_bars_i32_2 = num_bars as i32;
            let block_t = context_len.min(256) as u32;

            let lc = LaunchConfig {
                grid_dim: (num_episodes as u32, 1, 1),
                block_dim: (block_t, 1, 1),
                shared_mem_bytes: 0,
            };
            unsafe {
                stream.launch_builder(&kernels.build_trajectories)
                    .arg(&features_ptr)
                    .arg(&rtg_ptr)
                    .arg(&actions_all_ptr_rd)
                    .arg(&starts_ptr)
                    .arg(&traj_ptr)
                    .arg(&target_actions_ptr)
                    .arg(&t_i32)
                    .arg(&feat_dim_i32)
                    .arg(&input_dim_i32)
                    .arg(&num_episodes_i32)
                    .arg(&num_bars_i32_2)
                    .launch(lc)
                    .map_err(|e| MLError::ModelError(format!("DT build_trajectories: {e}")))?;
            }
        }

        stream.synchronize()
            .map_err(|e| MLError::ModelError(format!("DT build sync: {e}")))?;

        let traj_mb = (traj_total * 4) as f64 / (1024.0 * 1024.0);
        info!(
            num_episodes,
            traj_mb = format!("{traj_mb:.1}"),
            input_dim,
            "DT trajectories built on GPU"
        );

        Ok((trajectories_gpu, target_actions_gpu, num_complete_batches))
    }
}
