/* ══════════════════════════════════════════════════════════════════════════
 * SP18 v2 Phase 0 Task 0.1 (2026-05-08) — D-leg per-action reward
 * decomposition diagnostic.
 *
 * Reads the per-sample reward-component buffer + actions buffer produced
 * by `experience_kernels.cu`, bins per direction-axis (Hold / Long /
 * Short / Flat), and emits 5 per-bin statistics into a 20-float
 * mapped-pinned diagnostic buffer:
 *
 *   col 0: mean(r_micro)        per action bin   (component idx 3)
 *   col 1: mean(r_opp_cost)     per action bin   (component idx 4)
 *   col 2: mean(r_popart)       per action bin   (component idx 0)
 *   col 3: mean(|reward|)       per action bin   (sum of absolute values
 *                                                  across all 6 components)
 *   col 4: fire_rate            per action bin   (frac of bars with any
 *                                                  non-zero component)
 *
 * Buffer layout (row-major, 4 rows × 5 cols):
 *
 *     out[bin * 5 + col]
 *
 * with bin order:
 *
 *     bin 0 = Hold   (DIR_HOLD  = 1)
 *     bin 1 = Long   (DIR_LONG  = 2)
 *     bin 2 = Short  (DIR_SHORT = 0)
 *     bin 3 = Flat   (DIR_FLAT  = 3)
 *
 * The bin order is logical (Hold first — primary observable for SP18 KILL
 * CRITERION) and explicitly does NOT match the DIR_* enum order. Each
 * thread reads its sample's `actions[i]`, decodes the direction via
 * `action / (b1 * b2 * b3)`, and increments counters for the matching
 * bin.
 *
 * The reward-component layout is the SP15 Wave 2 reshape per
 * `reward_component_ema_kernel.cu`:
 *
 *   [c=0] popart    — final on-policy reward (PopArt input)
 *   [c=1] cf        — counterfactual replay
 *   [c=2] trail     — 0.0 placeholder
 *   [c=3] micro     — dense OFI-alignment micro-reward
 *   [c=4] opp_cost  — Flat opportunity cost
 *   [c=5] bonus     — trade-attempt / timing bonus
 *
 * Pearls + invariants:
 *
 *   - `feedback_no_atomicadd.md` — block tree-reduce; one block per bin,
 *     5 shmem accumulators each, no atomicAdd.
 *
 *   - `pearl_no_host_branches_in_captured_graph.md` — pure kernel launch,
 *     no host-side dispatch within the capture region (the call site is
 *     outside the captured graph anyway: per-step diagnostic at the same
 *     cadence as `launch_reward_component_ema_inplace`).
 *
 *   - `feedback_no_stubs.md` — full bin-aware reduction, no zero-output
 *     placeholder.
 *
 *   - `feedback_no_htod_htoh_only_mapped_pinned.md` — `out` is a mapped-
 *     pinned `MappedF32Buffer.dev_ptr`; host reads after stream sync.
 *
 *   - `pearl_first_observation_bootstrap.md` — the kernel writes raw
 *     means each launch (no Wiener-α blend at this Phase 0 cadence —
 *     diagnostic only). Host-side EMA blend is a separate concern; the
 *     Phase 0 KILL CRITERION reads the raw per-epoch means.
 *
 * Args:
 *
 *   rc        — `reward_components_per_sample[n_samples * 6]` device
 *               pointer. Row-major; row i is `[popart, cf, trail, micro,
 *               opp_cost, bonus]`.
 *
 *   actions   — `actions_out[n_samples]` device pointer (i32). Each
 *               int decodes via `dir = action / (b1 * b2 * b3)`.
 *
 *   n_samples — total sample count (production: `alloc_episodes ×
 *               alloc_timesteps`; the CF-doubled buffer half is NOT
 *               consumed by this kernel — pass the base count).
 *
 *   b1, b2, b3 — magnitude / order / urgency branch sizes (used to
 *                decode direction from the factored-action int).
 *
 *   out       — `out[20]` mapped-pinned f32, row-major 4×5 (4 bins × 5
 *               stats per bin). See top-of-file layout comment.
 *
 * Launch: grid=(4, 1, 1), block=(BLOCK_SIZE, 1, 1).
 * Shared mem: 5 × BLOCK_SIZE × sizeof(float) for the 5 reduction tiles.
 * ══════════════════════════════════════════════════════════════════════════ */

#include <cuda_runtime.h>
#include "state_layout.cuh"

#define SP18_DECOMP_BLOCK_SIZE 256

/* Component indices in `reward_components_per_sample` row. */
#define RCP_POPART_IDX   0
#define RCP_MICRO_IDX    3
#define RCP_OPP_COST_IDX 4
/* SP22 H6 Phase 3 (2026-05-13): bumped 6 → 7 for the new aux_align
 * reward component (c=6). The kernel's abs-sum aggregate (col 3
 * `mean(|reward|)`) now naturally includes r_aux_align in the per-bin
 * total — popart/micro/opp_cost per-bin means remain unchanged. */
#define RCP_NUM_COMPS    7

/* Output layout: 4 bins × 5 cols = 20 floats. */
#define SP18_DECOMP_NUM_BINS  4
#define SP18_DECOMP_NUM_COLS  5
#define SP18_DECOMP_OUT_SIZE  (SP18_DECOMP_NUM_BINS * SP18_DECOMP_NUM_COLS)

