/* ══════════════════════════════════════════════════════════════════════════
 * SP20 Phase 1.4 (2026-05-09) — sp20_aggregate_inputs kernel.
 *
 * Path C aggregation step. Reads per-env GPU arrays produced by
 * `experience_env_step` (trade_close, step_ret, hold_at_exit, actions)
 * + the upstream `sp20_stats_compute` outputs (aux_logits_p50/std) +
 * the existing `aux_dir_acc_reduce_kernel` output [0] (aux_dir_acc),
 * aggregates them into the 9-field `SP20EmaInputs` struct that
 * `sp20_emas_compute_kernel` reads.
 *
 * Aggregation rules (one struct, one rollout-step):
 *
 *   is_close = 1 if any env's trade_close_per_env[env] != 0, else 0.
 *   is_win   = 1 if (is_close == 1) AND fraction of closed envs with
 *              `is_win_per_env[env] != 0` is >= 0.5, else 0.
 *              SP20 Phase 2 Task 2.2-fix (2026-05-10): the predicate
 *              now reads the segment-level per-env `is_win_per_env`
 *              flag (1 iff `segment_return > 0` at trade close) rather
 *              than the per-bar `step_ret_per_env > 0` proxy. The
 *              per-bar step_return at close bars is dominated by
 *              tx-cost so `sr > 0` was almost always false, pinning
 *              `WR_EMA` at 0 in production despite real wins. NULL-
 *              tolerant fallback to the legacy predicate preserved for
 *              oracle tests without a segment-level producer wired.
 *   trade_duration = round(mean(hold_at_exit_per_env[env]) for env
 *                          where trade_close_per_env[env] != 0)
 *                    if is_close == 1, else 0.
 *   alpha    = mean(alpha_per_env[env]) over envs with
 *              trade_close_per_env[env] != 0; 0.0 if is_close == 0.
 *              (SP20 Phase 2 Task 2.2 — replaced the Phase 1.4 hardcoded
 *              0.0 forward-reference placeholder. The per-env producer
 *              writes alpha = R_event - hold_baseline at every trade-close
 *              from `sp20_compute_event_reward`. hold_baseline is still
 *              0.0f in Phase 2; Phase 3.2's Component 2 lands the real
 *              `hold_baseline_buffer.sum_over_trade_range()` consumer.)
 *   per_bar_hold_reward = 0.0 (Phase 3.2 forward reference — the
 *              `r_per_bar = -aux_conf * cost_scale` formula needs the
 *              per-bar aux_conf signal that Phase 3.2's per-bar aux
 *              forward emits. Phase 1.4 writes 0.0 so the HOLD_REWARD
 *              EMA's bootstrap counter advances on Hold-action bars
 *              while staying at its 0.0 sentinel until Phase 3.2.)
 *   aux_logits_p50 = aux_logits_p50_dev[0]   (sp20_stats output)
 *   aux_logits_std = aux_logits_std_dev[0]   (sp20_stats output)
 *   aux_dir_acc    = aux_dir_acc_dev[0]      (existing aux_dir_acc_buf[0]
 *                                              from aux_dir_acc_reduce_kernel)
 *   action_is_hold = 1 if count(actions_per_env[env] decoded direction
 *                    == HOLD) > n_envs / 2, else 0.
 *
 * The action decoding uses the production packed factored
 * action_idx encoding from `experience_action_select`:
 *
 *   action_idx = dir * (NUM_MAGNITUDES * NUM_ORD * NUM_URG)
 *              + mag * (NUM_ORD * NUM_URG)
 *              + ord * NUM_URG
 *              + urg
 *
 * so `dir = action_idx / (NUM_MAGNITUDES * NUM_ORD * NUM_URG)`. The
 * caller passes the divisor as `dir_divisor` (typically 27 = 3*3*3
 * for the production branch sizes) and the Hold direction id as
 * `hold_action_id` (= 1 = DIR_HOLD per `state_layout.cuh`).
 *
 * ── Pearls + invariants ────────────────────────────────────────────────
 *  - `feedback_no_atomicadd` — block tree-reduce only. The 3 reductions
 *    (closed-count, wins-count, hold-count) share one shmem tile
 *    sequentially with explicit `__syncthreads()` between phases. The
 *    `hold_at_exit` and SP20 Phase 2 Task 2.2 `alpha` means are computed
 *    as separate sum reductions immediately after; total of 5
 *    tree-reductions, all bdim×i32 or bdim×f32 in shmem (5 stripes).
 *  - `feedback_no_cpu_compute_strict` — every aggregation lives in
 *    this kernel; the launcher only passes pointers + n_envs +
 *    constants.
 *  - `feedback_no_htod_htoh_only_mapped_pinned` — the output struct
 *    buffer is mapped-pinned (or device-only — both work since the
 *    consumer kernel reads via dev_ptr). Kernel emits
 *    `__threadfence_system()` after the writes so the downstream
 *    `sp20_emas_compute_kernel` sees the updates after the in-stream
 *    ordering invariant.
 *  - `pearl_no_host_branches_in_captured_graph` — single-block kernel,
 *    no host branches; safe to capture in the per-step CUDA Graph.
 *  - `feedback_no_partial_refactor` — kernel + launcher + test +
 *    build entry land atomically with the Phase 1.2 buffer-arg
 *    refactor + the production wire-up + audit/spec/plan
 *    amendments in Phase 1.4.
 *  - `pearl_first_observation_bootstrap` — the placeholder fields
 *    (`alpha = 0.0`, `per_bar_hold_reward = 0.0`) emit the bootstrap
 *    sentinel so Phase 1.2's EMAs see "no firing" semantics for
 *    those EMAs until Phase 2 / Phase 3.2 wire the real producers.
 *
 * Launch contract:
 *   grid_dim  = (1, 1, 1)
 *   block_dim = (BLOCK_SIZE, 1, 1)   — typically 32 or 64 for n_envs
 *                                       up to a few hundred.
 *   shared mem = 5 × bdim × max(sizeof(int), sizeof(float)) = 5 × bdim × 4
 *                (4 stripes pre-Task-2.2 + 1 stripe for the alpha sum)
 * ══════════════════════════════════════════════════════════════════════════ */

