// multi_head_policy_backward.cu — Backward pass for the K-head policy mixture.
//
// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
//       ADDENDUM 2026-06-03 §R.4 (per-head aux KL priors), §R.6 (gradient flow)
// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-B)
//
// ── Forward recap (multi_head_policy_forward.cu) ────────────────────────
//   pi_logits_k[b,k,a] = W_heads[k,a,:] · h_t[b,:] + b_heads[k,a]
//   pi_probs_k[b,k,a]  = softmax_a(pi_logits_k[b,k,:])
//   pi_probs[b,a]      = Σ_k gate_probs[b,k] · pi_probs_k[b,k,a]
//   pi_logits[b,a]     = log(pi_probs[b,a] + ε)
//
// ── Backward chain implemented by this file (two kernels) ───────────────
//
// (1) `multi_head_policy_backward_pi` — distributes incoming `grad_pi_logits`
//     through the mixture and per-head softmax. Outputs:
//       a. grad_pi_logits_k [B × K × N_ACTIONS]  (for aux prior + diag)
//       b. grad_gate_probs  [B × K]              (consumed by kernel 2)
//       c. grad_W_heads_pb  [B × K × N_ACTIONS × HIDDEN_DIM]
//       d. grad_b_heads_pb  [B × K × N_ACTIONS]
//       e. grad_h_t_scratch [B × HIDDEN_DIM]     (OVERWRITE; caller folds
//                                                 via grad_h_accumulate_scaled)
//
// (2) `multi_head_policy_backward_gate` — distributes `grad_gate_probs`
//     through the gating softmax. Outputs:
//       a. grad_gate_logits  [B × K]                (diag)
//       b. grad_W_gate_pb    [B × K × REGIME_DIM]
//       c. grad_b_gate_pb    [B × K]
//       d. grad_regime_h     [B × REGIME_DIM]       (computed; NOT used
//                                                    upstream in Phase 2A-B)
//
// ── Math (kernel 1) ─────────────────────────────────────────────────────
//   grad_pi_probs[b,a]    = grad_pi_logits[b,a] / (pi_probs[b,a] + ε)
//   grad_pi_probs_k[b,k,a] = gate_probs[b,k] · grad_pi_probs[b,a]
//   grad_gate_probs[b,k]   = Σ_a pi_probs_k[b,k,a] · grad_pi_probs[b,a]
//
//   (per-head softmax-Jacobian; standard form):
//   grad_pi_logits_k[b,k,a] = pi_probs_k[b,k,a]
//                             · ( grad_pi_probs_k[b,k,a]
//                                 − Σ_{a'} pi_probs_k[b,k,a'] · grad_pi_probs_k[b,k,a'] )
//
// ── Math (kernel 2) ─────────────────────────────────────────────────────
//   grad_gate_logits[b,k] = gate_probs[b,k]
//                           · ( grad_gate_probs[b,k]
//                               − Σ_{k'} gate_probs[b,k'] · grad_gate_probs[b,k'] )
//
//   grad_W_gate_pb[b,k,r]  = grad_gate_logits[b,k] · regime_h[b,r]
//   grad_b_gate_pb[b,k]    = grad_gate_logits[b,k]
//   grad_regime_h[b,r]     = Σ_k W_gate[k,r] · grad_gate_logits[b,k]
//
// ── Per-batch grad scratch convention ───────────────────────────────────
//   Caller reduces `_pb` buffers via `reduce_axis0`:
//     grad_W_heads_pb [B × K·N·H]   → grad_W_heads [K·N·H]
//     grad_b_heads_pb [B × K·N]      → grad_b_heads [K·N]
//     grad_W_gate_pb  [B × K·R]      → grad_W_gate  [K·R]
//     grad_b_gate_pb  [B × K]        → grad_b_gate  [K]
//
// Per feedback_no_atomicadd: NO atomicAdd. Each output cell has exactly one
// writer per launch. Per-batch scratch + reduce_axis0 is the standard
// foxhunt pattern (see ppo_clipped_surrogate.cu::ppo_grad_w_b_h_t).
//
// Per feedback_no_nvrtc: pre-compiled cubin via build.rs.

