/**
 * GPU-resident curiosity forward model training kernels.
 *
 * cuBLAS GEMM pipeline (replaces the 1085ms curiosity_fwd_bwd_per_block kernel):
 *
 *   Forward:
 *     1. curiosity_shift_states     — build next_states from states[+1]
 *     2. curiosity_prepare_input    — build input[N, CUR_INPUT] (from inference cubin)
 *     3. cuBLAS GEMM1               — hidden[N, CUR_HIDDEN] = input @ W1^T
 *     4. curiosity_bias_leaky_relu  — in-place +b1, LeakyReLU (from inference cubin)
 *     5. cuBLAS GEMM2               — pred[N, CUR_OUTPUT] = hidden @ W2^T
 *     6. curiosity_mse_fwd_grad     — +b2, compute d_pred = 2/N*(pred+b2-target)
 *
 *   Backward:
 *     7. cuBLAS dW2                 — grad_w2 += d_pred^T @ hidden
 *     8. curiosity_bias_grad_reduce — grad_b2 = sum_batch(d_pred)
 *     9. cuBLAS d_hidden            — d_hidden = d_pred @ W2
 *    10. curiosity_leaky_relu_bwd   — mask d_hidden by post-activation sign
 *    11. cuBLAS dW1                 — grad_w1 += d_hidden^T @ input
 *    12. curiosity_bias_grad_reduce — grad_b1 = sum_batch(d_hidden)
 *    13. curiosity_adam_step        — Adam update per parameter group
 *
 * Architecture: [CUR_INPUT=MARKET_DIM+3] -> [CUR_HIDDEN=128] LeakyReLU(0.01) -> [CUR_OUTPUT=MARKET_DIM]
 * Total params: 128*(MARKET_DIM+3) + 128 + MARKET_DIM*128 + MARKET_DIM = 11306
 *
 * Requires common_device_functions.cuh to be prepended for CUR_INPUT/CUR_HIDDEN/CUR_OUTPUT.
 *
 * Native float everywhere.
 */

/* Total trainable parameters */
#define CUR_TOTAL_PARAMS (CUR_HIDDEN * CUR_INPUT + CUR_HIDDEN + CUR_OUTPUT * CUR_HIDDEN + CUR_OUTPUT)

/* ------------------------------------------------------------------ */
/*  Kernel 1: Build shifted next_states buffer                        */
/* ------------------------------------------------------------------ */

/**
 * Copy states[i + state_dim] -> next_states[i] for i in [0, N * state_dim).
 * Shifts the states buffer by one timestep to produce next_states.
 * Episode boundary noise (at most n_episodes samples out of total) is
 * negligible for a tiny auxiliary model.
 */
extern "C" __global__ void curiosity_shift_states(
    const float* __restrict__ states,
    float* __restrict__ next_states,
    int N,
    int state_dim
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * state_dim;
    if (i < total) {
        next_states[i] = states[i + state_dim];
    }
}

/* ------------------------------------------------------------------ */
/*  Kernel 2: Zero gradient buffers                                   */
/* ------------------------------------------------------------------ */

extern "C" __global__ void curiosity_zero_grads(float* grads, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) grads[i] = 0.0f;
}

/* ------------------------------------------------------------------ */
/*  Kernel 3: Forward + backward per-block (deterministic)            */
/* ------------------------------------------------------------------ */

/**
 * One thread per sample. Computes forward pass, MSE loss, and backpropagates
 * gradients. Each thread's gradient contributions are reduced within the block
 * using shared memory, then written to a per-block partial gradient buffer.
 * No atomicAdd -- fully deterministic.
 *
 * Output: partial_grads[blockIdx.x * CUR_TOTAL_PARAMS + param_idx] contains
 * the sum of gradient contributions from all threads in that block.
 *
 * The partial_grads layout is a flat vector of CUR_TOTAL_PARAMS per block:
 *   [0 .. w1_len)                          = grad_w1 partials
 *   [w1_len .. w1_len+b1_len)              = grad_b1 partials
 *   [w1_len+b1_len .. w1_len+b1_len+w2_len) = grad_w2 partials
 *   [w1_len+b1_len+w2_len .. total)        = grad_b2 partials
 *
 * Launch: grid=(ceil(N/BLOCK_SIZE)), block=(BLOCK_SIZE), shared_mem=BLOCK_SIZE*sizeof(float)
 */
