// crates/ml/src/cuda_pipeline/q_branch_stats_kernel.cu
//
// SP5 Task A1: shared signal source — per-branch Q-statistics reduction.
// Single-block, 4-thread (one per branch): Direction(4)/Magnitude(3)/Order(3)/Urgency(3).
// Reads q_out_buf [B, total_actions=13] row-major (produced by compute_expected_q),
// computes per-branch (mean, max_abs_dev, var, entropy) and writes 16 contiguous floats
// to producer_step_scratch_buf.
//
// Outputs feed Pearls 1, 2, 3, 5 via pearl_1_atom_kernel + apply_pearls_ad_kernel chain.
//
// No atomicAdd (feedback_no_atomicadd) — single-thread per branch, no cross-thread state.
// __threadfence_system() after all writes ensures visibility before apply_pearls_ad_kernel
// reads the outputs on the same stream.

#include <cuda_runtime.h>

extern "C" __global__ void q_branch_stats_update(
    const float* __restrict__ q_out_buf,     // [B, total_actions=13] row-major
    int   batch_size,
    const int* __restrict__ action_counts,   // [4]: {4, 3, 3, 3}
    const int* __restrict__ action_offsets,  // [4]: {0, 4, 7, 10}
    float* __restrict__ scratch_buf,
    int   scratch_idx_base                   // 4 outputs per branch × 4 branches = 16 contiguous
) {
    int branch = threadIdx.x;
    if (blockIdx.x != 0 || branch >= 4) return;

    int n_actions = action_counts[branch];
    int act_off   = action_offsets[branch];

    // Pass 1: accumulate sum and sum_sq over (batch × branch-actions)
    float sum    = 0.0f;
    float sum_sq = 0.0f;
    int total = batch_size * n_actions;

    for (int i = 0; i < batch_size; ++i) {
        for (int a = 0; a < n_actions; ++a) {
            float q = q_out_buf[i * 13 + act_off + a];
            sum    += q;
            sum_sq += q * q;
        }
    }
    float mean = sum / (float)total;
    float var  = (sum_sq / (float)total) - (mean * mean);
    // clamp var to 0 to prevent negative float rounding artifacts
    if (var < 0.0f) var = 0.0f;

    // Pass 2: max |Q - mean| across batch × branch-actions
    float max_abs_dev = 0.0f;
    for (int i = 0; i < batch_size; ++i) {
        for (int a = 0; a < n_actions; ++a) {
            float q   = q_out_buf[i * 13 + act_off + a];
            float dev = fabsf(q - mean);
            if (dev > max_abs_dev) max_abs_dev = dev;
        }
    }

    // Pass 3: action entropy from per-sample softmax over branch-actions
    // Averaged over batch to get mean entropy per branch.
    float entropy = 0.0f;
    for (int i = 0; i < batch_size; ++i) {
        // numerically-stable softmax: subtract max_q first
        float max_q = -1e30f;
        for (int a = 0; a < n_actions; ++a) {
            float q = q_out_buf[i * 13 + act_off + a];
            if (q > max_q) max_q = q;
        }
        float sum_exp = 0.0f;
        for (int a = 0; a < n_actions; ++a) {
            sum_exp += expf(q_out_buf[i * 13 + act_off + a] - max_q);
        }
        // accumulate -p * log(p) for each action
        for (int a = 0; a < n_actions; ++a) {
            float p = expf(q_out_buf[i * 13 + act_off + a] - max_q) / sum_exp;
            if (p > 1e-10f) entropy -= p * logf(p);
        }
    }
    entropy /= (float)batch_size;

    // Write 4 outputs per branch (contiguous in scratch):
    //   out_base + 0  → mean        (v_center for Pearl 1)
    //   out_base + 1  → max_abs_dev (v_half input for Pearl 1)
    //   out_base + 2  → var         (Q_VAR_PER_BRANCH shared signal)
    //   out_base + 3  → entropy     (BRANCH_ENTROPY for Pearl 3 input)
    int out_base = scratch_idx_base + branch * 4;
    scratch_buf[out_base + 0] = mean;
    scratch_buf[out_base + 1] = max_abs_dev;
    scratch_buf[out_base + 2] = var;
    scratch_buf[out_base + 3] = entropy;
    __threadfence_system();
}
