/**
 * Decision Transformer CUDA Kernels (Chen et al., 2021).
 *
 * Forward pass: embed → {QKV → causal attention → LN → FFN → LN} × N → action head → CE loss
 * Backward pass (Phase 1): CE backward → action head backward → residual passthrough → embed backward
 *
 * Compiled via NVRTC source concatenation (no #include).
 * The following constants are injected via NVRTC #define:
 *   DT_EMBED_DIM, DT_NUM_HEADS, DT_CONTEXT_LEN, DT_INPUT_DIM, DT_NUM_ACTIONS
 */

/* ── Compile-time defaults (overridden by NVRTC injection) ───────────── */
#ifndef DT_EMBED_DIM
#define DT_EMBED_DIM 128
#endif
#ifndef DT_NUM_HEADS
#define DT_NUM_HEADS 4
#endif
#ifndef DT_CONTEXT_LEN
#define DT_CONTEXT_LEN 20
#endif
#ifndef DT_INPUT_DIM
#define DT_INPUT_DIM 74
#endif
#ifndef DT_NUM_ACTIONS
#define DT_NUM_ACTIONS 9
#endif

#define DT_HEAD_DIM (DT_EMBED_DIM / DT_NUM_HEADS)
#define DT_FFN_DIM  (DT_EMBED_DIM * 4)

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 1: TOKEN EMBEDDING + POSITIONAL ENCODING
 *
 * Linear projection of each timestep's token (return-to-go, state, action)
 * plus learned positional encoding.
 *
 * out[b][t][d] = sum_j(traj[b][t][j] * W_embed[j][d]) + b_embed[d] + pos_embed[t][d]
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_embed_kernel(
    const float* __restrict__ trajectories,  /* [B, T, input_dim] */
    const float* __restrict__ W_embed,       /* [input_dim, embed_dim] */
    const float* __restrict__ b_embed,       /* [embed_dim] */
    const float* __restrict__ pos_embed,     /* [T, embed_dim] */
    float*       __restrict__ output,        /* [B, T, embed_dim] */
    int B, int T, int input_dim, int embed_dim
) {
    int bt = blockIdx.x;                /* combined (batch, timestep) index */
    if (bt >= B * T) return;
    int d = threadIdx.x;                /* embed dim index */

    int b = bt / T;
    int t = bt % T;

    const float* traj = trajectories + (b * T + t) * input_dim;
    float* out = output + (b * T + t) * embed_dim;

    for (int dd = d; dd < embed_dim; dd += blockDim.x) {
        float val = b_embed[dd] + pos_embed[t * embed_dim + dd];
        for (int j = 0; j < input_dim; j++) {
            val += traj[j] * W_embed[j * embed_dim + dd];
        }
        out[dd] = val;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 2: QKV PROJECTION
 *
 * Three linear projections for multi-head attention.
 * Q[b][t][d] = sum_j(input[b][t][j] * W_Q[j][d])   (no bias for Q/K/V in standard DT)
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_qkv_projection_kernel(
    const float* __restrict__ input,    /* [B, T, E] */
    const float* __restrict__ W_Q,      /* [E, E] */
    const float* __restrict__ W_K,      /* [E, E] */
    const float* __restrict__ W_V,      /* [E, E] */
    float*       __restrict__ Q_out,    /* [B, T, E] */
    float*       __restrict__ K_out,    /* [B, T, E] */
    float*       __restrict__ V_out,    /* [B, T, E] */
    int B, int T, int E
) {
    int bt = blockIdx.x;
    if (bt >= B * T) return;
    int d = threadIdx.x;

    const float* x = input + bt * E;

    for (int dd = d; dd < E; dd += blockDim.x) {
        float q = 0.0f, k = 0.0f, v = 0.0f;
        for (int j = 0; j < E; j++) {
            float xj = x[j];
            q += xj * W_Q[j * E + dd];
            k += xj * W_K[j * E + dd];
            v += xj * W_V[j * E + dd];
        }
        Q_out[bt * E + dd] = q;
        K_out[bt * E + dd] = k;
        V_out[bt * E + dd] = v;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 3: CAUSAL SELF-ATTENTION
 *
 * Per-head causal attention with softmax.
 * Grid:  (B, num_heads, 1) — one block per (batch, head)
 * Block: (T, 1, 1) — one thread per query timestep
 *
 * For each query position i:
 *   score[j] = Q[i] dot K[j] / sqrt(D_h)   for j <= i  (causal)
 *   score[j] = -inf                          for j > i
 *   attn[j] = softmax(score[j])
 *   out[i] = sum_j(attn[j] * V[j])
 *
 * Output projection W_O is also applied here:
 *   final[b][t][d] = sum_h( out_h[t][dh] * W_O[h*Dh + dh, d] )
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_causal_attention_kernel(
    const float* __restrict__ Q,        /* [B, T, E] */
    const float* __restrict__ K,        /* [B, T, E] */
    const float* __restrict__ V,        /* [B, T, E] */
    const float* __restrict__ W_O,      /* [E, E] output projection */
    const float* __restrict__ b_O,      /* [E] output bias */
    float*       __restrict__ output,   /* [B, T, E] */
    int B, int T, int E, int num_heads
) {
    int b = blockIdx.x;
    int h = blockIdx.y;
    if (b >= B || h >= num_heads) return;
    int t_idx = threadIdx.x;  /* query timestep */
    if (t_idx >= T) return;

    int Dh = E / num_heads;
    int head_off = h * Dh;

    /* Pointers for this batch element */
    const float* Q_b = Q + b * T * E;
    const float* K_b = K + b * T * E;
    const float* V_b = V + b * T * E;

    /* Shared memory for K and V of this head: [T, Dh] each */
    extern __shared__ float shmem[];
    float* sh_K = shmem;                /* [T, Dh] */
    float* sh_V = shmem + T * Dh;       /* [T, Dh] */

    /* Load K and V for this head into shared memory */
    for (int dd = 0; dd < Dh; dd++) {
        sh_K[t_idx * Dh + dd] = K_b[t_idx * E + head_off + dd];
        sh_V[t_idx * Dh + dd] = V_b[t_idx * E + head_off + dd];
    }
    __syncthreads();

    /* Compute attention scores Q[t_idx] dot K[j] / sqrt(Dh) for j <= t_idx */
    float inv_sqrt_dh = rsqrtf((float)Dh);

    /* Numerically stable softmax: find max score first */
    float max_score = -1e30f;
    for (int j = 0; j <= t_idx; j++) {
        float score = 0.0f;
        for (int dd = 0; dd < Dh; dd++) {
            score += Q_b[t_idx * E + head_off + dd] * sh_K[j * Dh + dd];
        }
        score *= inv_sqrt_dh;
        if (score > max_score) max_score = score;
    }

    /* Compute softmax weights and weighted sum of V */
    float attn_out[128]; /* max Dh — static array to avoid dynamic alloc */
    for (int dd = 0; dd < Dh; dd++) attn_out[dd] = 0.0f;

    float sum_exp = 0.0f;
    for (int j = 0; j <= t_idx; j++) {
        float score = 0.0f;
        for (int dd = 0; dd < Dh; dd++) {
            score += Q_b[t_idx * E + head_off + dd] * sh_K[j * Dh + dd];
        }
        score = expf(score * inv_sqrt_dh - max_score);
        sum_exp += score;
        for (int dd = 0; dd < Dh; dd++) {
            attn_out[dd] += score * sh_V[j * Dh + dd];
        }
    }

    /* Normalize by softmax denominator */
    float inv_sum = 1.0f / (sum_exp + 1e-8f);
    for (int dd = 0; dd < Dh; dd++) {
        attn_out[dd] *= inv_sum;
    }

    /* Store per-head attention output into temp shared buffer for output projection.
     * We use shmem after the sync point (K/V no longer needed for this thread). */
    __syncthreads();

    /* Reuse shmem for the concatenated head output: [T, E] would be too much.
     * Instead, each thread writes its head's output, then we project per-element. */
    /* Write per-head result into a shared buffer [T, Dh] */
    float* sh_head_out = shmem;  /* reuse: [T, Dh] */
    for (int dd = 0; dd < Dh; dd++) {
        sh_head_out[t_idx * Dh + dd] = attn_out[dd];
    }
    __syncthreads();

    /* Output projection: for this head's contribution to output[b][t_idx][d]
     * Each head contributes: sum_dh(head_out[t_idx][dh] * W_O[(h*Dh+dh)*E + d])
     * Since other heads aren't available here, we use atomicAdd on the output. */
    float* out = output + (b * T + t_idx) * E;

    /* Zero output on first head */
    if (h == 0) {
        for (int dd = 0; dd < E; dd++) {
            out[dd] = b_O[dd];  /* Initialize with bias */
        }
    }
    __syncthreads(); /* Ensure bias initialization is visible */

    /* Each head's projection contribution via atomicAdd */
    for (int dd = 0; dd < E; dd++) {
        float proj = 0.0f;
        for (int dh = 0; dh < Dh; dh++) {
            proj += sh_head_out[t_idx * Dh + dh] * W_O[(head_off + dh) * E + dd];
        }
        atomicAdd(&out[dd], proj);
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 4: LAYER NORMALIZATION
 *
 * output = gamma * (input - mean) / sqrt(var + eps) + beta
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_layernorm_kernel(
    const float* __restrict__ input,    /* [B*T, E] */
    const float* __restrict__ gamma,    /* [E] */
    const float* __restrict__ beta,     /* [E] */
    float*       __restrict__ output,   /* [B*T, E] */
    int N,  /* B*T */
    int E
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* x = input + n * E;
    float* out = output + n * E;

    /* Compute mean via block-parallel reduction */
    float local_sum = 0.0f;
    for (int d = tid; d < E; d += blockDim.x) {
        local_sum += x[d];
    }

    /* Warp-level reduction */
    for (int offset = 16; offset > 0; offset >>= 1)
        local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, offset);

    /* Cross-warp reduction via shared memory */
    __shared__ float warp_sums[8]; /* up to 8 warps */
    int warp_id = tid / 32;
    int lane = tid % 32;
    if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_sum;
    __syncthreads();

    float total_sum;
    if (warp_id == 0) {
        float val = (lane < (int)((blockDim.x + 31) / 32)) ? warp_sums[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        total_sum = val;
    }
    total_sum = __shfl_sync(0xFFFFFFFF, total_sum, 0);
    /* Broadcast from warp 0 lane 0 to all threads */
    __shared__ float sh_mean;
    if (tid == 0) sh_mean = total_sum / (float)E;
    __syncthreads();
    float mean = sh_mean;

    /* Compute variance */
    float local_var = 0.0f;
    for (int d = tid; d < E; d += blockDim.x) {
        float diff = x[d] - mean;
        local_var += diff * diff;
    }

    for (int offset = 16; offset > 0; offset >>= 1)
        local_var += __shfl_xor_sync(0xFFFFFFFF, local_var, offset);

    if (lane == 0 && warp_id < 8) warp_sums[warp_id] = local_var;
    __syncthreads();

    float total_var;
    if (warp_id == 0) {
        float val = (lane < (int)((blockDim.x + 31) / 32)) ? warp_sums[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        total_var = val;
    }
    total_var = __shfl_sync(0xFFFFFFFF, total_var, 0);
    __shared__ float sh_inv_std;
    if (tid == 0) sh_inv_std = rsqrtf(total_var / (float)E + 1e-5f);
    __syncthreads();
    float inv_std = sh_inv_std;

    /* Normalize + affine */
    for (int d = tid; d < E; d += blockDim.x) {
        out[d] = gamma[d] * (x[d] - mean) * inv_std + beta[d];
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 5: FFN (FEED-FORWARD NETWORK)
 *
 * hidden = GELU(input × W1 + b1)
 * output = hidden × W2 + b2 + input (residual)
 *
 * GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(E, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_ffn_kernel(
    const float* __restrict__ input,    /* [B*T, E] */
    const float* __restrict__ W1,       /* [E, 4E] */
    const float* __restrict__ b1,       /* [4E] */
    const float* __restrict__ W2,       /* [4E, E] */
    const float* __restrict__ b2,       /* [E] */
    float*       __restrict__ output,   /* [B*T, E] */
    int N,  /* B*T */
    int E
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    int FFN = E * 4;
    const float* x = input + n * E;
    float* out = output + n * E;

    /* Each thread computes its output elements via stride loop.
     * For the intermediate FFN dimension (4E), we need to compute
     * the full hidden vector. Use shared memory for the hidden activations. */
    extern __shared__ float sh_hidden[];  /* [4E] */

    /* Step 1: Compute hidden = GELU(x @ W1 + b1) */
    for (int ff = tid; ff < FFN; ff += blockDim.x) {
        float h = b1[ff];
        for (int j = 0; j < E; j++) {
            h += x[j] * W1[j * FFN + ff];
        }
        /* GELU approximation */
        float x3 = h * h * h;
        float inner = 0.7978845608f * (h + 0.044715f * x3);  /* sqrt(2/pi) ≈ 0.7979 */
        h = 0.5f * h * (1.0f + tanhf(inner));
        sh_hidden[ff] = h;
    }
    __syncthreads();

    /* Step 2: output = hidden @ W2 + b2 + input (residual) */
    for (int d = tid; d < E; d += blockDim.x) {
        float val = b2[d] + x[d];  /* bias + residual */
        for (int ff = 0; ff < FFN; ff++) {
            val += sh_hidden[ff] * W2[ff * E + d];
        }
        out[d] = val;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 6: ACTION HEAD (LINEAR PROJECTION)
 *
 * logits[b][t][a] = sum_d(input[b][t][d] * W_head[d][A]) + b_head[a]
 *
 * Takes the last timestep's embedding and projects to action logits.
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(A, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_action_head_kernel(
    const float* __restrict__ input,    /* [B*T, E] */
    const float* __restrict__ W_head,   /* [E, A] */
    const float* __restrict__ b_head,   /* [A] */
    float*       __restrict__ logits,   /* [B*T, A] */
    int N,  /* B*T */
    int E,
    int A   /* num_actions */
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* x = input + n * E;
    float* out = logits + n * A;

    for (int a = tid; a < A; a += blockDim.x) {
        float val = b_head[a];
        for (int d = 0; d < E; d++) {
            val += x[d] * W_head[d * A + a];
        }
        out[a] = val;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 7: CROSS-ENTROPY LOSS
 *
 * Numerically stable cross-entropy:
 *   loss = -logits[target] + log(sum(exp(logits - max)))  + max
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(A, 256), 1, 1) or (32, 1, 1) for warp reduction
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_cross_entropy_kernel(
    const float* __restrict__ logits,       /* [B*T, A] */
    const int*   __restrict__ targets,      /* [B*T] */
    float*       __restrict__ per_sample_loss, /* [B*T] */
    float*       __restrict__ total_loss,   /* [1] atomicAdd */
    int N,  /* B*T */
    int A   /* num_actions */
) {
    int n = blockIdx.x;
    if (n >= N) return;

    const float* lg = logits + n * A;
    int target = targets[n];

    /* Find max for numerical stability */
    float max_val = -1e30f;
    for (int a = 0; a < A; a++) {
        if (lg[a] > max_val) max_val = lg[a];
    }

    /* log-sum-exp */
    float sum_exp = 0.0f;
    for (int a = 0; a < A; a++) {
        sum_exp += expf(lg[a] - max_val);
    }

    float log_sum_exp = max_val + logf(sum_exp + 1e-8f);
    float target_logit = (target >= 0 && target < A) ? lg[target] : 0.0f;
    float loss = log_sum_exp - target_logit;

    /* Clamp to prevent NaN propagation */
    loss = fminf(fmaxf(loss, 0.0f), 100.0f);

    per_sample_loss[n] = loss;
    atomicAdd(total_loss, loss / (float)N);
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 8: CROSS-ENTROPY BACKWARD
 *
 * dL/d_logits = softmax(logits) - one_hot(target)
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(A, 256), 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_ce_backward_kernel(
    const float* __restrict__ logits,       /* [B*T, A] */
    const int*   __restrict__ targets,      /* [B*T] */
    float*       __restrict__ d_logits,     /* [B*T, A] */
    int N,  /* B*T */
    int A   /* num_actions */
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* lg = logits + n * A;
    float* d_lg = d_logits + n * A;
    int target = targets[n];

    /* Numerically stable softmax */
    float max_val = -1e30f;
    for (int a = 0; a < A; a++) {
        if (lg[a] > max_val) max_val = lg[a];
    }

    float sum_exp = 0.0f;
    for (int a = 0; a < A; a++) {
        sum_exp += expf(lg[a] - max_val);
    }
    float inv_sum = 1.0f / (sum_exp + 1e-8f);

    /* d_logits = softmax - one_hot, scaled by 1/N for mean loss */
    float scale = 1.0f / (float)N;
    for (int a = tid; a < A; a += blockDim.x) {
        float softmax_a = expf(lg[a] - max_val) * inv_sum;
        float one_hot = (a == target) ? 1.0f : 0.0f;
        d_lg[a] = (softmax_a - one_hot) * scale;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 9: LINEAR BACKWARD
 *
 * Given d_output[B*T, O] and input[B*T, I] and W[I, O]:
 *   dW[i][o] += sum_n(input[n][i] * d_output[n][o])     (weight gradient)
 *   db[o]    += sum_n(d_output[n][o])                     (bias gradient)
 *   d_input[n][i] = sum_o(d_output[n][o] * W[i][o])     (input gradient)
 *
 * Grid:  (B*T, 1, 1)
 * Block: (min(max(I,O), 256), 1, 1)
 *
 * Note: weight gradients use atomicAdd across samples.
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_linear_backward_kernel(
    const float* __restrict__ d_output, /* [N, O] */
    const float* __restrict__ input,    /* [N, I] */
    const float* __restrict__ W,        /* [I, O] */
    float*       __restrict__ dW,       /* [I, O] atomicAdd */
    float*       __restrict__ db,       /* [O] atomicAdd */
    float*       __restrict__ d_input,  /* [N, I] */
    int N,
    int I,
    int O
) {
    int n = blockIdx.x;
    if (n >= N) return;
    int tid = threadIdx.x;

    const float* dout = d_output + n * O;
    const float* x    = input + n * I;
    float* dx = d_input + n * I;

    /* Compute d_input[n][i] = sum_o(d_output[n][o] * W[i][o]) */
    for (int i = tid; i < I; i += blockDim.x) {
        float val = 0.0f;
        for (int o = 0; o < O; o++) {
            val += dout[o] * W[i * O + o];
        }
        dx[i] = val;
    }

    /* Accumulate weight + bias gradients via atomicAdd */
    for (int i = tid; i < I; i += blockDim.x) {
        float xi = x[i];
        for (int o = 0; o < O; o++) {
            atomicAdd(&dW[i * O + o], xi * dout[o]);
        }
    }

    /* Bias gradient: only one thread per sample contributes */
    for (int o = tid; o < O; o += blockDim.x) {
        atomicAdd(&db[o], dout[o]);
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 10: GRADIENT ZERO
 *
 * Zeroes a buffer before backward accumulation.
 *
 * Grid:  (ceil(n/256), 1, 1)
 * Block: (256, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_zero_kernel(
    float* __restrict__ buf,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) buf[i] = 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 11: RESIDUAL ADD
 *
 * output[i] += residual[i]
 *
 * Used after LayerNorm to add the residual connection.
 *
 * Grid:  (ceil(n/256), 1, 1)
 * Block: (256, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_residual_add_kernel(
    float* __restrict__ output,
    const float* __restrict__ residual,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) output[i] += residual[i];
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 12: RETURN-TO-GO (REVERSE CUMULATIVE SUM)
 *
 * R_t = r_t + gamma * R_{t+1}
 * Computed backward from T-1 to 0.  Each thread handles one episode.
 *
 * Grid:  (num_episodes, 1, 1)
 * Block: (1, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_return_to_go_kernel(
    const float* __restrict__ rewards,       /* [num_episodes, episode_len] */
    float*       __restrict__ returns_to_go, /* [num_episodes, episode_len] */
    float gamma,
    int episode_len,
    int num_episodes
) {
    int ep = blockIdx.x;
    if (ep >= num_episodes) return;

    const float* r   = rewards       + ep * episode_len;
    float*       rtg = returns_to_go + ep * episode_len;

    /* Backward scan: R_{T-1} = r_{T-1}, R_t = r_t + gamma * R_{t+1} */
    float cumul = 0.0f;
    for (int t = episode_len - 1; t >= 0; t--) {
        cumul = r[t] + gamma * cumul;
        rtg[t] = cumul;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 13: BUILD TRAJECTORY TOKENS
 *
 * Packs [return-to-go, state_features..., action] into the trajectory
 * tensor expected by dt_embed_kernel.
 *
 * For each (episode, timestep):
 *   out[ep * T * input_dim + t * input_dim + 0]              = rtg[ep * T + t]
 *   out[ep * T * input_dim + t * input_dim + 1..state_dim+1] = features[bar_idx * feat_dim + 0..feat_dim]
 *   out[ep * T * input_dim + t * input_dim + state_dim+1]    = action (as float)
 *
 * bar_idx = episode_start_indices[ep] + t
 *
 * Grid:  (num_episodes, 1, 1)
 * Block: (min(T, 256), 1, 1) — one thread per timestep within episode
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_build_trajectories_kernel(
    const float* __restrict__ features,              /* [num_bars, feat_dim] */
    const float* __restrict__ returns_to_go,         /* [num_episodes, T] */
    const int*   __restrict__ actions,               /* [num_bars] — global per-bar actions */
    const int*   __restrict__ episode_start_indices,  /* [num_episodes] */
    float*       __restrict__ trajectories,          /* [num_episodes, T, input_dim] */
    int*         __restrict__ target_actions,        /* [num_episodes * T] */
    int T,
    int feat_dim,
    int input_dim,
    int num_episodes,
    int num_bars
) {
    int ep = blockIdx.x;
    if (ep >= num_episodes) return;

    int start = episode_start_indices[ep];

    for (int t = threadIdx.x; t < T; t += blockDim.x) {
        int bar_idx = start + t;
        /* Clamp to valid bar range */
        if (bar_idx >= num_bars) bar_idx = num_bars - 1;

        float* out = trajectories + (ep * T + t) * input_dim;

        /* Slot 0: return-to-go */
        out[0] = returns_to_go[ep * T + t];

        /* Slots 1..feat_dim: state features */
        const float* feat = features + bar_idx * feat_dim;
        for (int f = 0; f < feat_dim; f++) {
            out[1 + f] = feat[f];
        }

        /* Slot state_dim+1: action as float (state_dim = feat_dim) */
        int act = actions[bar_idx];  /* index into global bar actions */
        out[feat_dim + 1] = (float)act;

        /* Target action for cross-entropy loss */
        target_actions[ep * T + t] = act;
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 14: COMPUTE BAR REWARDS + EXPERT ACTIONS
 *
 * Per bar: reward = close[t+1] / close[t] - 1
 * Expert action: if reward > threshold → Long(8), < -threshold → Short(0), else Flat(4)
 *
 * Grid:  (ceil(num_bars / 256), 1, 1)
 * Block: (256, 1, 1)
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_compute_rewards_actions_kernel(
    const float* __restrict__ targets,    /* [num_bars, 4] — O,H,L,C per bar */
    float*       __restrict__ rewards,    /* [num_bars] */
    int*         __restrict__ actions,    /* [num_bars] */
    int num_bars,
    float action_threshold
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= num_bars) return;

    /* Close price is at offset 3 (O=0, H=1, L=2, C=3) */
    float close_t  = targets[i * 4 + 3];

    float reward;
    if (i < num_bars - 1) {
        float close_t1 = targets[(i + 1) * 4 + 3];
        /* Avoid division by zero */
        reward = (close_t > 1e-8f) ? (close_t1 / close_t - 1.0f) : 0.0f;
    } else {
        reward = 0.0f;  /* Last bar has no future */
    }
    rewards[i] = reward;

    /* Expert action: Long=8 (max exposure), Short=0 (min exposure), Flat=4 (neutral) */
    int action;
    if (reward > action_threshold) {
        action = 8;  /* Long — b0_size-1 */
    } else if (reward < -action_threshold) {
        action = 0;  /* Short */
    } else {
        action = 4;  /* Flat — b0_size/2 */
    }
    actions[i] = action;
}
