// crates/ml/src/cuda_pipeline/alpha_window_push.cu
//
// Phase E.4.A.6 (2026-05-15): shift+insert circular-buffer push for the
// sliding-window state input to the Mamba2 temporal encoder. Mirrors
// production's `mamba2_update_history` (mamba2_temporal_kernel.cu):
// shift slots 1..K-1 to slots 0..K-2, then insert new state at slot K-1.
//
// Layout invariant: window[0..K-1, :] is CHRONOLOGICAL — slot 0 is the
// oldest observation, slot K-1 is the newest. This matches Mamba2Block's
// [B, K, in_dim] input contract directly (no reorder needed).
//
// Threads: state_dim threads, one per feature. Each thread sequentially
// shifts its column's K-1 slots and writes the new value at slot K-1.
// For typical state_dim=10, K=16: 10 threads × ~17 ops each, single
// block, no data races (each thread owns its column).

#include <cuda_runtime.h>

extern "C" __global__ void alpha_window_push_kernel(
    const float* __restrict__ state_in,    // [state_dim] (mapped-pinned)
    float*       __restrict__ window,      // [K, state_dim] row-major
    int          K,                        // window length
    int          state_dim                 // typical 10
) {
    const int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= state_dim) return;
    // Shift slot (t+1) → slot t for t = 0..K-2.
    for (int t = 0; t < K - 1; ++t) {
        window[t * state_dim + j] = window[(t + 1) * state_dim + j];
    }
    // Insert new state at slot K-1.
    window[(K - 1) * state_dim + j] = state_in[j];
}

// ----------------------------------------------------------------------
// Phase E.4.A.8.fix (2026-05-15): GPU-side slot copy for h_enriched.
// ----------------------------------------------------------------------
//
// Replaces the dtoh/htod sequence that was used to store per-step
// Mamba2 h_enriched into a buffer slot — a CPU-roundtrip violation of
// `feedback_cpu_is_read_only`. One thread per hidden_dim feature.
extern "C" __global__ void alpha_h_enriched_store_kernel(
    const float* __restrict__ src,         // [hidden_dim]
    float*       __restrict__ buf,         // [(horizon+1) * hidden_dim]
    int          slot_offset,              // slot index × hidden_dim
    int          hidden_dim
) {
    const int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= hidden_dim) return;
    buf[slot_offset + j] = src[j];
}

// ----------------------------------------------------------------------
// Phase E.4.A.T14.batched (2026-05-15): batched shift+insert push for
// parallel environments at eval time. Same chronological semantics as
// the single-env kernel, but across B independent windows of shape
// [B, K, state_dim]. Grid: (state_dim_blocks, B, 1).
// ----------------------------------------------------------------------
extern "C" __global__ void alpha_window_push_batched_kernel(
    const float* __restrict__ states_in,   // [B, state_dim]
    float*       __restrict__ windows,     // [B, K, state_dim]
    int          B,
    int          K,
    int          state_dim
) {
    const int j = blockIdx.x * blockDim.x + threadIdx.x;
    const int b = blockIdx.y;
    if (b >= B || j >= state_dim) return;
    const int base = b * K * state_dim;
    for (int t = 0; t < K - 1; ++t) {
        windows[base + t * state_dim + j] = windows[base + (t + 1) * state_dim + j];
    }
    windows[base + (K - 1) * state_dim + j] = states_in[b * state_dim + j];
}

// ----------------------------------------------------------------------
// Phase E.4.A.T15.batched (2026-05-15): batched h_enriched store for
// training-time transition collection. Copies cache.h_enriched
// [B, hidden_dim] into buf[step_offset_in_rows .. step_offset+B] of
// a [H+1, B, hidden_dim] buffer (step-outer layout, contiguous block
// per step). One thread per (env, feature).
// ----------------------------------------------------------------------
extern "C" __global__ void alpha_h_enriched_store_batched_kernel(
    const float* __restrict__ src,         // [B, hidden_dim]
    float*       __restrict__ buf,         // [..., B, hidden_dim]
    int          step_row_offset,          // step * B  (in rows of width hidden_dim)
    int          B,
    int          hidden_dim
) {
    const int j = blockIdx.x * blockDim.x + threadIdx.x;
    const int b = blockIdx.y;
    if (b >= B || j >= hidden_dim) return;
    const int dst_idx = (step_row_offset + b) * hidden_dim + j;
    const int src_idx = b * hidden_dim + j;
    buf[dst_idx] = src[src_idx];
}

// ----------------------------------------------------------------------
// Phase E.4.A.T10 (2026-05-15): batched train-window store. Copies the
// per-step windows tensor [B, K, state_dim] into an all-steps buffer
// [H*B, K, state_dim] at row offset step*B. Used so the training loop
// can re-run a single batched Mamba2 forward over all visited windows
// and recover a cache for backward.
// One thread per (env, slot, feature). Grid: (sd_blocks, K, B).
// ----------------------------------------------------------------------
extern "C" __global__ void alpha_train_window_store_batched_kernel(
    const float* __restrict__ src,         // [B, K, state_dim]
    float*       __restrict__ buf,         // [H*B, K, state_dim]
    int          step_row_offset,          // step * B  (row count of width K*state_dim)
    int          B,
    int          K,
    int          state_dim
) {
    const int j = blockIdx.x * blockDim.x + threadIdx.x;
    const int k = blockIdx.y;
    const int b = blockIdx.z;
    if (b >= B || k >= K || j >= state_dim) return;
    const int src_idx = (b * K + k) * state_dim + j;
    const int dst_idx = ((step_row_offset + b) * K + k) * state_dim + j;
    buf[dst_idx] = src[src_idx];
}
