/**
 * Pure-CUDA PPO forward for backtest evaluation.
 *
 * Requires common_device_functions.cuh prepended via NVRTC source concatenation.
 * Launch config: grid=(ceil(N/32), 1, 1), block=(32, 1, 1).
 * 32 independent windows per block — each thread handles its own window.
 *
 * Pipeline: actor_forward → softmax → exposure_collapse(45→5) → argmax → action
 *
 * H100 occupancy: 32 threads/block × multiple blocks/SM vs old 1 thread/block
 * that left 97% of each SM idle.
 */

/* Override NUM_ACTIONS for PPO: use full 45-action factored space */
#undef NUM_ACTIONS
#define NUM_ACTIONS 45  /* PPO_NUM_ACTIONS */

/* Actor layer sizes — overridable via NVRTC #define injection */
#ifndef ACTOR_H1
#define ACTOR_H1 128
#endif
#ifndef ACTOR_H2
#define ACTOR_H2  64
#endif

#define N_EXPOSURE_BINS 5
#define ACTIONS_PER_BIN 9  /* 3 order × 3 urgency */

extern "C" __global__ void backtest_forward_ppo_kernel(
    const float* __restrict__ states,  /* [N, STATE_DIM]                       */
    /* Actor weights (6 pointers) */
    const float* __restrict__ pw1,     /* [ACTOR_H1, STATE_DIM] = [128, S_DIM] */
    const float* __restrict__ pb1,     /* [ACTOR_H1]            = [128]        */
    const float* __restrict__ pw2,     /* [ACTOR_H2, ACTOR_H1]  = [64, 128]   */
    const float* __restrict__ pb2,     /* [ACTOR_H2]            = [64]         */
    const float* __restrict__ pw3,     /* [NUM_ACTIONS, ACTOR_H2] = [45, 64]  */
    const float* __restrict__ pb3,     /* [NUM_ACTIONS]           = [45]       */
    /* Output */
    int* __restrict__ out_actions,     /* [N] greedy exposure action (0-4)     */
    int N
) {
    int window_id = blockIdx.x * blockDim.x + threadIdx.x;
    if (window_id >= N) return;

    /* ---- Load state vector ---- */
    float state[STATE_DIM];
    int base = window_id * STATE_DIM;
    for (int i = 0; i < STATE_DIM; i++) {
        state[i] = states[base + i];
    }

    /* ---- Actor forward: state → logits ---- */
    float h1[ACTOR_H1];
    float h2[ACTOR_H2];
    float logits[NUM_ACTIONS];

    /* Hidden layer 1: state → h1 with LeakyReLU */
    matvec_leaky_relu(pw1, pb1, state, h1, STATE_DIM, ACTOR_H1, 1);

    /* Hidden layer 2: h1 → h2 with LeakyReLU */
    matvec_leaky_relu(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);

    /* Output layer: h2 → logits (no activation) */
    matvec_leaky_relu(pw3, pb3, h2, logits, ACTOR_H2, NUM_ACTIONS, 0);

    /* ---- Stable softmax: logits → probabilities ---- */
    float max_logit = logits[0];
    for (int i = 1; i < NUM_ACTIONS; i++) {
        if (logits[i] > max_logit) max_logit = logits[i];
    }

    float probs[NUM_ACTIONS];
    float sum_exp = 0.0f;
    for (int i = 0; i < NUM_ACTIONS; i++) {
        probs[i] = expf(logits[i] - max_logit);
        sum_exp += probs[i];
    }
    float inv_sum = 1.0f / fmaxf(sum_exp, 1e-8f);
    for (int i = 0; i < NUM_ACTIONS; i++) {
        probs[i] *= inv_sum;
    }

    /* ---- Exposure collapse: 45 → 5 bins ----
     * PPO 45 actions = 5 exposure × 3 order × 3 urgency.
     * Exposure bin k (k=0..4) sums probs of actions k*9 .. k*9+8. */
    float exposure_scores[N_EXPOSURE_BINS];
    for (int k = 0; k < N_EXPOSURE_BINS; k++) {
        float sum = 0.0f;
        int start = k * ACTIONS_PER_BIN;
        for (int j = 0; j < ACTIONS_PER_BIN; j++) {
            sum += probs[start + j];
        }
        exposure_scores[k] = sum;
    }

    /* ---- Argmax of exposure scores → action (0-4) ---- */
    int best = 0;
    float best_s = exposure_scores[0];
    for (int k = 1; k < N_EXPOSURE_BINS; k++) {
        if (exposure_scores[k] > best_s) {
            best_s = exposure_scores[k];
            best = k;
        }
    }
    out_actions[window_id] = best;
}
