// crates/ml/src/cuda_pipeline/dd_state_reduce_kernel.cu
//
// SP15 Phase 1.3.b-followup (2026-05-07) — DD state reduction kernel.
// SP15 Phase 1.3.b-followup-B (2026-05-07) — extended with per-env DD
// trajectory mean-aggregate for HEALTH_DIAG ISV[439] preservation.
//
// Aggregates the per-env DD tile written by `dd_state_kernel` into the
// 6 scalar ISV slots [401..407) plus the new `DD_PERSISTENCE_MAX_INDEX
// = 443` slot, AND the per-env trajectory output written by
// `dd_trajectory_decreasing_kernel` into ISV[DD_TRAJECTORY_DECREASING_
// INDEX = 439] (mean across envs — HEALTH_DIAG diagnostic preservation
// while production reads the per-env tile via insert-time
// `env_id = (j % (n_envs * lookback)) / lookback` lookup in
// `per_insert_pa`). Three distinct aggregation rules:
//
//   * Mean across envs for slots 401..407): smooth across-env summary
//     for HEALTH_DIAG diagnostic reporting. Preserves backward-
//     compatible "this is one DD scalar to look at" semantics for the
//     diagnostic stream — only the across-env aggregation rule
//     changed (Path A: env-0-canonical → Path B: mean across envs).
//     This is structural per `feedback_isv_for_adaptive_bounds` (the
//     aggregation rule itself is fixed mean — no adaptive smoothing
//     here).
//   * Mean across envs for slot 439 (DD_TRAJECTORY_DECREASING):
//     HEALTH_DIAG diagnostic preservation per Phase 1.3.b-followup-B.
//     The actual production consumer (`per_insert_pa`) reads the
//     per-env tile directly, NOT this scalar — so the slot's role is
//     diagnostic-only post-followup-B. Mean-aggregate matches the
//     other 6 ISV scalars' rule (smooth across-env summary).
//   * Max across envs for `DD_PERSISTENCE_MAX_INDEX = 443`: consumed
//     by `plasticity_injection_kernel` so the global advantage-head
//     reset fires when ANY env's persistence exceeds the threshold.
//     One set of advantage weights → global firing semantics → max-
//     aggregate is the only correct rule (mean-aggregate would silently
//     gate the trigger when only a fraction of envs are in long DD).
//
// Single-block tree-reduce (no atomicAdd per `feedback_no_atomicadd` —
// the existing block-tree-reduce pattern in `kelly_cap_update` /
// `cost_net_sharpe` baseline kernels). Block size `BLOCK = 256`; if
// `n_envs > 256` the reduction does a strided initial accumulation
// pass before the tree-reduce (the production sizes
// `optimal_n_episodes()` reach 32768 envs on H100, so this matters).
// `n_envs` is a launch-time scalar argument — fixed across the run, so
// the kernel grid shape is constant per `pearl_no_host_branches_in_
// captured_graph` (CUDA Graph capture remains compatible).
//
// `dd_trajectory_per_env` may be nullptr (test scaffold path that
// exercises the DD state reduction without the per-env trajectory
// buffer wired). When null, ISV[439] is left untouched (preserves any
// prior write by the kernel's caller — typically the FoldReset
// sentinel 0). Mirrors the `per_insert_pa` nullable-pointer pattern.
//
// Layout contract (the per-env tile offsets MUST mirror
// `dd_state_kernel.cu`'s writes; layout-fingerprint marker
// `DD_STATE_PER_ENV=sp15_phase_1_3_b_followup` in
// `layout_fingerprint_seed` enforces synchrony):
//   tile[env*6 + 0] = DD_CURRENT          → ISV[401] (mean)
//   tile[env*6 + 1] = DD_MAX              → ISV[402] (mean)
//   tile[env*6 + 2] = DD_RECOVERY_BARS    → ISV[403] (mean)
//   tile[env*6 + 3] = DD_PERSISTENCE      → ISV[404] (mean)
//                                         → ISV[443] (max — DD_PERSISTENCE_MAX)
//   tile[env*6 + 4] = CALMAR              → ISV[405] (mean)
//   tile[env*6 + 5] = DD_PCT              → ISV[406] (mean)
//   dd_trajectory_per_env[env]            → ISV[439] (mean — HEALTH_DIAG)
//
// Hard rules:
//   * `feedback_no_atomicadd` — single-block tree-reduce, no atomicAdd.
//   * `feedback_no_partial_refactor` — kernel + 5 downstream consumer
//     migrations land atomically with the per-env tile contract switch
//     (Phase 1.3.b-followup); per-env trajectory mean-aggregate lands
//     atomically with the per-env trajectory tile + per_insert_pa
//     migration (Phase 1.3.b-followup-B).
//   * `feedback_no_stubs` — kernel writes the actual mean+max
//     aggregates; no return-zero stubs.
//   * `feedback_isv_for_adaptive_bounds` — mean / max aggregation rules
//     are structural (fixed by the aggregation requirement: one shared
//     global advantage-head ⇒ max-aggregate for trigger, smooth
//     diagnostic ⇒ mean-aggregate for HEALTH_DIAG). Documented in
//     `docs/dqn-wire-up-audit.md` Phase 1.3.b-followup + 1.3.b-followup-B
//     entries.
//   * `pearl_no_host_branches_in_captured_graph` — `n_envs` is a kernel
//     launch arg (a scalar, not a host-side branch); kernel grid
//     shape is constant.

