/**
 * Update rolling history buffer: shift left by 1, insert [h_s2(SH2); ofi_embed(embed_dim)]
 * as one concatenated row at position K-1.
 *
 * Grid: ceil(N * (SH2 + embed_dim) / 256), Block: 256.
 */
extern "C" __global__ void mamba2_update_history(
    float* __restrict__ h_history,
    const float* __restrict__ h_s2,
    const float* __restrict__ ofi_embed,  /* [N, embed_dim] */
    int N,
    int K,
    int sh2,
    int embed_dim  /* 10 */
) {
    int total_width = sh2 + embed_dim;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * total_width;
    if (idx >= total) return;

    int i = idx / total_width;  /* sample index */
    int j = idx % total_width;  /* feature index within [SH2+embed_dim] */

    float* base = h_history + (long long)i * K * total_width;

    /* Shift: copy slot t+1 -> t for t=0..K-2 */
    for (int t = 0; t < K - 1; t++) {
        base[(long long)t * total_width + j] = base[(long long)(t + 1) * total_width + j];
    }

    /* Insert at K-1: first SH2 cols from h_s2, last embed_dim cols from ofi_embed */
    if (j < sh2) {
        base[(long long)(K - 1) * total_width + j] = h_s2[(long long)i * sh2 + j];
    } else {
        base[(long long)(K - 1) * total_width + j] = ofi_embed[i * embed_dim + (j - sh2)];
    }
}

/* ================================================================== */
/* Kernel: isv_temporal_route — ISV controls per-feature temporal depth*/
/* ================================================================== */

/**
 * Computes temporal_weight [SH2] from ISV embedding.
 * Each feature gets its own "how much history to use" weight.
 * temporal_weight[j] = sigmoid(w_temporal @ isv_emb + b_temporal)[j]
 *
 * Launch: grid=(1,1,1), block=(1,1,1). Single thread (SH2=256 iterations).
 */
extern "C" __global__ void isv_temporal_route(
    const float* __restrict__ isv_embedding,     /* [8] from isv_forward */
    const float* __restrict__ w_temporal_route,  /* [SH2, 8] */
    const float* __restrict__ b_temporal_route,  /* [SH2] */
    float* __restrict__ temporal_weight,         /* [SH2] output */
    int SH2
) {
    if (threadIdx.x != 0) return;
    for (int j = 0; j < SH2; j++) {
        float val = b_temporal_route[j];
        for (int k = 0; k < 8; k++)
            val += w_temporal_route[(long long)j * 8 + k] * isv_embedding[k];
        temporal_weight[j] = 1.0f / (1.0f + expf(-val));
    }
}

/* ================================================================== */
/* Kernel: mamba2_copy_enriched — graph-safe DtoD copy                 */
/* ================================================================== */

/**
 * Simple element-wise copy: src[idx] → dst[idx].
 * Replaces raw cudarc memcpy_dtod_async which is NOT captured in
 * CUDA Graph — causing stale data on replay.
 *
 * Grid: ceil(N / 256), Block: 256.
 */
extern "C" __global__ void mamba2_copy_enriched(
    const float* __restrict__ src,
    float* __restrict__ dst,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) dst[idx] = src[idx];
}


/* ================================================================== */
/* Kernel: mamba2_scan_projected_fwd                                    */
/* ================================================================== */

/**
 * Forward scan using pre-projected A_proj, B_proj from cuBLAS GEMMs.
 * W_C is passed directly for inline output projection (16 reads/thread).
 *
 * Each thread handles one (sample i, feature j):
 *   1. Sequential scan over K=8 steps using A_proj (gate) and B_proj (input)
 *   2. Compute temporal_context[j] = sum_s(w_c[j,s] * x[s])
 *   3. h_enriched[i,j] = h_s2[i,j] + temporal_weight[j] * ctx
 *
 * Grid: (B, ceil(SH2/256)), Block: 256.
 */
extern "C" __global__ void mamba2_scan_projected_fwd(
    const float* __restrict__ a_proj,
    const float* __restrict__ b_proj,
    const float* __restrict__ w_c,
    const float* __restrict__ h_s2,
    float* __restrict__ h_enriched,
    int N,
    int K,
    int sh2,
    int state_d,
    const float* __restrict__ isv_signals,
    const float* __restrict__ temporal_weight
) {
    int i = blockIdx.x;
    int j = blockIdx.y * blockDim.x + threadIdx.x;
    if (i >= N || j >= sh2) return;

    float x[16];
    for (int s = 0; s < state_d; s++)
        x[s] = 0.0f;

    float stability = 1.0f;
    if (isv_signals != NULL)
        stability = isv_signals[11];

    for (int t = 0; t < K; t++) {
        long long base = ((long long)i * K + t) * state_d;
        for (int s = 0; s < state_d; s++) {
            float gate = 1.0f / (1.0f + expf(-a_proj[base + s]));
            gate *= stability;
            x[s] = gate * x[s] + b_proj[base + s];
        }
    }

    float ctx = 0.0f;
    for (int s = 0; s < state_d; s++) {
        ctx += w_c[(long long)j * state_d + s] * x[s];
    }

    float tw = (temporal_weight != NULL) ? temporal_weight[j] : 1.0f;
    h_enriched[(long long)i * sh2 + j] = h_s2[(long long)i * sh2 + j] + tw * ctx;
}


