// rl_band_head_forward.cu — Phase 4-A band-head forward.
//
// Two entry points:
//   * `rl_band_head_linear_fwd` — small linear projection
//         band_pre[b, j] = b_band[j] + Σ_c W_band[j, c] · h_t[b, c]
//     for j ∈ {0 (b_l), 1 (b_u)}. Grid = (B, 2, 1), Block = (HIDDEN_DIM, 1, 1)
//     (matches the `ppo_policy_logits_fwd` launch shape for parity with
//     `PolicyHead::forward_logits`).
//   * `rl_band_apply_activation` — asymmetric ±|tanh| activation that
//     enforces `b_l ≤ 0 ≤ b_u` by construction:
//         b_l = -|tanh(band_pre[b, 0])| × N_max_eff
//         b_u = +|tanh(band_pre[b, 1])| × N_max_eff
//     where `N_max_eff` is read from `RL_HEAT_CAP_MAX_LOTS_INDEX`. Grid =
//     (ceil(B/32), 1, 1), Block = (32, 1, 1); single-thread-per-batch
//     element, deterministic.
//
// No PRNG, no atomicAdd, no shared mem. Pure read-deterministic;
// graph-capturable. The forward is independent of the master gate at slot
// 799 — the trainer skips the launch when the band is disabled.
//
// Per `pearl_determinism_achieved` discipline; per spec §1.1 / §1.3.

#include <stdint.h>
#include <math.h>

#define HIDDEN_DIM                        128
#define BAND_OUT                          2
#define RL_HEAT_CAP_MAX_LOTS_INDEX       504

// ── Linear forward: band_pre[b, j] = b_band[j] + Σ_c W_band[j, c] · h_t[b, c]
//
// Block-reduce over HIDDEN_DIM threads via shared memory. Matches the
// existing `ppo_policy_logits_fwd` reduction pattern; deterministic
// tree-reduce (no atomics) so output is bit-equal across runs.
extern "C" __global__ void rl_band_head_linear_fwd(
    const float* __restrict__ w_band,    // [BAND_OUT × HIDDEN_DIM]
    const float* __restrict__ b_band,    // [BAND_OUT]
    const float* __restrict__ h_t,       // [B × HIDDEN_DIM]
    int                       b_size,
    float* __restrict__       band_pre   // [B × BAND_OUT]
) {
    const int b = blockIdx.x;
    const int j = blockIdx.y;
    if (b >= b_size || j >= BAND_OUT) return;

    const int tid = threadIdx.x;

    __shared__ float partial[HIDDEN_DIM];
    const float* w_row = w_band + j * HIDDEN_DIM;
    const float* h_row = h_t + b * HIDDEN_DIM;

    partial[tid] = w_row[tid] * h_row[tid];
    __syncthreads();

    // Tree-reduce over HIDDEN_DIM (assumed power-of-two = 128).
    for (int stride = HIDDEN_DIM / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            partial[tid] += partial[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        band_pre[b * BAND_OUT + j] = partial[0] + b_band[j];
    }
}

// ── Asymmetric ±|tanh| activation, scaled by N_max_eff.
//
// Output layout (per spec §1.1):
//   band_out[b, 0] = -|tanh(band_pre[b, 0])| × N_max_eff    (≤ 0)
//   band_out[b, 1] = +|tanh(band_pre[b, 1])| × N_max_eff    (≥ 0)
//
// Guarantees `b_l ≤ 0 ≤ b_u` for every batch element by construction —
// Davis-Norman optimality theorem requires position 0 (flat) to lie inside
// the no-transaction region; this activation enforces it.
extern "C" __global__ void rl_band_apply_activation(
    const float* __restrict__ band_pre,  // [B × BAND_OUT]
    const float* __restrict__ isv,
    int                       b_size,
    float* __restrict__       band_out   // [B × BAND_OUT]
) {
    const int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= b_size) return;

    const float n_max_eff = isv[RL_HEAT_CAP_MAX_LOTS_INDEX];

    const float pre_l = band_pre[b * BAND_OUT + 0];
    const float pre_u = band_pre[b * BAND_OUT + 1];

    const float t_l = tanhf(pre_l);
    const float t_u = tanhf(pre_u);

    band_out[b * BAND_OUT + 0] = -fabsf(t_l) * n_max_eff;
    band_out[b * BAND_OUT + 1] = +fabsf(t_u) * n_max_eff;
}