#include <stdint.h>

#define HIDDEN_DIM    128
#define N_ACTIONS     11
#define REGIME_DIM    6
#define MAX_K_HEADS   8       // Compile-time max; runtime K from ISV
#define LOG_PROB_EPS  1e-12f  // mirrors forward kernel's `+1e-12` guard

// ─────────────────────────────────────────────────────────────────────
// Kernel 1: backward through the mixture + per-head softmax.
//
// Block layout:
//   grid  = (B, 1, 1)
//   block = (HIDDEN_DIM = 128, 1, 1)
//   shared = MAX_K_HEADS*N_ACTIONS*3 + MAX_K_HEADS + N_ACTIONS ≈ 304 floats
//
// Thread c handles one hidden-dim index. Threads 0..N_ACTIONS-1 first
// reconstruct pi_probs[b,:] and the softmax-Jacobian — only threads with
// `c < N_ACTIONS` participate in the small reductions. All HIDDEN_DIM
// threads participate in the per-batch weight-grad emit.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void multi_head_policy_backward_pi(
    const float* __restrict__ grad_pi_logits, // [B × N_ACTIONS] incoming grad on log-mixture
    const float* __restrict__ h_t,            // [B × HIDDEN_DIM]
    const float* __restrict__ W_heads,        // [K × N_ACTIONS × HIDDEN_DIM]
    const float* __restrict__ gate_probs,     // [B × K]  forward output
    const float* __restrict__ pi_probs_k,     // [B × K × N_ACTIONS]  forward output
    const int    B,
    const int    K,
    float*       __restrict__ grad_pi_logits_k,   // [B × K × N_ACTIONS] OUT (for aux + diag)
    float*       __restrict__ grad_gate_probs,    // [B × K]             OUT (→ kernel 2)
    float*       __restrict__ grad_W_heads_pb,    // [B × K × N_ACTIONS × HIDDEN_DIM]
    float*       __restrict__ grad_b_heads_pb,    // [B × K × N_ACTIONS]
    float*       __restrict__ grad_h_t_scratch    // [B × HIDDEN_DIM] (OVERWRITE)
) {
    const int b = blockIdx.x;
    if (b >= B) return;
    const int c = threadIdx.x;  // hidden-dim index, 0..HIDDEN_DIM-1

    // ── Shared state (per-batch) ──────────────────────────────────────
    __shared__ float s_pi_probs[N_ACTIONS];                       // mixture probs
    __shared__ float s_grad_pi_probs[N_ACTIONS];                  // dL/dp[a]
    __shared__ float s_gate_probs[MAX_K_HEADS];                   // forward gate
    __shared__ float s_pi_probs_k[MAX_K_HEADS * N_ACTIONS];       // forward per-head probs
    __shared__ float s_grad_pi_probs_k[MAX_K_HEADS * N_ACTIONS];  // step 2 intermediate
    __shared__ float s_grad_pi_logits_k[MAX_K_HEADS * N_ACTIONS]; // final per-head logit grad
    __shared__ float s_dot_k[MAX_K_HEADS];                        // softmax-Jacobian scalar per head

    // ── Stage forward outputs (small, per-batch) into shared mem ───────
    if (c < K) {
        s_gate_probs[c] = gate_probs[b * K + c];
    }
    if (c < N_ACTIONS) {
        const int gpi_idx = b * N_ACTIONS + c;
        s_grad_pi_probs[c] = 0.0f;  // accumulator, populated below
        // Stage pi_probs_k[b, :, c] across all heads.
        for (int k = 0; k < K; ++k) {
            s_pi_probs_k[k * N_ACTIONS + c] =
                pi_probs_k[(b * K + k) * N_ACTIONS + c];
        }
        // Reconstruct pi_probs[b, c] = Σ_k gate · pi_probs_k.
        // Sequential summation order matches the forward kernel for
        // determinism.
        float p = 0.0f;
        #pragma unroll
        for (int k = 0; k < K; ++k) {
            p += gate_probs[b * K + k] * pi_probs_k[(b * K + k) * N_ACTIONS + c];
        }
        s_pi_probs[c] = p;
        // dL/dp[a] = dL/dlog(p+ε) × 1/(p+ε)
        const float gl = grad_pi_logits[gpi_idx];
        s_grad_pi_probs[c] = gl / (s_pi_probs[c] + LOG_PROB_EPS);
    }
    __syncthreads();

    // ── Step 2: grad_pi_probs_k + grad_gate_probs ─────────────────────
    // grad_pi_probs_k[b,k,a] = gate_probs[b,k] · grad_pi_probs[b,a]
    // grad_gate_probs[b,k]   = Σ_a pi_probs_k[b,k,a] · grad_pi_probs[b,a]
    if (c < N_ACTIONS) {
        for (int k = 0; k < K; ++k) {
            s_grad_pi_probs_k[k * N_ACTIONS + c] =
                s_gate_probs[k] * s_grad_pi_probs[c];
        }
    }
    __syncthreads();

    // Thread k computes Σ_a pi_probs_k[k,a] · grad_pi_probs[a] for its head.
    if (c < K) {
        float gp = 0.0f;
        #pragma unroll
        for (int a = 0; a < N_ACTIONS; ++a) {
            gp += s_pi_probs_k[c * N_ACTIONS + a] * s_grad_pi_probs[a];
        }
        grad_gate_probs[b * K + c] = gp;
    }
    __syncthreads();

    // ── Step 3: per-head softmax-Jacobian → grad_pi_logits_k ──────────
    // dot_k[b,k] = Σ_{a'} pi_probs_k[k,a'] · grad_pi_probs_k[k,a']
    if (c < K) {
        float dot = 0.0f;
        #pragma unroll
        for (int a = 0; a < N_ACTIONS; ++a) {
            dot += s_pi_probs_k[c * N_ACTIONS + a]
                 * s_grad_pi_probs_k[c * N_ACTIONS + a];
        }
        s_dot_k[c] = dot;
    }
    __syncthreads();

    if (c < N_ACTIONS) {
        for (int k = 0; k < K; ++k) {
            const int kak = k * N_ACTIONS + c;
            const float gl = s_pi_probs_k[kak]
                           * ( s_grad_pi_probs_k[kak] - s_dot_k[k] );
            s_grad_pi_logits_k[kak] = gl;
            // Persist to global memory for aux prior kernel + diagnostics.
            grad_pi_logits_k[(b * K + k) * N_ACTIONS + c] = gl;
        }
    }
    __syncthreads();

    // ── Step 4: per-batch weight grads + grad_h_t ─────────────────────
    // Each thread c is the sole writer of:
    //   grad_W_heads_pb[b, k, a, c]   for all (k, a)
    //   grad_h_t_scratch[b, c]
    // Threads c<K * N_ACTIONS additionally write grad_b_heads_pb[b, k, a].
    if (c < HIDDEN_DIM) {
        const float h_bc = h_t[b * HIDDEN_DIM + c];
        float gh_acc = 0.0f;
        for (int k = 0; k < K; ++k) {
            for (int a = 0; a < N_ACTIONS; ++a) {
                const int kak = k * N_ACTIONS + a;
                const float g = s_grad_pi_logits_k[kak];
                // grad_W_heads_pb[b, k, a, c] = g * h[b, c]
                grad_W_heads_pb[(long long)((b * K + k) * N_ACTIONS + a) * HIDDEN_DIM + c]
                    = g * h_bc;
                // Accumulate grad_h_t[b, c] = Σ_{k,a} g · W_heads[k, a, c]
                gh_acc += g * W_heads[(long long)((k * N_ACTIONS) + a) * HIDDEN_DIM + c];
            }
        }
        // OVERWRITE — caller folds via grad_h_accumulate_scaled.
        grad_h_t_scratch[b * HIDDEN_DIM + c] = gh_acc;
    }

    // Bias grads — sole-writer per (b, k, a). Use the first K*N_ACTIONS
    // threads to cover the per-(k, a) slots in one pass.
    if (c < K * N_ACTIONS) {
        const int k = c / N_ACTIONS;
        const int a = c % N_ACTIONS;
        grad_b_heads_pb[(b * K + k) * N_ACTIONS + a] =
            s_grad_pi_logits_k[k * N_ACTIONS + a];
    }
}