extern "C" __global__ void curiosity_fwd_bwd_per_block(
    const float* __restrict__ states,       /* [N, state_dim]            */
    const int*   __restrict__ actions,       /* [N]                       */
    const float* __restrict__ next_states,   /* [N, state_dim]            */
    const float* __restrict__ w1,            /* [CUR_HIDDEN, CUR_INPUT]   */
    const float* __restrict__ b1,            /* [CUR_HIDDEN]              */
    const float* __restrict__ w2,            /* [CUR_OUTPUT, CUR_HIDDEN]  */
    const float* __restrict__ b2,            /* [CUR_OUTPUT]              */
    float* __restrict__ partial_grads,       /* [num_blocks, CUR_TOTAL_PARAMS] */
    int N,
    int state_dim
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int lane = threadIdx.x;
    int block_size = blockDim.x;

    /* Shared memory for intra-block reduction: one float per thread */
    extern __shared__ float shmem[];

    /* Base offset for this block's partial gradient output */
    float* block_out = partial_grads + blockIdx.x * CUR_TOTAL_PARAMS;

    /* ---- Forward + backward per sample ---- */

    /* We'll compute each thread's gradient contribution into local arrays,
     * then reduce across the block for each parameter. Since CUR_TOTAL_PARAMS
     * is ~11K, we can't store the full gradient vector per thread. Instead,
     * we compute forward+backward, and then iterate over each parameter index,
     * contributing to the block-level reduction one parameter at a time.
     *
     * But that would require re-running forward+backward for each parameter --
     * too expensive. Instead, each thread stores its intermediate values
     * (input, hidden, pre_act, d_pred, d_hidden) on the stack, then we
     * iterate over parameters and reduce. */

    /* Forward pass -- only if this thread has a valid sample */
    float input[CUR_INPUT];
    float pre_act[CUR_HIDDEN];
    float hidden[CUR_HIDDEN];
    float d_pred[CUR_OUTPUT];
    float d_hidden[CUR_HIDDEN];
    int valid = (tid < N) ? 1 : 0;

    if (valid) {
        const float* state_bf = states + tid * state_dim;
        int action_idx = actions[tid];
        const float* next_state_bf = next_states + tid * state_dim;

        /* Build input: first MARKET_DIM state features + 3-class action one-hot */
        for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf[i]; }
        input[MARKET_DIM + 0] = 0.0f;
        input[MARKET_DIM + 1] = 0.0f;
        input[MARKET_DIM + 2] = 0.0f;

        /* Action to category one-hot -- decode exposure from factored action.
         * Factored action = exposure_idx * (b1*b2) + order_idx * b2 + urgency_idx.
         * With b0=9, b1=3, b2=3: exposure_idx = action / 9.
         * Categories: Short (idx 0-3), Flat (idx 4), Long (idx 5-8). */
        int exposure_idx = action_idx / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS);
        int category;
        if (exposure_idx < DQN_NUM_ACTIONS / 2)       category = 0;  /* Short */
        else if (exposure_idx == DQN_NUM_ACTIONS / 2)  category = 1;  /* Flat */
        else                                           category = 2;  /* Long */
        input[MARKET_DIM + category] = 1.0f;

        /* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, alpha=0.01) */
        for (int h = 0; h < CUR_HIDDEN; h++) {
            float sum = b1[h];
            for (int i = 0; i < CUR_INPUT; i++) {
                sum += w1[h * CUR_INPUT + i] * input[i];
            }
            pre_act[h] = sum;
            hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum;
        }

        /* Layer 2: pred = w2 * hidden + b2 (no activation) */
        float pred[CUR_OUTPUT];
        for (int o = 0; o < CUR_OUTPUT; o++) {
            float sum = b2[o];
            for (int h = 0; h < CUR_HIDDEN; h++) {
                sum += w2[o * CUR_HIDDEN + h] * hidden[h];
            }
            pred[o] = sum;
        }

        /* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */
        /* d_loss/d_pred[o] = 2 * (pred[o] - next_state_bf[o]) / CUR_OUTPUT */
        float inv_out = (2.0f / (float)CUR_OUTPUT);
        for (int o = 0; o < CUR_OUTPUT; o++) {
            d_pred[o] = (pred[o] - next_state_bf[o]) * inv_out;
        }

        /* ---- Backward: compute d_hidden from Layer 2 ---- */
        for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f;
        for (int o = 0; o < CUR_OUTPUT; o++) {
            for (int h = 0; h < CUR_HIDDEN; h++) {
                d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o];
            }
        }

        /* ---- Backward: LeakyReLU ---- */
        for (int h = 0; h < CUR_HIDDEN; h++) {
            if (pre_act[h] <= 0.0f) d_hidden[h] = d_hidden[h] * 0.01f;
        }
    }

    /* ---- Block-level reduction for each parameter ---- */
    /* Iterate over all CUR_TOTAL_PARAMS parameters. For each parameter,
     * each thread computes its gradient contribution, then we do a
     * block-level tree reduction in shared memory. Thread 0 writes the
     * result to the per-block output buffer.
     *
     * Layout: [grad_w1 | grad_b1 | grad_w2 | grad_b2]
     * grad_w1[h, i] = d_hidden[h] * input[i]   (for Layer 1)
     * grad_b1[h]    = d_hidden[h]               (for Layer 1)
     * grad_w2[o, h] = d_pred[o] * hidden[h]     (for Layer 2)
     * grad_b2[o]    = d_pred[o]                  (for Layer 2)
     */

    int w1_len = CUR_HIDDEN * CUR_INPUT;
    int b1_len = CUR_HIDDEN;
    int w2_len = CUR_OUTPUT * CUR_HIDDEN;
    /* b2_len = CUR_OUTPUT */

    for (int p = 0; p < CUR_TOTAL_PARAMS; p++) {
        /* Compute this thread's gradient contribution for parameter p */
        float my_grad = 0.0f;
        if (valid) {
            if (p < w1_len) {
                /* grad_w1[h][i]: h = p / CUR_INPUT, i = p % CUR_INPUT */
                int h = p / CUR_INPUT;
                int i = p % CUR_INPUT;
                my_grad = d_hidden[h] * input[i];
            } else if (p < w1_len + b1_len) {
                /* grad_b1[h] */
                int h = p - w1_len;
                my_grad = d_hidden[h];
            } else if (p < w1_len + b1_len + w2_len) {
                /* grad_w2[o][h]: idx = p - w1_len - b1_len, o = idx / CUR_HIDDEN, h = idx % CUR_HIDDEN */
                int idx = p - w1_len - b1_len;
                int o = idx / CUR_HIDDEN;
                int h = idx % CUR_HIDDEN;
                my_grad = d_pred[o] * hidden[h];
            } else {
                /* grad_b2[o] */
                int o = p - w1_len - b1_len - w2_len;
                my_grad = d_pred[o];
            }
        }

        /* Block-level tree reduction in shared memory */
        shmem[lane] = my_grad;
        __syncthreads();

        for (int stride = block_size / 2; stride > 0; stride >>= 1) {
            if (lane < stride) {
                shmem[lane] += shmem[lane + stride];
            }
            __syncthreads();
        }

        /* Thread 0 writes the block's partial sum */
        if (lane == 0) {
            block_out[p] = shmem[0];
        }
        __syncthreads();
    }
}

