// rl_dueling_q_forward.cu — Phase 4 Independent Dueling Q head forward.
//
// Per spec docs/superpowers/specs/2026-05-30-phase4-independent-dueling-head-design.md.
//
// Architecture: parallel head to C51/IQN/π/value_head with ZERO shared
// state with downstream consumers (ensemble, distill, action selection).
// Trains its own V + A weights via Bellman loss on composed_Q at taken
// action. V output feeds PPO advantage baseline (Phase 4.3) — composed_Q
// is internal to this head's loss path and never consumed elsewhere.
//
// Why this design (vs Phase 2 v2 / Phase 3.x failures):
//   - Phase 2 v2 (scalar V + categorical CE): N/A here — uses scalar
//     Bellman loss, no softmax math eating V.
//   - Phase 3.2 (V_IQN cross-architecture calibration): V_dq trained
//     on same Bellman reward signal as C51, scales align by construction.
//   - Phase 3.1 (composed Q perturbs ensemble): composed_Q_dq never
//     feeds ensemble. Only feeds its own Bellman loss + diag.
//   - Phase 3.1-fix (gradient structure mismatch contaminates weights):
//     DuelingQHead has its OWN weights — mean-zero A grad pattern only
//     affects DuelingQHead's training, no shared weights with C51/IQN.
//
// Forward computation:
//   V[b]            = Σ_c w_v[c] × h_t[b, c] + b_v[0]
//   A[b, a]         = Σ_c w_a[c, a] × h_t[b, c] + b_a[a]    for a in 0..N
//   composed_Q[b, a] = V[b] + A[b, a] − (1/N) Σ_a' A[b, a']
//
// Block layout:
//   grid  = (B, 1, 1)
//   block = (HIDDEN_DIM = 128, 1, 1)
//   One block per batch. Each thread computes one hidden-dim term and
//   participates in tree-reduce. Then thread 0 broadcasts to A
//   computation. Each thread computes one action's matmul contribution.
//   Mean reduction over N_ACTIONS done in shared mem.
//
// Memory:
//   shared float s_h[HIDDEN_DIM]                — cached h_t row
//   shared float s_v                            — scalar V value
//   shared float s_a[N_ACTIONS]                 — A values (post-bias)
//
// Per feedback_no_atomicadd: sole-writer per output cell.
// Per feedback_cpu_is_read_only: pure device kernel.

#define HIDDEN_DIM 128
#define N_ACTIONS  11

extern "C" __global__ void rl_dueling_q_forward(
    const float* __restrict__ h_t,            // [B × HIDDEN_DIM]
    const float* __restrict__ w_v,            // [HIDDEN_DIM]
    const float* __restrict__ b_v,            // [1]
    const float* __restrict__ w_a,            // [HIDDEN_DIM × N_ACTIONS], row-major
    const float* __restrict__ b_a,            // [N_ACTIONS]
    int          B,
    float*       __restrict__ v_out,          // [B]
    float*       __restrict__ a_out,          // [B × N_ACTIONS]
    float*       __restrict__ q_composed_out  // [B × N_ACTIONS]
) {
    const int b = blockIdx.x;
    const int c = threadIdx.x;
    if (b >= B)          return;
    if (c >= HIDDEN_DIM) return;

    extern __shared__ float s_h[];  // [HIDDEN_DIM], sized by smem arg

    // ── Load h_t[b] into shared mem ──
    s_h[c] = h_t[b * HIDDEN_DIM + c];
    __syncthreads();

    // ── V projection (block-reduce) ──
    // V[b] = Σ_c w_v[c] × s_h[c] + b_v[0]
    // Tree reduction over HIDDEN_DIM threads.
    __shared__ float s_v_partial[HIDDEN_DIM];
    s_v_partial[c] = w_v[c] * s_h[c];
    __syncthreads();
    // Tree reduce
    for (int stride = HIDDEN_DIM / 2; stride > 0; stride >>= 1) {
        if (c < stride) {
            s_v_partial[c] += s_v_partial[c + stride];
        }
        __syncthreads();
    }
    __shared__ float s_v;
    if (c == 0) {
        s_v = s_v_partial[0] + b_v[0];
        v_out[b] = s_v;
    }
    __syncthreads();

    // ── A projection (each thread handles one action) ──
    // A[b, a] = Σ_c w_a[c, a] × s_h[c] + b_a[a]
    // Threads 0..N_ACTIONS-1 compute one action each.
    // Other threads idle for this section.
    __shared__ float s_a[N_ACTIONS];
    if (c < N_ACTIONS) {
        float acc = 0.0f;
        #pragma unroll
        for (int i = 0; i < HIDDEN_DIM; ++i) {
            acc += w_a[i * N_ACTIONS + c] * s_h[i];
        }
        acc += b_a[c];
        s_a[c] = acc;
        a_out[b * N_ACTIONS + c] = acc;
    }
    __syncthreads();

    // ── Mean over actions (thread 0 only — small N) ──
    __shared__ float s_mean_a;
    if (c == 0) {
        float sum = 0.0f;
        #pragma unroll
        for (int i = 0; i < N_ACTIONS; ++i) {
            sum += s_a[i];
        }
        s_mean_a = sum * (1.0f / (float)N_ACTIONS);
    }
    __syncthreads();

    // ── Compose Q (each thread for one action writes the result) ──
    if (c < N_ACTIONS) {
        q_composed_out[b * N_ACTIONS + c] = s_v + s_a[c] - s_mean_a;
    }
}
