/**
 * rl_per_push_flush.cu — Coordinated coalesced replay write.
 *
 * Second kernel in the split rl_per_push pipeline. After rl_per_push_ring has
 * written data to the n-step ring and set flush_flags[b] for completed n-step
 * transitions, this kernel coordinates slot allocation and performs the actual
 * coalesced write to the replay buffer.
 *
 * Architecture:
 *   - Grid=(b_size, 1, 1), Block=(128, 1, 1). One block per batch element.
 *   - Block 0 runs the prefix-sum to assign write slots (single-threaded, O(B)).
 *   - All other blocks spin on a volatile global-memory ready flag (no atomics).
 *   - Once ready, each flushing block does a coalesced 128-thread copy of h_t
 *     and h_tp1, then thread 0 computes the n-step return and writes scalars +
 *     priority + resets the ring count.
 *
 * Coordination: block-0 prefix-sum + volatile ready flag.
 * No atomicAdd — slot assignment is deterministic serial scan (feedback_no_atomicadd).
 * Pre-compiled cubin only (feedback_no_nvrtc).
 *
 * Launch: CUDA_HOME=/usr/local/cuda nvcc -cubin -arch sm_86 -O3 \
 *         --generate-line-info rl_per_push_flush.cu -o rl_per_push_flush.cubin
 */

#define HIDDEN_DIM       128
#define N_STEP_MAX       16
#define SCALARS_PER_TRANSITION 7

/* ISV slot indices for RL config */
#define RL_GAMMA_INDEX   400
#define RL_N_STEP_INDEX  403