/* ------------------------------------------------------------------ */
/*  Kernel 3b: Reduce per-block partial gradients (deterministic)     */
/* ------------------------------------------------------------------ */

/**
 * One thread per parameter. Loops over all block partials and sums them
 * into the final gradient buffer. Fully deterministic -- fixed summation
 * order (block 0, 1, 2, ...).
 *
 * The gradient is written into 4 separate buffers (grad_w1, grad_b1,
 * grad_w2, grad_b2) matching the Adam optimizer's per-group layout.
 *
 * Launch: grid=(ceil(CUR_TOTAL_PARAMS / 256)), block=(256)
 */
extern "C" __global__ void curiosity_grad_reduce(
    const float* __restrict__ partial_grads,  /* [num_blocks, CUR_TOTAL_PARAMS] */
    float* __restrict__ grad_w1,              /* [CUR_HIDDEN * CUR_INPUT]       */
    float* __restrict__ grad_b1,              /* [CUR_HIDDEN]                   */
    float* __restrict__ grad_w2,              /* [CUR_OUTPUT * CUR_HIDDEN]      */
    float* __restrict__ grad_b2,              /* [CUR_OUTPUT]                   */
    int num_blocks
) {
    int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= CUR_TOTAL_PARAMS) return;

    int w1_len = CUR_HIDDEN * CUR_INPUT;
    int b1_len = CUR_HIDDEN;
    int w2_len = CUR_OUTPUT * CUR_HIDDEN;

    /* Sum across all blocks for this parameter */
    float sum = 0.0f;
    for (int blk = 0; blk < num_blocks; blk++) {
        sum += partial_grads[blk * CUR_TOTAL_PARAMS + p];
    }

    /* Write to the correct gradient buffer */
    if (p < w1_len) {
        grad_w1[p] = sum;
    } else if (p < w1_len + b1_len) {
        grad_b1[p - w1_len] = sum;
    } else if (p < w1_len + b1_len + w2_len) {
        grad_w2[p - w1_len - b1_len] = sum;
    } else {
        grad_b2[p - w1_len - b1_len - w2_len] = sum;
    }
}

