// crates/ml/src/cuda_pipeline/alpha_linear_q.cu
//
// Phase E.1 Task 12a (2026-05-15): linear Q-network kernels for the
// H=600 DQN smoke. Single linear layer:
//
//   Q[b, i] = sum_j W[i, j] * X[b, j] + b[i]
//
// No hidden layer — the smoke's purpose is to validate whether ANY signal
// propagates through the env + reward shaping + kill-criteria pipeline.
// If linear Q can't learn, no architecture upgrade will save it. The
// 10-dim state has meaningful direct features (alpha_logit, spread_bps,
// position, ...) so a linear model has real expressive power for this
// kill-criteria gate.
//
// Three kernels in this file:
//   1. alpha_linear_q_forward_kernel — batched Q forward pass
//   2. alpha_linear_q_grad_kernel    — gradient of MSE-style TD loss
//                                      w.r.t. W and b (sparse over actions)
//   3. alpha_linear_q_sgd_step_kernel — plain SGD update (W -= lr · dW)
//
// **GPU contract:** all three are simple element-wise / inner-loop work.
// No atomicAdd (per `feedback_no_atomicadd`). No host branches inside
// kernels (per `pearl_no_host_branches_in_captured_graph`). Tiny grid /
// block configs fit per-step launch overhead.

#include <cuda_runtime.h>

// ----------------------------------------------------------------------
// (1) Forward: Q = X · W^T + b
// ----------------------------------------------------------------------
//
// Threads: one per (b, i) where b ∈ [0, batch) and i ∈ [0, n_actions).
// Grid:    (ceil(batch * n_actions / 256), 1, 1); Block: (256, 1, 1).
//
// Each thread loops over state_dim accumulating W[i, j] * X[b, j].
extern "C" __global__ void alpha_linear_q_forward_kernel(
    const float* __restrict__ W,    // [n_actions, state_dim] row-major
    const float* __restrict__ b,    // [n_actions]
    const float* __restrict__ X,    // [batch, state_dim] row-major
    float*       __restrict__ Q,    // [batch, n_actions] row-major
    int          batch,
    int          state_dim,
    int          n_actions
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = batch * n_actions;
    if (idx >= total) return;

    const int sample = idx / n_actions;
    const int action = idx % n_actions;

    float acc = b[action];
    const int w_row_base = action * state_dim;
    const int x_row_base = sample * state_dim;
    #pragma unroll 4
    for (int j = 0; j < state_dim; ++j) {
        acc += W[w_row_base + j] * X[x_row_base + j];
    }
    Q[idx] = acc;
}

// ----------------------------------------------------------------------
// (2) Gradient: ∂L/∂W and ∂L/∂b for sparse TD-style loss
// ----------------------------------------------------------------------
//
// Per the standard DQN gradient, only the action taken at sample b
// contributes to the gradient. For each (i, j):
//
//   dW[i, j] = scale · sum_{b: actions[b] == i} (Q[b, i] − target[b]) · X[b, j]
//   db[i]    = scale · sum_{b: actions[b] == i} (Q[b, i] − target[b])
//
// `scale` defaults to 1/batch (standard MSE batching) but is passed as
// a kernel arg so a controller could ISV-drive it (e.g., IS-weighted
// prioritized replay).
//
// Threads: 1 per (i, j) for dW, 1 per i for db. Combined index into a
// single (n_actions * state_dim + n_actions)-thread grid for simplicity.
// At batch=64 and (9, 10): 99 threads — single block, fine.
extern "C" __global__ void alpha_linear_q_grad_kernel(
    const float* __restrict__ Q,        // [batch, n_actions] — current Q from forward
    const float* __restrict__ target,   // [batch] — TD target for the taken action
    const int*   __restrict__ actions,  // [batch] — action taken per sample
    const float* __restrict__ X,        // [batch, state_dim] — input states
    float*       __restrict__ dW,       // [n_actions, state_dim] — output, overwritten
    float*       __restrict__ db,       // [n_actions] — output, overwritten
    int          batch,
    int          state_dim,
    int          n_actions,
    float        scale                  // typ. 1/batch
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total_w = n_actions * state_dim;

    if (idx < total_w) {
        const int i = idx / state_dim; // action
        const int j = idx % state_dim; // input feature dim
        float grad = 0.0f;
        for (int b_idx = 0; b_idx < batch; ++b_idx) {
            // Sparse: only contribute when sample b's action matches i.
            if (actions[b_idx] == i) {
                const float td_err = Q[b_idx * n_actions + i] - target[b_idx];
                grad += td_err * X[b_idx * state_dim + j];
            }
        }
        dW[idx] = grad * scale;
    } else if (idx < total_w + n_actions) {
        const int i = idx - total_w;
        float grad = 0.0f;
        for (int b_idx = 0; b_idx < batch; ++b_idx) {
            if (actions[b_idx] == i) {
                grad += Q[b_idx * n_actions + i] - target[b_idx];
            }
        }
        db[i] = grad * scale;
    }
}

// ----------------------------------------------------------------------
// (3) SGD step: W -= lr · dW (element-wise, flat over W or b)
// ----------------------------------------------------------------------
//
// Generic enough to update either the weight matrix or the bias vector
// — caller passes `n = numel` of the buffer.
extern "C" __global__ void alpha_linear_q_sgd_step_kernel(
    float*       __restrict__ params,   // [n] — overwritten
    const float* __restrict__ grad,     // [n]
    float        lr,
    int          n
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    params[idx] -= lr * grad[idx];
}

// ----------------------------------------------------------------------
// (4) Element-wise clip-in-place: x = clamp(x, -bound, +bound)
// ----------------------------------------------------------------------
//
// Stabilizer for the Munchausen-target gradient. The Q-network can grow
// rapidly when ‖dW‖ exceeds the per-step LR's safe operating range
// (lr · ‖dW‖ → weight delta). Without clipping, a single burst step can
// double the weight magnitude, after which the V_soft bootstrap target
// also doubles, gradient doubles → exponential blow-up.
//
// Clipping per-element is the simplest safety net (vs L2-norm clipping
// which needs a reduction). For the linear Q-network with 99 parameters,
// per-element clipping is appropriate.
extern "C" __global__ void alpha_clip_inplace_kernel(
    float* __restrict__ x,    // [n] — overwritten in place
    float  bound,
    int    n
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    const float v = x[idx];
    x[idx] = fmaxf(-bound, fminf(bound, v));
}