// ─────────────────────────────────────────────────────────────────────
// Kernel 2: backward through the gating softmax.
//
// Block layout:
//   grid  = (B, 1, 1)
//   block = (max(K, REGIME_DIM) = 8, 1, 1)
//
// Tiny block — K≤8 and REGIME_DIM=6. Thread layout is per-head for the
// softmax-Jacobian reduction, per-regime-dim for the regime-grad output.
// We use the same threads for both phases since both dimensions are ≤8.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void multi_head_policy_backward_gate(
    const float* __restrict__ grad_gate_probs, // [B × K] from kernel 1
    const float* __restrict__ gate_probs,      // [B × K] forward output
    const float* __restrict__ regime_h,        // [B × REGIME_DIM]
    const float* __restrict__ W_gate,          // [K × REGIME_DIM]
    const int    B,
    const int    K,
    float*       __restrict__ grad_gate_logits,// [B × K]            OUT (for diag)
    float*       __restrict__ grad_W_gate_pb,  // [B × K × REGIME_DIM]
    float*       __restrict__ grad_b_gate_pb,  // [B × K]
    float*       __restrict__ grad_regime_h    // [B × REGIME_DIM] (computed; unused upstream)
) {
    const int b = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;

    __shared__ float s_gate_probs[MAX_K_HEADS];
    __shared__ float s_grad_gate_probs[MAX_K_HEADS];
    __shared__ float s_grad_gate_logits[MAX_K_HEADS];
    __shared__ float s_dot;  // Σ_{k'} gate · grad_gate_probs

    // ── Stage forward + incoming grad ─────────────────────────────────
    if (tid < K) {
        s_gate_probs[tid] = gate_probs[b * K + tid];
        s_grad_gate_probs[tid] = grad_gate_probs[b * K + tid];
    }
    __syncthreads();

    // ── Softmax-Jacobian reduction (single thread for determinism) ────
    if (tid == 0) {
        float dot = 0.0f;
        for (int k = 0; k < K; ++k) {
            dot += s_gate_probs[k] * s_grad_gate_probs[k];
        }
        s_dot = dot;
    }
    __syncthreads();

    // ── grad_gate_logits[b, k] = p · (dp − dot) ───────────────────────
    if (tid < K) {
        const float gl = s_gate_probs[tid] * (s_grad_gate_probs[tid] - s_dot);
        s_grad_gate_logits[tid] = gl;
        grad_gate_logits[b * K + tid] = gl;
        grad_b_gate_pb[b * K + tid] = gl;
    }
    __syncthreads();

    // ── grad_W_gate_pb[b, k, r] = grad_gate_logits[k] · regime_h[r] ───
    // Cover all (k, r) slots — K*REGIME_DIM ≤ 48 — using a single loop
    // per active thread. Thread `tid` handles its column across all k.
    if (tid < REGIME_DIM) {
        const float r_b = regime_h[b * REGIME_DIM + tid];
        for (int k = 0; k < K; ++k) {
            grad_W_gate_pb[((b * K + k) * REGIME_DIM) + tid] =
                s_grad_gate_logits[k] * r_b;
        }
    }

    // ── grad_regime_h[b, r] = Σ_k W_gate[k, r] · grad_gate_logits[k] ──
    // Computed but NOT propagated upstream in Phase 2A-B (regime_h is
    // loader-precomputed in foxhunt). Allocate + write the buffer so a
    // future phase can route it if needed.
    if (tid < REGIME_DIM) {
        float gr = 0.0f;
        #pragma unroll
        for (int k = 0; k < K; ++k) {
            gr += W_gate[k * REGIME_DIM + tid] * s_grad_gate_logits[k];
        }
        grad_regime_h[b * REGIME_DIM + tid] = gr;
    }
}