/* ------------------------------------------------------------------ */
/*  Kernel 4: Adam optimizer step (single param group)                */
/* ------------------------------------------------------------------ */

/**
 * One thread per parameter. Applies Adam update with bias correction.
 * Gradients are divided by batch_size to compute mean gradient.
 * Host scalar params (lr, beta1, beta2, eps) stay float, converted on entry.
 */
extern "C" __global__ void curiosity_adam_step(
    float* __restrict__ params,       /* weights/biases to update      */
    const float* __restrict__ grads,  /* accumulated gradients         */
    float* __restrict__ m,            /* first moment                  */
    float* __restrict__ v,            /* second moment                 */
    int num_params,
    int batch_size,                   /* divide grads by this for mean */
    float lr,
    float beta1,
    float beta2,
    float eps,
    int step,                         /* 1-based step counter          */
    /* SP4 Layer B: ISV-driven |p| bound (`ISV[WEIGHT_BOUND[group]].max(EPS_CLAMP_FLOOR)`)
     * — same pattern wired into all five Adam kernels. 0=disabled (debug only). */
    float weight_clamp_max_abs,
    /* SP4 Task A14 (2026-04-30): Pearl C engagement counter via warp+
     * block tree-reduce. Curiosity has 4 sub-buffer launches (W1, b1,
     * W2, b2) per step. Each sub-launch's grid is sized for its own
     * param count, so distinct engage_buf_offsets are required to
     * avoid overwriting peers' counts:
     *   W1 → SP4_ENGAGE_OFFSET_CURIOSITY_W1 (= 7 × MAX_BLOCKS = 1792)
     *   b1 → SP4_ENGAGE_OFFSET_CURIOSITY_B1 (= 8 × MAX_BLOCKS = 2048)
     *   W2 → SP4_ENGAGE_OFFSET_CURIOSITY_W2 (= 9 × MAX_BLOCKS = 2304)
     *   b2 → SP4_ENGAGE_OFFSET_CURIOSITY_B2 (= 10 × MAX_BLOCKS = 2560)
     * Host-side `pearl_c_post_adam_engagement_check(Curiosity, N)` sums
     * all four ranges. `engage_buf_offset < 0` skips the writeback. */
    int* __restrict__ nan_flags_buf,
    int diag_slot,
    int* __restrict__ clamp_engage_per_block_buf,
    int engage_buf_offset
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    bool active = (i < num_params);
    int local_engage = 0;

    if (active) {
        float lr_bf     = lr;
        float beta1_bf  = beta1;
        float beta2_bf  = beta2;
        float eps_bf    = eps;
        float one_bf    = 1.0f;

        float g = grads[i] / (float)batch_size;

        /* Skip NaN/Inf gradients — fall through to engagement reduce so
         * the per-block count stays correct. */
        if (isfinite((float)g)) {
            m[i] = beta1_bf * m[i] + (one_bf - beta1_bf) * g;
            v[i] = beta2_bf * v[i] + (one_bf - beta2_bf) * g * g;
            float m_hat = m[i] / (1.0f - powf(beta1, (float)step));
            float v_hat = v[i] / (1.0f - powf(beta2, (float)step));
            float p = params[i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf);

            /* SP3 Mech 9 + SP4 Task A14: post-Adam |p| clamp +
             * engagement counter — last step before writeback. */
            if (weight_clamp_max_abs > 0.0f) {
                if (fabsf(p) > weight_clamp_max_abs) {
                    if (nan_flags_buf != nullptr && diag_slot >= 0) {
                        nan_flags_buf[diag_slot] = 1;
                    }
                    local_engage = 1;
                }
                p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
            }

            params[i] = p;
        }
    }

    /* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
    int e = local_engage;
    for (int offset = 16; offset > 0; offset >>= 1)
        e += __shfl_xor_sync(0xFFFFFFFF, e, offset);

    __shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
    int warp_id   = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_engage[warp_id] = e;
    __syncthreads();

    if (warp_id == 0) {
        int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
            clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
        }
    }
}

/* ------------------------------------------------------------------ */
/*  Kernel 5: MSE forward + d_pred computation (cuBLAS training path) */
/* ------------------------------------------------------------------ */

/**
 * One thread per element in pred[N, CUR_OUTPUT].
 *
 * 1. Adds b2 in-place to pred (broadcast along batch).
 * 2. Computes d_pred[n, o] = (2.0 / CUR_OUTPUT) * (pred[n,o] - target[n,o]).
 *    The 1/N mean is applied later by curiosity_adam_step (divides by batch_size).
 *
 * cuBLAS GEMM 2 writes pred[N, CUR_OUTPUT] without bias. We add bias here
 * and compute the MSE gradient for the backward pass.
 *
 * Launch: grid=(ceil(N*CUR_OUTPUT/256)), block=(256)
 */
extern "C" __global__ void curiosity_mse_fwd_grad(
    float*       __restrict__ pred,       /* [N, CUR_OUTPUT] — in-place +b2, then d_pred out */
    const float* __restrict__ b2,         /* [CUR_OUTPUT]                                    */
    const float* __restrict__ target,     /* [N, state_dim] (uses first CUR_OUTPUT elements) */
    int N,
    int state_dim
) {
    int total = N * CUR_OUTPUT;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;

    int n = i / CUR_OUTPUT;   /* sample index */
    int o = i % CUR_OUTPUT;   /* output index */

    float p = pred[i] + b2[o];
    float t = target[n * state_dim + o];
    float d = (2.0f / (float)CUR_OUTPUT) * (p - t);

    /* Write back d_pred; pred buffer is reused as gradient buffer */
    pred[i] = d;
}

/* ------------------------------------------------------------------ */
/*  Kernel 6: LeakyReLU backward (in-place on d_hidden)               */
/* ------------------------------------------------------------------ */

/**
 * One thread per element in d_hidden[N, CUR_HIDDEN].
 *
 * Gates d_hidden by the sign of the post-activation hidden values:
 *   d_hidden[i] *= (hidden_post[i] > 0) ? 1.0 : 0.01
 *
 * This is valid because LeakyReLU preserves the sign of the input:
 * f(x) > 0 iff x > 0, so the mask is equivalent to (pre_act > 0).
 *
 * Launch: grid=(ceil(N*CUR_HIDDEN/256)), block=(256)
 */
extern "C" __global__ void curiosity_leaky_relu_bwd(
    float*       __restrict__ d_hidden,    /* [N, CUR_HIDDEN] — modified in-place */
    const float* __restrict__ hidden_post, /* [N, CUR_HIDDEN] — post-activation   */
    int total_elements                     /* N * CUR_HIDDEN                       */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_elements) return;
    float mask = (hidden_post[i] > 0.0f) ? 1.0f : 0.01f;
    d_hidden[i] *= mask;
}