#define BLOCK_SIZE 256

extern "C" __global__ void dd_state_reduce_kernel(
    int                       n_envs,
    const float* __restrict__ dd_state_per_env,        /* [n_envs * 6] tile input  */
    const float* __restrict__ dd_trajectory_per_env,   /* [n_envs] (nullable) trajectory tile */
    float* __restrict__       isv                      /* [≥444] ISV bus            */
) {
    // Single-block (gridDim.x == 1 enforced by the launcher's
    // `grid_dim: (1, 1, 1)`).
    const int tid = threadIdx.x;

    // Per-thread accumulators (one per output statistic).
    // Six mean-aggregates (one per tile field) + one max-aggregate (for
    // persistence) + one mean-aggregate for trajectory. Initialise to
    // identity (0 for sums, -inf for max).
    float sum_dd_current      = 0.0f;
    float sum_dd_max          = 0.0f;
    float sum_recovery_bars   = 0.0f;
    float sum_persistence     = 0.0f;
    float sum_calmar          = 0.0f;
    float sum_dd_pct          = 0.0f;
    float max_persistence     = -1.0f;  // persistence is non-negative; -1 < any valid value
    float sum_trajectory      = 0.0f;

    /* Strided initial accumulation pass: each thread sums every (n_envs /
     * BLOCK_SIZE)-th env starting from `tid`. For n_envs ≤ BLOCK_SIZE
     * each thread handles at most one env (with the remaining threads
     * leaving the identity initialisers untouched). For n_envs > BLOCK_SIZE
     * the strided loop accumulates many envs into per-thread state
     * before the tree-reduce. */
    for (int env_id = tid; env_id < n_envs; env_id += BLOCK_SIZE) {
        const int base = env_id * 6;
        sum_dd_current     += dd_state_per_env[base + 0];
        sum_dd_max         += dd_state_per_env[base + 1];
        sum_recovery_bars  += dd_state_per_env[base + 2];
        const float persistence = dd_state_per_env[base + 3];
        sum_persistence    += persistence;
        max_persistence     = fmaxf(max_persistence, persistence);
        sum_calmar         += dd_state_per_env[base + 4];
        sum_dd_pct         += dd_state_per_env[base + 5];
        if (dd_trajectory_per_env != nullptr) {
            sum_trajectory += dd_trajectory_per_env[env_id];
        }
    }

    /* Block-shared tree-reduce. We reduce 8 quantities; allocate a
     * `[BLOCK_SIZE × 8]` shared-memory tile so each output keeps its own
     * lane during the reduction. Layout (column-major over output):
     *   sdata[k * BLOCK_SIZE + tid] for output k ∈ [0..8).
     */
    __shared__ float sdata[BLOCK_SIZE * 8];

    sdata[0 * BLOCK_SIZE + tid] = sum_dd_current;
    sdata[1 * BLOCK_SIZE + tid] = sum_dd_max;
    sdata[2 * BLOCK_SIZE + tid] = sum_recovery_bars;
    sdata[3 * BLOCK_SIZE + tid] = sum_persistence;
    sdata[4 * BLOCK_SIZE + tid] = sum_calmar;
    sdata[5 * BLOCK_SIZE + tid] = sum_dd_pct;
    sdata[6 * BLOCK_SIZE + tid] = max_persistence;
    sdata[7 * BLOCK_SIZE + tid] = sum_trajectory;
    __syncthreads();

    /* Tree-reduce: halving stride, sum for outputs 0..6 + 7 and max for
     * output 6 (persistence-max). The branchy max vs sum pattern fits in
     * one loop because we hard-code which lane uses which combiner. */
    for (int stride = BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[0 * BLOCK_SIZE + tid] += sdata[0 * BLOCK_SIZE + tid + stride];
            sdata[1 * BLOCK_SIZE + tid] += sdata[1 * BLOCK_SIZE + tid + stride];
            sdata[2 * BLOCK_SIZE + tid] += sdata[2 * BLOCK_SIZE + tid + stride];
            sdata[3 * BLOCK_SIZE + tid] += sdata[3 * BLOCK_SIZE + tid + stride];
            sdata[4 * BLOCK_SIZE + tid] += sdata[4 * BLOCK_SIZE + tid + stride];
            sdata[5 * BLOCK_SIZE + tid] += sdata[5 * BLOCK_SIZE + tid + stride];
            sdata[6 * BLOCK_SIZE + tid] = fmaxf(
                sdata[6 * BLOCK_SIZE + tid],
                sdata[6 * BLOCK_SIZE + tid + stride]
            );
            sdata[7 * BLOCK_SIZE + tid] += sdata[7 * BLOCK_SIZE + tid + stride];
        }
        __syncthreads();
    }

    /* Thread 0 writes the final reduced values to the ISV bus. Mean
     * aggregates are sum / n_envs (n_envs ≥ 1 enforced by the launcher
     * — single-env smoke configs still pass n_envs = 1). The CALMAR
     * slot remains a denominator floor (mean of `max(dd_max_i, 1e-4)`
     * stays ≥ 1e-4 because every term is ≥ 1e-4, so the mean cannot
     * drop below the floor). */
    if (tid == 0) {
        const float inv_n = 1.0f / (float)n_envs;
        isv[/*DD_CURRENT_INDEX*/        401] = sdata[0 * BLOCK_SIZE] * inv_n;
        isv[/*DD_MAX_INDEX*/            402] = sdata[1 * BLOCK_SIZE] * inv_n;
        isv[/*DD_RECOVERY_BARS_INDEX*/  403] = sdata[2 * BLOCK_SIZE] * inv_n;
        isv[/*DD_PERSISTENCE_INDEX*/    404] = sdata[3 * BLOCK_SIZE] * inv_n;
        isv[/*CALMAR_INDEX*/            405] = sdata[4 * BLOCK_SIZE] * inv_n;
        isv[/*DD_PCT_INDEX*/            406] = sdata[5 * BLOCK_SIZE] * inv_n;
        isv[/*DD_PERSISTENCE_MAX_INDEX*/ 443] = sdata[6 * BLOCK_SIZE];
        if (dd_trajectory_per_env != nullptr) {
            isv[/*DD_TRAJECTORY_DECREASING_INDEX*/ 439] = sdata[7 * BLOCK_SIZE] * inv_n;
        }
        __threadfence_system();
    }
}