extern "C" __global__ void rl_per_push_flush(
    /* ── Flush coordination ── */
    const int* __restrict__ flush_flags,              /* [B] from ring kernel: 1=flush, 0=skip */
    int* __restrict__ write_offsets,                  /* [B+2] output: per-batch offsets, base, ready flag */

    /* ── N-step ring (read source) ── */
    const float* __restrict__ nstep_scalars,          /* [B * N_STEP_MAX * 3]: (r_scaled, r_raw, done) per ring slot */
    const float* __restrict__ nstep_h_t,              /* [B * N_STEP_MAX * HIDDEN_DIM] */
    const unsigned int* __restrict__ nstep_write_idx, /* [B] ring write cursor */
    const unsigned int* __restrict__ nstep_count,     /* [B] items in ring */

    /* ── Current step data (for h_tp1) ── */
    const float* __restrict__ h_tp1_current,          /* [B * HIDDEN_DIM] */
    const float* __restrict__ log_pi_old,             /* [B] */
    const int*   __restrict__ actions,                /* [B] */
    const float* __restrict__ isv,                    /* ISV array */

    /* ── Replay storage (write destination) ── */
    float* __restrict__ replay_h_t,                   /* [capacity * HIDDEN_DIM] */
    float* __restrict__ replay_h_tp1,                 /* [capacity * HIDDEN_DIM] */
    float* __restrict__ replay_scalars,               /* [capacity * SCALARS_PER_TRANSITION] */
    float* __restrict__ priority_tree,                /* [2 * capacity] (leaf layer at offset capacity) */
    unsigned int* __restrict__ write_head,            /* [1] circular write cursor */
    unsigned int* __restrict__ replay_len,            /* [1] current buffer occupancy */
    float* __restrict__ max_priority,                 /* [1] running max priority for new inserts */

    /* ── N-step count reset (written on flush) ── */
    unsigned int* __restrict__ nstep_count_out,       /* [B] — set to 0 for flushed batches */

    int b_size,
    int capacity
)
{
    /* ====================================================================
     * PHASE 1: Block 0 coordination (thread 0 only).
     *
     * Computes prefix-sum of flush_flags to assign contiguous write slots,
     * advances write_head and replay_len, stores base for other blocks,
     * and signals ready via volatile write.
     * ==================================================================== */
    if (blockIdx.x == 0 && threadIdx.x == 0) {
        unsigned int base = write_head[0];

        /* Exclusive prefix-sum: write_offsets[i] = number of flushes before batch i */
        int total = 0;
        for (int i = 0; i < b_size; ++i) {
            write_offsets[i] = total;
            total += flush_flags[i];
        }

        /* Advance write head and replay length */
        if (total > 0) {
            write_head[0] = (base + (unsigned int)total) % (unsigned int)capacity;
            unsigned int len = replay_len[0] + (unsigned int)total;
            if (len > (unsigned int)capacity) len = (unsigned int)capacity;
            replay_len[0] = len;
        }

        /* Store base write head for other blocks at slot [b_size] */
        write_offsets[b_size] = (int)base;

        /* Fence: ensure all writes above are globally visible before signaling */
        __threadfence();

        /* Signal ready at slot [b_size + 1] */
        write_offsets[b_size + 1] = 1;
        __threadfence();
    }

    /* ====================================================================
     * PHASE 2: All blocks wait for ready signal.
     *
     * Non-block-0 blocks spin on a volatile read of the ready flag.
     * Block 0 already wrote it, so it proceeds immediately.
     * ==================================================================== */
    if (threadIdx.x == 0 && blockIdx.x != 0) {
        volatile int* ready = (volatile int*)&write_offsets[b_size + 1];
        while (*ready == 0) {
            /* spin — volatile load, no atomic */
        }
    }
    __syncthreads();

    /* ====================================================================
     * PHASE 3: Coalesced replay write (128 threads per block).
     *
     * Each flushing block writes:
     *   - h_t[HIDDEN_DIM]: from oldest ring entry (coalesced, all 128 threads)
     *   - h_tp1[HIDDEN_DIM]: from current step (coalesced, all 128 threads)
     *   - scalars[7]: n-step return + metadata (thread 0 only)
     *   - priority: max_priority into leaf layer (thread 0 only)
     *   - reset: nstep_count_out[b] = 0 (thread 0 only)
     * ==================================================================== */
    const int b = blockIdx.x;

    /* Early exit: nothing to flush for this batch */
    if (flush_flags[b] == 0) return;

    /* Compute target replay slot */
    const unsigned int base = (unsigned int)write_offsets[b_size];
    const unsigned int my_offset = (unsigned int)write_offsets[b];
    const unsigned int my_slot = (base + my_offset) % (unsigned int)capacity;
    const int tid = threadIdx.x;

    /* Read n-step config from ISV */
    int n_step_raw = (int)isv[RL_N_STEP_INDEX];
    const int n_step = (n_step_raw < 1) ? 1 : ((n_step_raw > N_STEP_MAX) ? N_STEP_MAX : n_step_raw);
    const float gamma = isv[RL_GAMMA_INDEX];

    /* Determine oldest ring entry to read h_t from */
    const unsigned int count = nstep_count[b];
    const int actual_n = ((int)count < n_step) ? (int)count : n_step;
    const unsigned int oldest_ring = (nstep_write_idx[b] + (unsigned int)N_STEP_MAX - (unsigned int)actual_n)
                                     % (unsigned int)N_STEP_MAX;

    /* Coalesced h_t copy: oldest ring entry -> replay (128 threads, 1 float each) */
    replay_h_t[my_slot * HIDDEN_DIM + tid] =
        nstep_h_t[(b * N_STEP_MAX + oldest_ring) * HIDDEN_DIM + tid];

    /* Coalesced h_tp1 copy: current hidden state -> replay (128 threads, 1 float each) */
    replay_h_tp1[my_slot * HIDDEN_DIM + tid] =
        h_tp1_current[b * HIDDEN_DIM + tid];

    /* Thread 0: compute n-step return, write scalars, set priority, reset count */
    if (tid == 0) {
        /* N-step discounted return computation */
        float r_scaled = 0.0f;
        float r_raw = 0.0f;
        float gamma_pow = 1.0f;
        float n_step_gamma = 1.0f;
        int any_done = 0;

        for (int k = 0; k < actual_n; ++k) {
            unsigned int rk = (oldest_ring + (unsigned int)k) % (unsigned int)N_STEP_MAX;
            int sc_idx = (b * N_STEP_MAX + (int)rk) * 3;
            r_scaled += gamma_pow * nstep_scalars[sc_idx + 0];
            r_raw += gamma_pow * nstep_scalars[sc_idx + 1];
            if (nstep_scalars[sc_idx + 2] > 0.5f) {
                any_done = 1;
                n_step_gamma = 0.0f;
                break;
            }
            gamma_pow *= gamma;
        }
        if (!any_done) {
            n_step_gamma = gamma_pow; /* gamma^actual_n for bootstrap */
        }

        /* Write transition scalars to replay buffer */
        int sc_out = (int)my_slot * SCALARS_PER_TRANSITION;
        replay_scalars[sc_out + 0] = (float)actions[b];       /* action index */
        replay_scalars[sc_out + 1] = r_scaled;                /* n-step scaled return */
        replay_scalars[sc_out + 2] = r_raw;                   /* n-step raw return */
        replay_scalars[sc_out + 3] = (float)any_done;         /* terminal flag */
        replay_scalars[sc_out + 4] = log_pi_old[b];           /* log pi(a|s) at collection */
        replay_scalars[sc_out + 5] = n_step_gamma;            /* discount for bootstrap */
        replay_scalars[sc_out + 6] = r_raw;                   /* raw return (duplicate for diagnostics) */

        /* Set priority to current max (standard PER new-insert convention).
         * priority_tree leaf layer starts at offset capacity. */
        float max_p = max_priority[0];
        if (max_p < 1e-9f) max_p = 1.0f;
        priority_tree[capacity + (int)my_slot] = max_p;

        /* Reset n-step count for this batch (ring consumed) */
        nstep_count_out[b] = 0;
    }
}
