// tlob_kernel.cu — TLOB attention kernels (D.8 Plan 2 Task 6C).
//
// Compresses OFI[32] → TLOB[16] using a single-head scaled-dot-product attention
// over the feature dimension, composed from existing cuBLAS Q/K/V GEMMs + this
// per-sample SDP kernel.
//
// Architecture per sample:
//   1. Q/K/V projections: cuBLAS SGEMM [16,32] × [32,B] → [16,B]  (done in GpuTlob Rust host)
//   2. Scaled dot-product: attn_scores[16] = softmax(Q^T K / sqrt(16)), out[16] = V × attn
//   3. Output projection: cuBLAS SGEMM [16,16] × [16,B] → [16,B]   (done in GpuTlob Rust host)
//   4. tlob_write_states: scatter [16,B] into states_buf at SL_TLOB_START
//
// All kernels follow zero-atomicAdd rule (per feedback_no_atomicadd.md).
// Per-sample arrays; host reduces when needed.

#include "state_layout.cuh"

// ── tlob_sdp_forward ─────────────────────────────────────────────────────────
//
// Scaled dot-product attention for a batch of samples over the TLOB feature dim.
// Each thread handles one sample b ∈ [0, B).
//
// Inputs (all column-major [D, B]):
//   proj_q  [SL_TLOB_DIM, B]  — Q projection output
//   proj_k  [SL_TLOB_DIM, B]  — K projection output
//   proj_v  [SL_TLOB_DIM, B]  — V projection output
// Outputs:
//   attn_out [SL_TLOB_DIM, B] — attention output (V weighted sum)
//   attn_scores [SL_TLOB_DIM, B] — softmax weights (saved for backward)
//
// Grid: ceil(B/256), Block: 256.  One thread per sample.
extern "C" __global__ void tlob_sdp_forward(
    const float* __restrict__ proj_q,       // [SL_TLOB_DIM, B]
    const float* __restrict__ proj_k,       // [SL_TLOB_DIM, B]
    const float* __restrict__ proj_v,       // [SL_TLOB_DIM, B]
    float*       __restrict__ attn_out,     // [SL_TLOB_DIM, B] output
    float*       __restrict__ attn_scores,  // [SL_TLOB_DIM, B] softmax weights saved for backward
    int B
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    const int D = SL_TLOB_DIM;
    const float scale = 1.0f / sqrtf((float)D);

    // Read Q/K/V for this sample (column-major: element [d, b] = ptr[d * B + b])
    float q[SL_TLOB_DIM], k[SL_TLOB_DIM], v[SL_TLOB_DIM];
    for (int d = 0; d < D; d++) {
        q[d] = proj_q[d * B + b];
        k[d] = proj_k[d * B + b];
        v[d] = proj_v[d * B + b];
    }

    // Single-head SDP: score[d] = q[d] * k[d] * scale (element-wise; feature-dim attention)
    // This is a feature-gating attention: each dimension gates itself independently.
    // Equivalent to a per-feature learned gate driven by the Q-K inner product.
    float scores[SL_TLOB_DIM];
    float max_score = -1e38f;
    for (int d = 0; d < D; d++) {
        scores[d] = q[d] * k[d] * scale;
        if (scores[d] > max_score) max_score = scores[d];
    }

    // Softmax (numerically stable)
    float sum_exp = 0.0f;
    for (int d = 0; d < D; d++) {
        scores[d] = expf(scores[d] - max_score);
        sum_exp += scores[d];
    }
    float inv_sum = (sum_exp > 1e-12f) ? (1.0f / sum_exp) : 0.0f;
    for (int d = 0; d < D; d++) {
        scores[d] *= inv_sum;
    }

    // Output = softmax_scores ⊙ V (element-wise weighted values)
    for (int d = 0; d < D; d++) {
        attn_out[d * B + b]    = scores[d] * v[d];
        attn_scores[d * B + b] = scores[d];  // save for backward
    }
}