/* ------------------------------------------------------------------ */
/*  Kernel 7: Bias gradient reduction — 2-phase shared-memory reduce  */
/* ------------------------------------------------------------------ */

/**
 * Phase 1: block-level shared-memory reduce over batch dimension.
 *
 * Each block reduces up to 256 samples for one output dimension d.
 * Grid.y selects the output dimension; Grid.x tiles over the batch.
 *
 * Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float).
 */
extern "C" __global__ void curiosity_bias_grad_reduce_p1(
    const float* __restrict__ dy,       /* [N, out_dim]                    */
    float*       __restrict__ partials, /* [num_blocks, out_dim]           */
    int N,
    int out_dim
) {
    extern __shared__ float sdata[];
    int d   = blockIdx.y;
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    float val = (gid < N) ? dy[gid * out_dim + d] : 0.0f;
    sdata[tid] = val;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid == 0) partials[blockIdx.x * out_dim + d] = sdata[0];
}

/**
 * Phase 2: final reduce across block partials.
 *
 * One thread per output dimension d. Loops over num_blocks partials
 * to produce the final bias gradient.
 *
 * Grid: ceil(out_dim/256), Block: 256.
 */
extern "C" __global__ void curiosity_bias_grad_reduce_p2(
    const float* __restrict__ partials, /* [num_blocks, out_dim] */
    float*       __restrict__ grad_b,   /* [out_dim]             */
    int out_dim,
    int num_blocks
) {
    int d = blockIdx.x * blockDim.x + threadIdx.x;
    if (d >= out_dim) return;
    float sum = 0.0f;
    for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + d];
    grad_b[d] = sum;
}