#include <cuda_runtime.h>

/* Same struct layout as `sp20_emas_compute_kernel.cu`. The two .cu
 * files are compiled separately into different cubins; redeclaring the
 * struct here keeps each cubin self-contained and avoids a shared
 * .cuh header that would couple the build manifest. The Rust side's
 * single source of truth (`EmaInputs` in `sp20_emas_compute.rs` with
 * `#[repr(C)]`) ensures both kernels see the same byte layout. */
struct SP20EmaInputs {
    int   is_close;
    /* SP21 T3.1 (2026-05-10): fractional win rate ∈ [0, 1] = wins_count /
     * closed_count when is_close == 1, else 0.0f. Replaces the prior
     * binary `int is_win` (majority-vote) field which structurally pinned
     * WR_EMA at 0 whenever true WR < 0.5 (the binary "majority won this
     * step" indicator was 0 most of the time when actual win rate was
     * < 50%). The EMA blends over this fractional value so WR_EMA
     * converges to the true population win rate. */
    float win_fraction;
    int   trade_duration;
    /* SP21 T3.2 (2026-05-10): fractional hold rate ∈ [0, 1] = hold_count /
     * n_envs. Replaces the prior binary `int action_is_hold` (strict
     * majority vote) which structurally pinned HOLD_PCT_EMA at 0 whenever
     * fewer than half the envs were in Hold (the canonical case during
     * directional trading). HOLD_REWARD_EMA's gate preserves the strict-
     * majority semantic via `hold_fraction > 0.5f` test in the consumer. */
    float hold_fraction;
    float alpha;
    float per_bar_hold_reward;
    float aux_logits_p50;
    float aux_logits_std;
    float aux_dir_acc;
};

#ifndef SP20_AGG_BLOCK
#define SP20_AGG_BLOCK 64
#endif