/* ================================================================== */
/* Kernel: mamba2_scan_projected_bwd                                    */
/* ================================================================== */

/**
 * Backward scan using pre-projected A_proj, B_proj from forward pass.
 *
 * Each thread handles one (sample b, state_dim s):
 *   1. Forward replay to get x_states[0..K] and a_raw[0..K-1]
 *   2. Compute d_x_K from d_h_enriched via W_C
 *   3. Reverse scan to get d_gate and d_x per time step
 *   4. Store d_context = x_K for W_C weight gradient GEMM
 *
 * Grid: (B, ceil(STATE_D/32)), Block: 32.
 */
extern "C" __global__ void mamba2_scan_projected_bwd(
    const float* __restrict__ a_proj,
    const float* __restrict__ b_proj,
    const float* __restrict__ d_h_enriched,
    const float* __restrict__ w_c,
    const float* __restrict__ temporal_weight,
    float* __restrict__ d_gate,
    float* __restrict__ d_x_out,
    float* __restrict__ d_context,
    int B,
    int K,
    int sh2,
    int state_d,
    const float* __restrict__ isv_signals
) {
    int b = blockIdx.x;
    int s = blockIdx.y * blockDim.x + threadIdx.x;
    if (b >= B || s >= state_d) return;

    float stability = 1.0f;
    if (isv_signals != NULL)
        stability = isv_signals[11];

    float x_states[9];
    float a_raw[8];
    x_states[0] = 0.0f;

    for (int t = 0; t < K; t++) {
        long long base = ((long long)b * K + t) * state_d;
        float a_val = a_proj[base + s];
        float b_val = b_proj[base + s];
        a_raw[t] = a_val;
        float gate = 1.0f / (1.0f + expf(-a_val));
        gate *= stability;
        x_states[t + 1] = gate * x_states[t] + b_val;
    }

    const float* d_out = d_h_enriched + (long long)b * sh2;
    float d_x_s = 0.0f;
    for (int j = 0; j < sh2; j++) {
        float tw = (temporal_weight != NULL) ? temporal_weight[j] : 1.0f;
        d_x_s += d_out[j] * w_c[(long long)j * state_d + s] * tw;
    }

    d_context[(long long)b * state_d + s] = x_states[K];

    for (int t = K - 1; t >= 0; t--) {
        long long base = ((long long)b * K + t) * state_d;

        float a_val = a_raw[t];
        float gate = 1.0f / (1.0f + expf(-a_val));
        float sigmoid_deriv = gate * (1.0f - gate);

        float dg = d_x_s * x_states[t] * sigmoid_deriv * stability;
        d_gate[base + s] = dg;
        d_x_out[base + s] = d_x_s;

        d_x_s = d_x_s * gate * stability;
    }
}


/* ================================================================== */
/* Kernel: mamba2_scale_d_enriched                                      */
/* ================================================================== */

/**
 * Element-wise: d_tw[idx] = d_h_enriched[idx] * temporal_weight[j].
 * Produces the scaled gradient for d_W_C weight gradient GEMM.
 * If temporal_weight is NULL, copies d_h_enriched unchanged.
 *
 * Grid: ceil(N*SH2/256), Block: 256.
 */
extern "C" __global__ void mamba2_scale_d_enriched(
    const float* __restrict__ d_h_enriched,
    const float* __restrict__ temporal_weight,
    float* __restrict__ d_tw,
    int N,
    int sh2
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * sh2;
    if (idx >= total) return;

    int j = idx % sh2;
    float tw = (temporal_weight != NULL) ? temporal_weight[j] : 1.0f;
    d_tw[idx] = d_h_enriched[idx] * tw;
}


/* ================================================================== */
/* Kernel: extract_ofi_embed_grad                                      */
/* ================================================================== */

/**
 * Extract last embed_dim rows of d_h_history and accumulate across K
 * timesteps to produce d_ofi_embed[embed_dim, B].
 *
 * d_h_history is col-major [h_width, B*K] where h_width = SH2 + embed_dim.
 * The OFI embed gradient lives in the last embed_dim rows (rows SH2..SH2+embed_dim-1).
 * We sum over the K timestep dimension for each sample:
 *   d_ofi_embed[d, b] = sum_{k=0}^{K-1} d_h_history[(SH2+d), b*K+k]
 *
 * Grid: (ceil(B/256), embed_dim), Block: 256.
 */
extern "C" __global__ void extract_ofi_embed_grad(
    const float* __restrict__ d_h_history,  /* [h_width, B*K] col-major */
    float* __restrict__ d_ofi_embed,        /* [embed_dim, B] col-major output */
    int B, int K, int sh2, int embed_dim
) {
    int d = blockIdx.y;   /* which embed dimension */
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B || d >= embed_dim) return;
    int h_width = sh2 + embed_dim;
    float sum = 0.0f;
    for (int k = 0; k < K; k++) {
        /* d_h_history col-major [h_width, B*K]: element (row, col) = row + col * h_width */
        sum += d_h_history[(sh2 + d) + (long long)(b * K + k) * h_width];
    }
    d_ofi_embed[d + b * embed_dim] = sum;
}
