// rl_band_head_backward.cu — Phase 4-B band-head backward chain.
//
// Spec: docs/superpowers/specs/2026-06-03-no-transaction-band-architecture.md
// §3.3 (backward chain into encoder).
//
// Forward chain (recap, see `rl_band_head_forward.cu`):
//     band_pre[b, j] = b_band[j] + Σ_c W_band[j, c] · h_t[b, c]
//     b_l = -|tanh(band_pre[b, 0])| × N_max_eff
//     b_u = +|tanh(band_pre[b, 1])| × N_max_eff
//
// Backward chain (this file):
//   Given dL/d(b_l) = grad_band[b, 0] and dL/d(b_u) = grad_band[b, 1] from
//   the turnover-loss kernel:
//
//     dL/d(band_pre[b, 0]) = grad_band[b, 0] · d(b_l)/d(pre)
//                          = grad_band[b, 0] · (-sign(t_l)) · (1 - t_l²) · N_max
//                          = grad_band[b, 0] · (-sign(t_l)) · sech²(pre_l) · N_max
//     dL/d(band_pre[b, 1]) = grad_band[b, 1] · d(b_u)/d(pre)
//                          = grad_band[b, 1] · (+sign(t_u)) · sech²(pre_u) · N_max
//
//   Then the linear backward:
//     dL/d(W_band[j, c]) per batch = dL/d(pre[b, j]) · h_t[b, c]
//     dL/d(b_band[j])    per batch = dL/d(pre[b, j])
//     dL/d(h_t[b, c])             = Σ_j dL/d(pre[b, j]) · W_band[j, c]
//
// Outputs (per-batch scratch; caller reduces via `reduce_axis0`):
//   grad_w_per_batch [B × BAND_OUT × HIDDEN_DIM]
//   grad_b_per_batch [B × BAND_OUT]
//   grad_h_t         [B × HIDDEN_DIM] (OVERWRITE — caller folds into encoder
//                                      grad via `grad_h_accumulate_scaled`)
//
// Disciplines:
//   * `feedback_no_atomicadd`        — per-batch scratch + reduce_axis0
//   * `feedback_no_nvrtc`            — pre-compiled cubin via build.rs
//   * `pearl_determinism_achieved`   — sole-writer per (b, j, c); no PRNG
//   * `pearl_no_host_branches_in_captured_graph` — only kernel args, no
//     host scalars in control flow
//
// Launch:
//   grid  = (B, 1, 1)
//   block = (HIDDEN_DIM, 1, 1)
//   shared_mem_bytes = BAND_OUT × sizeof(float)   (s_dpre[2] staging)
//
// Per-block: thread c writes its own grad_w[b, 0, c], grad_w[b, 1, c],
// grad_h_t[b, c]. Thread 0 additionally writes grad_b[b, 0] and grad_b[b, 1].

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

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

extern "C" __global__ void rl_band_head_backward(
    const float* __restrict__ w_band,             // [BAND_OUT × HIDDEN_DIM]
    const float* __restrict__ band_pre,           // [B × BAND_OUT]
    const float* __restrict__ h_t,                // [B × HIDDEN_DIM]
    const float* __restrict__ grad_band,          // [B × BAND_OUT]
    const float* __restrict__ isv,
    int                       b_size,
    float*       __restrict__ grad_w_per_batch,   // [B × BAND_OUT × HIDDEN_DIM]
    float*       __restrict__ grad_b_per_batch,   // [B × BAND_OUT]
    float*       __restrict__ grad_h_t            // [B × HIDDEN_DIM] (OVERWRITE)
) {
    const int b = blockIdx.x;
    const int c = threadIdx.x;
    if (b >= b_size || c >= HIDDEN_DIM) return;

    __shared__ float s_dpre[BAND_OUT];

    // Master gate: write zeros and return so the caller's reduce_axis0 +
    // accumulate_grad_h chain produces no encoder-grad contribution. The
    // host-side branch in the trainer ALSO short-circuits the launch when
    // disabled — this device-side guard is defense-in-depth so the kernel
    // is safe to launch unconditionally during graph capture.
    const float enabled = isv[RL_BAND_ENABLED_INDEX];
    if (enabled <= 0.5f) {
        grad_w_per_batch[(b * BAND_OUT + 0) * HIDDEN_DIM + c] = 0.0f;
        grad_w_per_batch[(b * BAND_OUT + 1) * HIDDEN_DIM + c] = 0.0f;
        grad_h_t[b * HIDDEN_DIM + c] = 0.0f;
        if (c == 0) {
            grad_b_per_batch[b * BAND_OUT + 0] = 0.0f;
            grad_b_per_batch[b * BAND_OUT + 1] = 0.0f;
        }
        return;
    }

    // ── Stage 1: thread 0 computes the two activation-derivative scalars ─
    if (c == 0) {
        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);
        // sech²(x) = 1 − tanh²(x)
        const float sech2_l = 1.0f - t_l * t_l;
        const float sech2_u = 1.0f - t_u * t_u;
        // d(b_l)/d(pre_l) = -sign(t_l) · sech²(pre_l) · N_max
        // d(b_u)/d(pre_u) = +sign(t_u) · sech²(pre_u) · N_max
        // sign(0) = 0 — at pre = 0 the gradient is zero (the activation
        // is non-differentiable at the cusp |tanh| → 0). Treat as zero;
        // upstream signs propagate cleanly.
        const float sign_l = (t_l > 0.0f) ? 1.0f : ((t_l < 0.0f) ? -1.0f : 0.0f);
        const float sign_u = (t_u > 0.0f) ? 1.0f : ((t_u < 0.0f) ? -1.0f : 0.0f);

        const float g_b_l = grad_band[b * BAND_OUT + 0];
        const float g_b_u = grad_band[b * BAND_OUT + 1];
        // dL/d(pre_l) = g_b_l · (-sign_l · sech2_l · N_max)
        // dL/d(pre_u) = g_b_u · (+sign_u · sech2_u · N_max)
        s_dpre[0] = g_b_l * (-sign_l) * sech2_l * n_max_eff;
        s_dpre[1] = g_b_u * (+sign_u) * sech2_u * n_max_eff;

        // grad_b_band per batch is just dL/d(pre[b, j]).
        grad_b_per_batch[b * BAND_OUT + 0] = s_dpre[0];
        grad_b_per_batch[b * BAND_OUT + 1] = s_dpre[1];
    }
    __syncthreads();

    const float dpre_l = s_dpre[0];
    const float dpre_u = s_dpre[1];

    // ── Stage 2: per-(b, c) writes ───────────────────────────────────
    // grad_w_band[j, c] per batch = dpre[j] · h_t[b, c]
    const float h_bc = h_t[b * HIDDEN_DIM + c];
    grad_w_per_batch[(b * BAND_OUT + 0) * HIDDEN_DIM + c] = dpre_l * h_bc;
    grad_w_per_batch[(b * BAND_OUT + 1) * HIDDEN_DIM + c] = dpre_u * h_bc;

    // grad_h_t[b, c] = dpre_l · W_band[0, c] + dpre_u · W_band[1, c]
    // OVERWRITE — caller's `grad_h_accumulate_scaled` folds into encoder
    // grad via additive +=.
    const float w_lc = w_band[0 * HIDDEN_DIM + c];
    const float w_uc = w_band[1 * HIDDEN_DIM + c];
    grad_h_t[b * HIDDEN_DIM + c] = dpre_l * w_lc + dpre_u * w_uc;
}