/* Logical bin → DIR_* enum lookup. Order: Hold, Long, Short, Flat. */
__device__ __forceinline__ int sp18_decomp_bin_to_dir(int bin) {
    switch (bin) {
        case 0: return DIR_HOLD;
        case 1: return DIR_LONG;
        case 2: return DIR_SHORT;
        case 3: return DIR_FLAT;
        default: return -1;   /* unreachable */
    }
}

extern "C" __global__
void reward_decomp_diag(
    const float* __restrict__ rc,        /* [n_samples × 6] */
    const int*   __restrict__ actions,   /* [n_samples] */
    int           n_samples,
    int           b1,
    int           b2,
    int           b3,
    float*        __restrict__ out       /* [20] mapped-pinned */
) {
    extern __shared__ float s_buf[];
    float* s_micro    = s_buf;                                  /* [BLOCK_SIZE] */
    float* s_opp      = s_buf + SP18_DECOMP_BLOCK_SIZE;         /* [BLOCK_SIZE] */
    float* s_popart   = s_buf + 2 * SP18_DECOMP_BLOCK_SIZE;     /* [BLOCK_SIZE] */
    float* s_abs_sum  = s_buf + 3 * SP18_DECOMP_BLOCK_SIZE;     /* [BLOCK_SIZE] */
    float* s_count    = s_buf + 4 * SP18_DECOMP_BLOCK_SIZE;     /* [BLOCK_SIZE] */

    const int tid     = threadIdx.x;
    const int bin     = blockIdx.x;
    if (bin >= SP18_DECOMP_NUM_BINS) return;
    const int target_dir = sp18_decomp_bin_to_dir(bin);
    const int dir_div    = b1 * b2 * b3;

    /* Per-thread accumulators — strided over n_samples. */
    float l_micro   = 0.0f;
    float l_opp     = 0.0f;
    float l_popart  = 0.0f;
    float l_abs_sum = 0.0f;
    float l_count   = 0.0f;        /* count of bars matching this bin */
    float l_fire    = 0.0f;        /* count of bars with non-zero component */

    for (int i = tid; i < n_samples; i += SP18_DECOMP_BLOCK_SIZE) {
        const int dir_i = (dir_div > 0) ? (actions[i] / dir_div) : -1;
        if (dir_i != target_dir) continue;
        const int row = i * RCP_NUM_COMPS;
        const float r_popart = rc[row + RCP_POPART_IDX];
        const float r_micro  = rc[row + RCP_MICRO_IDX];
        const float r_opp    = rc[row + RCP_OPP_COST_IDX];

        /* sum of |components| → mean of |reward| across all 6 components.
         * Read all 6 explicitly so any future component addition fails
         * loudly via the row-stride contract instead of silently skipping. */
        float abs_total = 0.0f;
        #pragma unroll
        for (int c = 0; c < RCP_NUM_COMPS; ++c) {
            abs_total += fabsf(rc[row + c]);
        }

        l_micro   += r_micro;
        l_opp     += r_opp;
        l_popart  += r_popart;
        l_abs_sum += abs_total;
        l_count   += 1.0f;
        if (abs_total > 0.0f) {
            l_fire += 1.0f;
        }
    }

    /* Track count for division. We reuse s_count to reduce both the
     * matching-bin count AND the firing-bar count with a two-phase pattern:
     * first phase reduces matching count (denominator for means); second
     * phase reduces firing count (numerator for fire_rate). */
    s_micro[tid]   = l_micro;
    s_opp[tid]     = l_opp;
    s_popart[tid]  = l_popart;
    s_abs_sum[tid] = l_abs_sum;
    s_count[tid]   = l_count;
    __syncthreads();

    /* Block tree-reduce. */
    for (int s = SP18_DECOMP_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_micro[tid]   += s_micro[tid + s];
            s_opp[tid]     += s_opp[tid + s];
            s_popart[tid]  += s_popart[tid + s];
            s_abs_sum[tid] += s_abs_sum[tid + s];
            s_count[tid]   += s_count[tid + s];
        }
        __syncthreads();
    }

    const float total_micro   = s_micro[0];
    const float total_opp     = s_opp[0];
    const float total_popart  = s_popart[0];
    const float total_abs_sum = s_abs_sum[0];
    const float total_count   = s_count[0];
    __syncthreads();

    /* Phase 2: reduce firing count using the same s_count tile. */
    s_count[tid] = l_fire;
    __syncthreads();
    for (int s = SP18_DECOMP_BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_count[tid] += s_count[tid + s];
        }
        __syncthreads();
    }
    const float total_fire = s_count[0];

    if (tid == 0) {
        /* Means default to 0.0 when no samples landed in the bin (rather
         * than NaN — the consumer reads these as means and would mis-
         * interpret NaN as a kernel failure). fire_rate is bounded
         * [0, 1] by construction. */
        const float denom = (total_count > 0.0f) ? total_count : 1.0f;
        const int   o     = bin * SP18_DECOMP_NUM_COLS;
        out[o + 0] = (total_count > 0.0f) ? (total_micro   / denom) : 0.0f;
        out[o + 1] = (total_count > 0.0f) ? (total_opp     / denom) : 0.0f;
        out[o + 2] = (total_count > 0.0f) ? (total_popart  / denom) : 0.0f;
        out[o + 3] = (total_count > 0.0f) ? (total_abs_sum / denom) : 0.0f;
        out[o + 4] = (total_count > 0.0f) ? (total_fire    / denom) : 0.0f;
        __threadfence_system();   /* PCIe-visible to mapped-pinned host_ptr */
    }
}