extern "C" __global__ void sp20_aggregate_inputs_kernel(
    /* Per-env input arrays (read at offset env_id for env_id in [0, n_envs)).
     * Each is a slice of the [N, L] per-sample buffers at the current
     * timestep — the Rust launcher passes a pointer offset by
     * `env_offset_for_step = current_t` (since the per-sample layout is
     * `[i, t]` flat with stride L between envs, we pass the full base
     * pointer + a stride-L iteration through env_id below). */
    const int* __restrict__ trade_close_per_env,    /* [n_envs * stride] */
    const float* __restrict__ step_ret_per_env,     /* [n_envs * stride] */
    const float* __restrict__ hold_at_exit_per_env, /* [n_envs * stride] */
    const int* __restrict__ actions_per_env,        /* [n_envs * stride] */
    /* Stride between consecutive envs in the four per-env arrays. The
     * `experience_env_step` flat layout is `out_off = env*L + t`, so
     * the launcher passes `stride = L` (alloc_timesteps) and adds
     * `current_t` to each base pointer before launch — i.e., the
     * arrays here are already pointing at the slice for the current
     * timestep, and we stride by L between envs. */
    int env_stride,
    /* Upstream sp20_stats_compute outputs — single-element device
     * pointers for the median / std of aux_logits at this step. */
    const float* __restrict__ aux_logits_p50_dev,    /* [1] */
    const float* __restrict__ aux_logits_std_dev,    /* [1] */
    /* Existing aux_dir_acc_reduce_kernel output [0] — the directional
     * accuracy of the aux head this step. */
    const float* __restrict__ aux_dir_acc_dev,       /* [1] */
    /* SP20 Phase 2 Task 2.2 (2026-05-09) — per-env trade-close alpha
     * input. `[n_envs]` f32 device buffer; the env's slot is read only
     * when `trade_close_per_env[env * env_stride] != 0` so non-close-
     * step writes don't reach the EMA. NULL-tolerant — when NULL the
     * kernel emits `out_inputs->alpha = 0.0f` (matches Phase 1.4
     * behavior, used by tests that haven't wired the alpha producer). */
    const float* __restrict__ alpha_per_env,         /* [n_envs] or NULL */
    /* SP20 Phase 2 Task 2.2-fix (2026-05-10) — per-env trade-close
     * win flag. `[n_envs]` i32 device buffer; the env's slot is read
     * only when `trade_close_per_env[env * env_stride] != 0`. Each
     * slot is `1` if the segment's realized P&L was strictly positive
     * at trade close, else `0`. Replaces the broken `sr > 0` per-bar
     * step-return predicate (which conflates the close bar's per-bar
     * MTM tick with trade outcome — at close bars the per-bar tick is
     * dominated by tx-cost so the `sr > 0` predicate is essentially
     * always false, pinning WR_EMA at 0).
     *
     * NULL-tolerant — when NULL the kernel falls back to the legacy
     * `step_ret_per_env[env] > 0` predicate (matches the pre-2.2-fix
     * behavior used by aggregate-kernel oracle tests that don't wire
     * the segment-level producer). Production callers pass non-NULL
     * with the per-env buffer populated by `experience_env_step`'s
     * segment_complete branch as `(segment_return > 0.0f) ? 1 : 0`. */
    const int* __restrict__ is_win_per_env,          /* [n_envs] or NULL */
    /* SP21 T3.3 (2026-05-10) — per-env per-bar Hold opportunity-cost
     * input. `[n_envs]` f32 device buffer; written every step by
     * `experience_env_step` as `per_bar_opp_cost = -aux_conf × cost_scale ∈
     * [-0.25, 0]`. The aggregator reads `per_bar_opp_cost_per_env[env]`
     * only for envs in `dir == hold_action_id` (Hold-direction envs) so
     * non-Hold envs' opp_costs don't contaminate the average. The
     * downstream HOLD_REWARD_EMA's gate (`hold_fraction > 0.5f` at the
     * EMA kernel) ensures the EMA fires only when majority Hold —
     * preserving the strict-majority semantic of the pre-T3.2 binary
     * `action_is_hold` field.
     *
     * NULL-tolerant — when NULL the kernel emits
     * `out_inputs->per_bar_hold_reward = 0.0f` (matches the pre-T3.3
     * Phase 3.2-forward-reference contract used by aggregate-kernel
     * oracle tests that don't wire the per-bar producer). */
    const float* __restrict__ per_bar_opp_cost_per_env, /* [n_envs] or NULL */
    int n_envs,
    /* Action decoding parameters (production packed factored action
     * encoding). `dir_divisor = NUM_MAGNITUDES * NUM_ORD * NUM_URG`,
     * `hold_action_id = DIR_HOLD = 1`. */
    int dir_divisor,
    int hold_action_id,
    /* Output struct (single 36-byte struct on device; the EMA kernel
     * reads it as `SP20EmaInputs* ema_inputs`). */
    SP20EmaInputs* __restrict__ out_inputs)
{
    /* Single-block contract — guard against accidental multi-block
     * launches that would corrupt the shmem reductions. */
    if (blockIdx.x != 0) return;

    const int tid  = threadIdx.x;
    const int bdim = blockDim.x;

    /* Dynamic shared memory: 6 stripes of bdim × 4 bytes each (5 stripes
     * pre-T3.3 + 1 stripe for the SP21 T3.3 per-bar Hold opp-cost sum).
     * We use `int` storage for the 3 count reductions and `float` storage
     * for the hold-at-exit, alpha, and opp-cost sums (alias-compatible
     * because both are 4 bytes wide; we cast between views as needed). */
    extern __shared__ int sh_int[];
    /* Layout: [closed_count | wins_count | hold_count |
     *          hold_at_exit_sum_f32 | alpha_sum_f32 | opp_cost_sum_f32] */
    int*   sh_closed_count     = &sh_int[0];
    int*   sh_wins_count       = &sh_int[bdim];
    int*   sh_hold_count       = &sh_int[2 * bdim];
    float* sh_hold_at_exit_sum = (float*)&sh_int[3 * bdim];
    float* sh_alpha_sum        = (float*)&sh_int[4 * bdim];
    float* sh_opp_cost_sum     = (float*)&sh_int[5 * bdim];

    /* Per-thread strided accumulation over [0, n_envs). */
    int   local_closed_count   = 0;
    int   local_wins_count     = 0;
    int   local_hold_count     = 0;
    float local_hold_at_exit_sum = 0.0f;
    float local_alpha_sum        = 0.0f;
    float local_opp_cost_sum     = 0.0f;

    for (int env = tid; env < n_envs; env += bdim) {
        const int   tc = trade_close_per_env[env * env_stride];
        const float sr = step_ret_per_env[env * env_stride];
        const float he = hold_at_exit_per_env[env * env_stride];
        const int   action_idx = actions_per_env[env * env_stride];
        /* dir = action_idx / dir_divisor — same as
         * `hold_rate_observer_kernel.cu`'s decoding. */
        const int   dir = action_idx / dir_divisor;

        if (tc != 0) {
            local_closed_count += 1;
            /* SP20 Phase 2 Task 2.2-fix (2026-05-10): wins predicate
             * derives from the segment-level realized-return sign
             * via the per-env i32 `is_win_per_env` buffer (1 iff
             * segment_return > 0 at trade close), populated by
             * experience_env_step's segment_complete branch alongside
             * `alpha_per_env`. The legacy per-bar `sr > 0` predicate
             * is preserved as the NULL-tolerant fallback for oracle-
             * test scaffolds that don't wire a segment-level producer.
             * Bug class: per-bar step_ret at close bars is
             * `position*(close_t-close_{t-1}) - tx_cost`; the per-bar
             * tick is small while tx_cost is fixed → `sr > 0` is
             * almost always false at close bars even when the trade
             * closed profitably overall, pinning WR_EMA at 0. */
            int seg_win;
            if (is_win_per_env != NULL) {
                seg_win = (is_win_per_env[env] != 0) ? 1 : 0;
            } else {
                seg_win = (sr > 0.0f) ? 1 : 0;
            }
            local_wins_count += seg_win;
            local_hold_at_exit_sum += he;
            /* SP20 Phase 2 Task 2.2: gated alpha read on closed-env
             * slots only. Non-close-step alpha_per_env values are
             * stale (or sentinel 0 pre-first-close) and MUST NOT
             * reach the ALPHA_EMA. NULL-tolerant — when alpha_per_env
             * is NULL, the read is skipped and `local_alpha_sum`
             * stays at 0 → `out_inputs->alpha = 0.0f`, matching the
             * Phase 1.4 placeholder behavior used by tests without
             * the alpha producer wired. */
            if (alpha_per_env != NULL) {
                local_alpha_sum += alpha_per_env[env];
            }
        }
        if (dir == hold_action_id) {
            local_hold_count += 1;
            /* SP21 T3.3 (2026-05-10): accumulate per-bar Hold opp-cost
             * for Hold-direction envs only. NULL-tolerant — when the
             * producer isn't wired (oracle tests / pre-T3.3 callers),
             * the accumulator stays at 0 → `out_inputs->per_bar_hold_reward
             * = 0.0f`, matching the Phase 3.2 forward-reference contract. */
            if (per_bar_opp_cost_per_env != NULL) {
                local_opp_cost_sum += per_bar_opp_cost_per_env[env];
            }
        }
    }

    sh_closed_count[tid]      = local_closed_count;
    sh_wins_count[tid]        = local_wins_count;
    sh_hold_count[tid]        = local_hold_count;
    sh_hold_at_exit_sum[tid]  = local_hold_at_exit_sum;
    sh_alpha_sum[tid]         = local_alpha_sum;
    sh_opp_cost_sum[tid]      = local_opp_cost_sum;
    __syncthreads();

    /* Tree reduction across `bdim`. Standard log2(bdim) pattern. We
     * reduce all six stripes in one loop body — single barrier per
     * pass keeps the cost equal to a single reduction. */
    for (int s = bdim / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sh_closed_count[tid]     += sh_closed_count[tid + s];
            sh_wins_count[tid]       += sh_wins_count[tid + s];
            sh_hold_count[tid]       += sh_hold_count[tid + s];
            sh_hold_at_exit_sum[tid] += sh_hold_at_exit_sum[tid + s];
            sh_alpha_sum[tid]        += sh_alpha_sum[tid + s];
            sh_opp_cost_sum[tid]     += sh_opp_cost_sum[tid + s];
        }
        __syncthreads();
    }

    /* Thread 0 writes the output struct. */
    if (tid == 0) {
        const int   closed_count       = sh_closed_count[0];
        const int   wins_count         = sh_wins_count[0];
        const int   hold_count         = sh_hold_count[0];
        const float hold_at_exit_sum_f = sh_hold_at_exit_sum[0];
        const float alpha_sum_f        = sh_alpha_sum[0];
        const float opp_cost_sum_f     = sh_opp_cost_sum[0];

        const int is_close_out = (closed_count > 0) ? 1 : 0;

        /* SP21 T3.1 (2026-05-10): fractional win rate, NOT binary majority.
         * `win_fraction = wins_count / closed_count ∈ [0, 1]`. The prior
         * binary majority-vote `(2 * wins_count >= closed_count)` predicate
         * was structurally incapable of representing WR < 0.5 — it emitted
         * 0 whenever fewer than half the closed envs won, regardless of
         * the actual ratio, pinning WR_EMA at 0 in production whenever
         * true WR < 0.5 (the canonical case). The downstream EMA blends
         * over this float directly so WR_EMA converges to the population
         * win rate.
         *
         * Emit 0.0f when `is_close_out == 0` for cleanliness — the consumer
         * gates the blend on `is_close == 1` so the value is moot in that
         * branch. */
        float win_fraction_out = 0.0f;
        if (is_close_out == 1) {
            win_fraction_out = (float)wins_count / (float)closed_count;
        }

        /* trade_duration: round(mean of hold_at_exit over closed envs). */
        int trade_duration_out = 0;
        if (is_close_out == 1) {
            const float mean_he = hold_at_exit_sum_f / (float)closed_count;
            trade_duration_out = (int)__float2int_rn(mean_he);
        }

        /* SP20 Phase 2 Task 2.2: alpha = mean(alpha_per_env) over closed
         * envs. Replaces the Phase 1.4 hardcoded `0.0f` placeholder
         * (audit doc Task 2.2 forward-reference closure). When
         * `is_close_out == 0` (no env closed this step) the EMA gate
         * is closed at the consumer, so this slot's value is moot — but
         * we still emit 0.0 for cleanliness so downstream debug reads
         * don't see stale sums from the reduction. NULL-tolerant via
         * the per-thread accumulation gate above. */
        float alpha_out = 0.0f;
        if (is_close_out == 1) {
            alpha_out = alpha_sum_f / (float)closed_count;
        }

        /* SP21 T3.2 (2026-05-10): fractional hold rate, NOT binary majority.
         * `hold_fraction = hold_count / n_envs ∈ [0, 1]`. The prior strict-
         * majority `(hold_count * 2 > n_envs)` predicate was structurally
         * incapable of representing hold rates < 0.5 — it emitted 0 whenever
         * fewer than half the envs were in Hold, pinning HOLD_PCT_EMA at 0
         * during directional trading and cascade-pinning HOLD_COST_SCALE at
         * its floor (the controller saw `hold_pct_ema < lower` every step
         * and ramped × 0.95 to clamp). Downstream HOLD_REWARD_EMA gate
         * preserves the strict-majority semantic via `hold_fraction > 0.5f`. */
        const float hold_fraction_out = (n_envs > 0)
            ? (float)hold_count / (float)n_envs
            : 0.0f;

        /* Read the upstream stats / dir-acc outputs. Stream-ordering
         * invariant guarantees these are visible to this kernel's
         * read because the launcher places sp20_stats_compute and
         * aux_dir_acc_reduce BEFORE this aggregation kernel on the
         * same stream (see Phase 1.4 collector wire-up). */
        const float p50_in     = aux_logits_p50_dev[0];
        const float std_in     = aux_logits_std_dev[0];
        const float dir_acc_in = aux_dir_acc_dev[0];

        /* Write the 9 fields. `alpha` is now the load-bearing SP20
         * Phase 2 Task 2.2 producer output (mean alpha over closed
         * envs); `per_bar_hold_reward` remains a Phase 3.2 forward
         * reference until Component 2 lands the per-bar Hold opp-cost
         * dual emission. */
        out_inputs->is_close            = is_close_out;
        out_inputs->win_fraction        = win_fraction_out;
        out_inputs->trade_duration      = trade_duration_out;
        out_inputs->hold_fraction       = hold_fraction_out;  /* SP21 T3.2 */
        out_inputs->alpha               = alpha_out;       /* Task 2.2 — was 0.0f */
        /* SP21 T3.3 (2026-05-10): per_bar_hold_reward = mean per-bar Hold
         * opp-cost over Hold-direction envs at this step. Replaces the
         * Phase 3.2 forward-reference `0.0f` placeholder. The accumulator
         * `opp_cost_sum_f` only summed over envs with `dir == hold_action_id`
         * (per-thread gate above), so dividing by `hold_count` gives the
         * Hold-conditional mean. When no envs are in Hold OR the producer
         * is NULL (oracle tests), the sum is 0 and we emit 0.0f matching
         * pre-T3.3 contract. The downstream HOLD_REWARD_EMA's gate
         * (`hold_fraction > 0.5f` at the EMA kernel) ensures the EMA fires
         * only when majority Hold — preserving strict-majority semantic. */
        float per_bar_hold_reward_out = 0.0f;
        if (hold_count > 0) {
            per_bar_hold_reward_out = opp_cost_sum_f / (float)hold_count;
        }
        out_inputs->per_bar_hold_reward = per_bar_hold_reward_out;
        out_inputs->aux_logits_p50      = p50_in;
        out_inputs->aux_logits_std      = std_in;
        out_inputs->aux_dir_acc         = dir_acc_in;

        __threadfence_system();
    }
}