// ── tlob_sdp_backward ────────────────────────────────────────────────────────
//
// Backward pass for TLOB SDP.
// Inputs:
//   d_attn_out [SL_TLOB_DIM, B]   — upstream gradient
//   attn_scores [SL_TLOB_DIM, B]  — saved softmax weights from forward
//   proj_v  [SL_TLOB_DIM, B]      — saved V from forward
//   proj_q  [SL_TLOB_DIM, B]      — saved Q from forward
//   proj_k  [SL_TLOB_DIM, B]      — saved K from forward
// Outputs:
//   d_proj_q [SL_TLOB_DIM, B]     — gradient for Q projection
//   d_proj_k [SL_TLOB_DIM, B]     — gradient for K projection
//   d_proj_v [SL_TLOB_DIM, B]     — gradient for V projection
//
// Grid: ceil(B/256), Block: 256.  One thread per sample.
extern "C" __global__ void tlob_sdp_backward(
    const float* __restrict__ d_attn_out,   // [SL_TLOB_DIM, B]
    const float* __restrict__ attn_scores,  // [SL_TLOB_DIM, B]
    const float* __restrict__ proj_v,       // [SL_TLOB_DIM, B]
    const float* __restrict__ proj_q,       // [SL_TLOB_DIM, B]
    const float* __restrict__ proj_k,       // [SL_TLOB_DIM, B]
    float*       __restrict__ d_proj_q,     // [SL_TLOB_DIM, B]
    float*       __restrict__ d_proj_k,     // [SL_TLOB_DIM, B]
    float*       __restrict__ d_proj_v,     // [SL_TLOB_DIM, B]
    int B,
    float inv_batch
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    const int D = SL_TLOB_DIM;
    const float scale = 1.0f / sqrtf((float)D);

    // Read saved forward values for this sample
    float s[SL_TLOB_DIM], v[SL_TLOB_DIM], q[SL_TLOB_DIM], k[SL_TLOB_DIM];
    float dout[SL_TLOB_DIM];
    for (int d = 0; d < D; d++) {
        s[d]    = attn_scores[d * B + b];
        v[d]    = proj_v[d * B + b];
        q[d]    = proj_q[d * B + b];
        k[d]    = proj_k[d * B + b];
        dout[d] = d_attn_out[d * B + b];
    }

    // d_v[d] = dout[d] * s[d]  (element-wise, no cross-dim dependency)
    for (int d = 0; d < D; d++) {
        d_proj_v[d * B + b] = dout[d] * s[d] * inv_batch;
    }

    // d_scores[d] = dout[d] * v[d]
    float ds[SL_TLOB_DIM];
    for (int d = 0; d < D; d++) {
        ds[d] = dout[d] * v[d];
    }

    // Softmax backward: d_pre_softmax[d] = s[d] * (ds[d] - dot(s, ds))
    float dot_s_ds = 0.0f;
    for (int d = 0; d < D; d++) dot_s_ds += s[d] * ds[d];
    float d_pre[SL_TLOB_DIM];
    for (int d = 0; d < D; d++) {
        d_pre[d] = s[d] * (ds[d] - dot_s_ds);
    }

    // pre-softmax = q[d]*k[d]*scale → d_q[d] = d_pre[d]*k[d]*scale, d_k[d] = d_pre[d]*q[d]*scale
    for (int d = 0; d < D; d++) {
        d_proj_q[d * B + b] = d_pre[d] * k[d] * scale * inv_batch;
        d_proj_k[d * B + b] = d_pre[d] * q[d] * scale * inv_batch;
    }
}

// ── tlob_write_states ────────────────────────────────────────────────────────
//
// Scatter TLOB features [SL_TLOB_DIM, B] (col-major) into states_buf
// at slot SL_TLOB_START for each sample.
//
// states_buf layout: [B, state_dim_padded]  (row-major per sample, padded stride)
//
// Grid: ceil(B/256), Block: 256.
extern "C" __global__ void tlob_write_states(
    const float* __restrict__ tlob_out,   // [SL_TLOB_DIM, B] col-major
    float*       __restrict__ states_buf, // [B, state_dim_padded] row-major
    int B,
    int state_dim_padded
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    float* out_row = states_buf + (long long)b * state_dim_padded;
    for (int d = 0; d < SL_TLOB_DIM; d++) {
        out_row[SL_TLOB_START + d] = tlob_out[d * B + b];
    }
}

// ── tlob_read_grad_from_states ───────────────────────────────────────────────
//
// Extract TLOB gradient slice from the states gradient buffer.
// d_states layout after trunk backward: [B, concat_dim] (row-major).
// TLOB sits at offset tlob_concat_off within concat (= bn_dim + TLOB_START - market_dim).
//
// Output: d_tlob_out [SL_TLOB_DIM, B] col-major (matching SDP backward input format).
//
// Grid: ceil(B/256), Block: 256.
extern "C" __global__ void tlob_read_grad_from_states(
    const float* __restrict__ d_states,   // [B, concat_dim] row-major
    float*       __restrict__ d_tlob_out, // [SL_TLOB_DIM, B] col-major
    int B,
    int concat_dim,
    int tlob_concat_off  // offset of TLOB within concat row
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    const float* row = d_states + (long long)b * concat_dim;
    for (int d = 0; d < SL_TLOB_DIM; d++) {
        d_tlob_out[d * B + b] = row[tlob_concat_off + d];
    }
}
