// crates/ml/src/cuda_pipeline/aux_outcome_softmax_to_per_env_kernel.cu
//
// SP22 H6 vNext Phase C (2026-05-14): aux→policy state bridge — per-env
// trade-outcome probability extractor.
//
// Role
// ────
// At rollout step t, after the K=3 trade-outcome aux head's softmax
// forward produces the per-env tile `aux_outcome_softmax [n_envs, K=3]`,
// this kernel copies all three probability columns into the per-env
// cache buffer `prev_aux_outcome_probs [n_envs, 3]`. The cache is read
// by `experience_state_gather` at step t+1 to populate state slots:
//
//   state[AUX_OUTCOME_PROFIT_INDEX  = 121] ← prev_aux_outcome_probs[env, 0]
//   state[AUX_OUTCOME_STOP_INDEX    = 122] ← prev_aux_outcome_probs[env, 1]
//   state[AUX_OUTCOME_TIMEOUT_INDEX = 123] ← prev_aux_outcome_probs[env, 2]
//
// Replaces the K=2 head's single-slot bridge
// (`aux_softmax_to_per_env_kernel` → `prev_aux_dir_prob` → state[121] =
// `2*p_up - 1`). The K=3 head's 3-slot bridge supersedes that semantic
// at step t+1 once the state gather consumer is updated (Phase C-2,
// follow-up commit). For Phase C (this commit) the buffer is populated
// but no consumer reads it yet — dead-code scaffolding.
//
// Encoding
// ────────
// Slots [121..124) hold raw softmax probabilities in [0, 1] (sum to 1
// per row). Unlike the K=2 head's slot 121 which used a recentered
// `2*p - 1` encoding to land in [-1, +1], the K=3 head's three slots
// keep the natural [0, 1] range:
//   - No structural bound violation (softmax outputs are inherently
//     [0, 1] and sum to 1)
//   - Three-slot probabilistic representation is more informative than
//     a single "p_up − p_down" scalar
//   - Cold-start sentinel: all three slots = 0.0. The policy reads
//     "no signal" for all three classes, which is the conservative
//     baseline (vs. the K=2 head's 0.0 = neutral encoding). After
//     first rollout step the producer overwrites with real K=3 probs.
//
// Discipline
// ──────────
// - No atomicAdd (no reduction, just a strided gather)
// - GPU-only per `feedback_cpu_is_read_only.md`
// - Stream-ordered with the K=3 forward producer (same-stream barrier
//   suffices — no explicit sync needed)
//
// Launch
// ──────
// grid = ((n_envs + 255) / 256), block = (256). One thread per env.
// Each thread writes 3 consecutive output elements (K=3 unrolled).
//
// Cost
// ────
// Three coalesced 4-byte loads + three 4-byte stores per env. n_envs is
// O(1k); ~1-2µs total. Negligible vs the existing K=2 sibling.

#include <cuda_runtime.h>

extern "C" __global__ void aux_outcome_softmax_to_per_env_kernel(
    /* Per-env K=3 softmax tile [n_envs, K=3]. Written by
     * `aux_trade_outcome_forward` during the captured rollout-step
     * graph; valid after same-stream sync. Layout row-major; this
     * kernel reads `aux_outcome_softmax[env * K + 0/1/2]` per env. */
    const float* __restrict__ aux_outcome_softmax,
    /* Per-env cache buffer [n_envs, K=3] row-major. Written here, read
     * by `experience_state_gather` (Phase C-2) at the next step's
     * state assembly. Slot 0 = p_Profit, slot 1 = p_Stop, slot 2 =
     * p_Timeout per the K=3 outcome label convention from
     * `trade_outcome_label_kernel`. */
    float*       __restrict__ prev_aux_outcome_probs,
    /* Number of envs (rollout-time batch width = `alloc_episodes`). */
    int n_envs,
    /* Softmax K (= AUX_OUTCOME_K = 3 in Phase A3+). Kernel reads slots
     * [0, K). Passed runtime so the kernel signature stays stable if
     * the head's K flips in a future redesign. */
    int K
) {
    int env = blockIdx.x * blockDim.x + threadIdx.x;
    if (env >= n_envs) return;

    /* Phase C (2026-05-14): write raw softmax probabilities [0, 1]
     * into 3 per-env slots. The state gather consumer (Phase C-2)
     * will copy these into state[121..124). Unlike the K=2 head's
     * recentered encoding, the K=3 head's natural softmax range is
     * already in [0, 1] and the three-slot representation preserves
     * the full categorical distribution.
     *
     * The kernel writes K slots (passed runtime) so the loop tolerates
     * future K flips; in production K=3 always so the loop trips
     * exactly 3 times. Bounds: K can be any positive value, but the
     * cache buffer must be sized [n_envs, K] by the caller. */
    const size_t src_base = (size_t)env * (size_t)K;
    const size_t dst_base = (size_t)env * (size_t)K;
    for (int k = 0; k < K; ++k) {
        prev_aux_outcome_probs[dst_base + (size_t)k] =
            aux_outcome_softmax[src_base + (size_t)k];
    }
}
